diff --git a/pyproject.toml b/pyproject.toml index 50d5677..908e8b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "kfp-deployer" -version = "0.1.0" +version = "0.1.1" description = "Deploy the KFP ML Pipeline from CLI." authors = ["Karno "] readme = "README.md" diff --git a/src/kfp_deployer/main.py b/src/kfp_deployer/main.py index 4cc35f2..768f33b 100644 --- a/src/kfp_deployer/main.py +++ b/src/kfp_deployer/main.py @@ -3,15 +3,21 @@ import argparse from datetime import datetime -from typing import Optional -from pytz import timezone -import kfp +from typing import Iterable, NamedTuple, Optional -LIST_PAGE_SIZE = 1000 +import kfp +from pytz import timezone KNOWN_TIMEZONE_TABLE = {"JST": "Asia/Tokyo"} +class Pipeline(NamedTuple): + """Describes a pipeline deployed on the KFP instance.""" + + id: str + name: str + + def main() -> None: """Entrypoint.""" parser = _build_argparser() @@ -68,17 +74,29 @@ def get_pipeline_id(client: kfp.Client, pipeline_name: str) -> Optional[str]: Returns: Optional[str]: If found, return Pipeline ID. If not, return None. """ - pipelines_list = client.list_pipelines(page_size=LIST_PAGE_SIZE) - if pipelines_list.pipelines: - # check pipelines_list.pipelines is not null - for p in pipelines_list.pipelines: - if p.name == pipeline_name: - # found - return p.id + for p in _iterate_pipelines(client): + if p.name == pipeline_name: + return p.id # not found return None +def _iterate_pipelines(client: kfp.Client) -> Iterable[Pipeline]: + token: Optional[str] = None + while True: + pipelines_list = client.list_pipelines(page_token=token) + # handle the response + for p in pipelines_list.pipelines: + p_obj = Pipeline(p.id, p.name) + yield p_obj + # go on to the next page + token = pipelines_list.next_page_token + if token is None: + # when the reading page reached to the end, + # next_page_token will be None. + return + + def deploy_new_pipeline( client: kfp.Client, pipeline_name: str, pipeline_file_path: str ) -> str: @@ -139,4 +157,4 @@ def create_version_str( if __name__ == "__main__": - main() \ No newline at end of file + main()