From 11232abc09d57eeac009cfeee74242be016c283f Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 18 Jan 2024 13:13:33 -0800 Subject: [PATCH] Skip errors when running pyflyte register (#2111) Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/register.py | 10 ++++++++++ flytekit/tools/repo.py | 22 ++++++++++++++-------- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/flytekit/clis/sdk_in_container/register.py b/flytekit/clis/sdk_in_container/register.py index 4f8dddbc18..45e41efe47 100644 --- a/flytekit/clis/sdk_in_container/register.py +++ b/flytekit/clis/sdk_in_container/register.py @@ -117,6 +117,14 @@ callback=key_value_callback, help="Environment variables to set in the container, of the format `ENV_NAME=ENV_VALUE`", ) +@click.option( + "--skip-errors", + "--skip-error", + default=False, + is_flag=True, + help="Skip errors during registration. This is useful when registering multiple packages and you want to skip " + "errors for some packages.", +) @click.argument("package-or-module", type=click.Path(exists=True, readable=True, resolve_path=True), nargs=-1) @click.pass_context def register( @@ -135,6 +143,7 @@ def register( dry_run: bool, activate_launchplans: bool, env: typing.Optional[typing.Dict[str, str]], + skip_errors: bool, ): """ see help @@ -187,6 +196,7 @@ def register( env=env, dry_run=dry_run, activate_launchplans=activate_launchplans, + skip_errors=skip_errors, ) except Exception as e: raise e diff --git a/flytekit/tools/repo.py b/flytekit/tools/repo.py index 195e3eea17..c0055eb87d 100644 --- a/flytekit/tools/repo.py +++ b/flytekit/tools/repo.py @@ -220,6 +220,7 @@ def register( env: typing.Optional[typing.Dict[str, str]], dry_run: bool = False, activate_launchplans: bool = False, + skip_errors: bool = False, ): detected_root = find_common_root(package_or_module) click.secho(f"Detected Root {detected_root}, using this to create deployable package...", fg="yellow") @@ -276,14 +277,19 @@ def register( secho(og_id, "") try: if not dry_run: - i = remote.raw_register( - cp_entity, serialization_settings, version=version, create_default_launchplan=False - ) - secho(i) - if is_lp and activate_launchplans: - secho(og_id, "", op="Activation") - remote.activate_launchplan(i) - secho(i, reason="activated", op="Activation") + try: + i = remote.raw_register( + cp_entity, serialization_settings, version=version, create_default_launchplan=False + ) + secho(i, state="success") + if is_lp and activate_launchplans: + secho(og_id, "", op="Activation") + remote.activate_launchplan(i) + secho(i, reason="activated", op="Activation") + except Exception as e: + if not skip_errors: + raise e + secho(og_id, state="failed") else: secho(og_id, reason="Dry run Mode!") except RegistrationSkipped: