Skip errors when running pyflyte register (#2111)

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Jan 18, 2024 · 11232ab · 11232ab
1 parent 6279b81
commit 11232ab
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git 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
@@ -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: