Add save models to registry setting of a model to CLI methods (zenm…

…l-io#2447) * Fix conditional statements in GitHub workflows (zenml-io#2404) * Fix conditional statements in GitHub workflows * rename core CI flows * slow CI check doesn't happen when draft * Auto-update of Starter template * fix double conditional --------- Co-authored-by: GitHub Actions <actions@github.com> * Ensure proper spacing in error messages (zenml-io#2399) * Ensure proper spacing in error messages * update TOC (zenml-io#2406) --------- Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com> * Fix hyperai markdown table (zenml-io#2426) * build: ⬆️ Upgrade min required google-cloud-aiplatform to 1.34.0 (zenml-io#2428) * Close code block left open in the docs (zenml-io#2432) * Fix docs * wrong ticks! * Simplify HF example and notify when cache is down (zenml-io#2300) * starter files for the new CI paradigm * disable fast/slow ci on base branch * disable core workflow * Fast/slow CI core scaffold (zenml-io#2274) * give darglint check its own job * fastCI * add slowCI * reenable fast CI * remove comment * add integration tests * fix spellcheck context * enable slow CI for testing * remove unit test dependency * fix dependency installations * yamlfixed * Comment-driven CI (zenml-io#2275) * test comment-driven approach * delete unused test file * slow CI is comment-driven * restore CI * conditionally respond to comments depending on team status * add the whole team * Update .github/workflows/ci-slow.yml Co-authored-by: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> --------- Co-authored-by: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> * delete old CI --------- Co-authored-by: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> * remove spellcheck from slow CI * update spellcheck run conditions * Add GitHub issue creation on cache miss * test failing cache * Update Minio endpoint in setup_environment action.yml * Update minio-service endpoint in setup_environment action.yml * Sharded integration tests for Ubuntu (zenml-io#2286) * add pytest-shard dev dependency * update script for sharded testing * add ubuntu sharding * fix naming * Use `pytest-split` to shard CI (zenml-io#2296) * add split test to action * Update user authentication logic * Fix bug in login functionality * Refactor test coverage script * Update excluded directories in pyproject.toml * Update integration test script to include shard number * Update integration test script to use matrix.shard * Update caching key in setup_environment action.yml * Update durations path in test-coverage-xml.sh * Update cache key in setup_environment action.yml * Auto-update of Starter template * Fix formatting issue in setup_environment action.yml * Refactor code to improve performance and readability --------- Co-authored-by: GitHub Actions <actions@github.com> * Update pyproject.toml Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com> * add extra final line * make workflows use normal ubuntu) * fix durations path and update docstring --------- Co-authored-by: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Co-authored-by: GitHub Actions <actions@github.com> * Auto-update of Starter template * Auto-update of NLP template * Auto-update of E2E template * add docker testing back in * Auto-update of E2E template * temporarily trigger slow CI * revert to comment-driven CI * run full slow CI * CI as it should be * pyyaml fix * fix docker compose installation * test * test * update templates test * ubuntu-unit tests * restore unit tests back to normal * fix matrix for slow CI * uncomment the conditional checks * add input variable * remove mac and windows for testing * split out slow and fast integration testing * naming fix * confirm mac and windows ok" * improve hf and neuralprofet example * update the issue * update TOC (zenml-io#2406) * Correct docstring in integration init file (zenml-io#2408) * Fixed precedence * adding the new version to the migration tests (zenml-io#2411) * update js code for github cache miss * update context to github * add discord webhooks * Add Discord webhook support for notifications * allow fallback of cache failure * ignore if weebhock fails to to many request * Add PYTORCH integration to DockerSettings --------- Co-authored-by: Alex Strick van Linschoten <stricksubscriptions@fastmail.fm> Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com> Co-authored-by: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Co-authored-by: GitHub Actions <actions@github.com> Co-authored-by: Christian Versloot <c.versloot@infoplaza.nl> Co-authored-by: Hamza Tahir <htahir111@gmail.com> Co-authored-by: Barış Can Durak <36421093+bcdurak@users.noreply.github.com> * Adding the latest version id and name to the artifact response (zenml-io#2430) * update TOC (zenml-io#2406) * Correct docstring in integration init file (zenml-io#2408) * Fixed precedence * adding the new version to the migration tests (zenml-io#2411) * adding latest version name and id to artifact response * removed optional column from the conftest --------- Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com> Co-authored-by: Christian Versloot <c.versloot@infoplaza.nl> Co-authored-by: Hamza Tahir <htahir111@gmail.com> * Adding the ID of the producer pipeline run to artifact versions (zenml-io#2431) * adding producer pipeline run id to artifact versions * reverting one of the changes * fixing type * Add vulnerability notice to README (zenml-io#2437) * Add security vulnerability notice to README * add CVE ID * Allow more recent `adlfs` and `s3fs` versions (zenml-io#2402) * bump azure integration * bump s3 * Add new property for filtering service account events (zenml-io#2405) * add new property for filtering service account activities * Auto-update of Starter template --------- Co-authored-by: GitHub Actions <actions@github.com> * Add `download_files` method for `ArtifactVersion` (zenml-io#2434) * add save_binary method * Fix file overwrite issue in save_artifact_binary_from_response() and improve error handling * refactor * tests ofc * add docs * linting * mypy fixes * ruff fix * coderabbit suggestions * missing docstring * docstring fix * Update artifact method name from save_binary to download_binary * more renaming (save -> download) * final rename (binary -> files) * update settings syntax * Fixing `update_model`s and revert zenml-io#2402 (zenml-io#2440) * fixing update models * reverting the update model changes * linting * linting * revert zenml-io#2402 * revert adlfs changes --------- Co-authored-by: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Co-authored-by: Alex Strick van Linschoten <stricksubscriptions@fastmail.fm> * Prepare release 0.55.3 (zenml-io#2445) * alembic migration and bump version * release notes * add `save_models_to_registry` to CLI --------- Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com> Co-authored-by: GitHub Actions <actions@github.com> Co-authored-by: Christian Versloot <c.versloot@infoplaza.nl> Co-authored-by: François SERRA <francois.serra@adeo.com> Co-authored-by: jlopezpena <jlopezpena@users.noreply.github.com> Co-authored-by: Safoine El Khabich <34200873+safoinme@users.noreply.github.com> Co-authored-by: Alex Strick van Linschoten <stricksubscriptions@fastmail.fm> Co-authored-by: Hamza Tahir <htahir111@gmail.com> Co-authored-by: Barış Can Durak <36421093+bcdurak@users.noreply.github.com> Co-authored-by: Jayesh Sharma <wjayesh@outlook.com>
adtygan · Mar 20, 2024 · 8492b0d · 8492b0d
1 parent 2060067
commit 8492b0d
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 10 deletions.
diff --git a/src/zenml/cli/model.py b/src/zenml/cli/model.py
@@ -42,6 +42,9 @@ def _model_to_print(model: ModelResponse) -> Dict[str, Any]:
         "latest_version": model.latest_version_name,
         "description": model.description,
         "tags": [t.name for t in model.tags],
+        "save_to_registry": ":white_check_mark:"
+        if model.save_models_to_registry
+        else "",
         "use_cases": model.use_cases,
         "audience": model.audience,
         "limitations": model.limitations,
@@ -168,6 +171,14 @@ def list_models(**kwargs: Any) -> None:
     required=False,
     multiple=True,
 )
+@click.option(
+    "--save-models-to-registry",
+    "-s",
+    help="Whether to automatically save model artifacts to the model registry.",
+    type=click.BOOL,
+    required=False,
+    default=True,
+)
 def register_model(
     name: str,
     license: Optional[str],
@@ -178,6 +189,7 @@ def register_model(
     ethical: Optional[str],
     limitations: Optional[str],
     tag: Optional[List[str]],
+    save_models_to_registry: Optional[bool],
 ) -> None:
     """Register a new model in the Model Control Plane.
 
@@ -191,18 +203,25 @@ def register_model(
         ethical: The ethical implications of the model.
         limitations: The know limitations of the model.
         tag: Tags associated with the model.
+        save_models_to_registry: Whether to save the model to the
+            registry.
     """
     try:
         model = Client().create_model(
-            name=name,
-            license=license,
-            description=description,
-            audience=audience,
-            use_cases=use_cases,
-            trade_offs=tradeoffs,
-            ethics=ethical,
-            limitations=limitations,
-            tags=tag,
+            **remove_none_values(
+                dict(
+                    name=name,
+                    license=license,
+                    description=description,
+                    audience=audience,
+                    use_cases=use_cases,
+                    trade_offs=tradeoffs,
+                    ethics=ethical,
+                    limitations=limitations,
+                    tags=tag,
+                    save_models_to_registry=save_models_to_registry,
+                )
+            )
         )
     except (EntityExistsError, ValueError) as e:
         cli_utils.error(str(e))
@@ -282,6 +301,14 @@ def register_model(
     required=False,
     multiple=True,
 )
+@click.option(
+    "--save-models-to-registry",
+    "-s",
+    help="Whether to automatically save model artifacts to the model registry.",
+    type=click.BOOL,
+    required=False,
+    default=True,
+)
 def update_model(
     model_name_or_id: str,
     name: Optional[str],
@@ -294,6 +321,7 @@ def update_model(
     limitations: Optional[str],
     tag: Optional[List[str]],
     remove_tag: Optional[List[str]],
+    save_models_to_registry: Optional[bool],
 ) -> None:
     """Register a new model in the Model Control Plane.
 
@@ -309,6 +337,8 @@ def update_model(
         limitations: The know limitations of the model.
         tag: Tags to be added to the model.
         remove_tag: Tags to be removed from the model.
+        save_models_to_registry: Whether to save the model to the
+            registry.
     """
     model_id = Client().get_model(model_name_or_id=model_name_or_id).id
     update_dict = remove_none_values(
@@ -323,6 +353,7 @@ def update_model(
             limitations=limitations,
             add_tags=tag,
             remove_tags=remove_tag,
+            save_models_to_registry=save_models_to_registry,
         )
     )
     model = Client().update_model(model_name_or_id=model_id, **update_dict)

diff --git a/src/zenml/client.py b/src/zenml/client.py
@@ -4745,6 +4745,7 @@ def update_model(
         ethics: Optional[str] = None,
         add_tags: Optional[List[str]] = None,
         remove_tags: Optional[List[str]] = None,
+        save_models_to_registry: Optional[bool] = None,
     ) -> ModelResponse:
         """Updates an existing model in Model Control Plane.
 
@@ -4760,6 +4761,8 @@ def update_model(
             ethics: The ethical implications of the model.
             add_tags: Tags to add to the model.
             remove_tags: Tags to remove from to the model.
+            save_models_to_registry: Whether to save the model to the
+                registry.
 
         Returns:
             The updated model.
@@ -4779,6 +4782,7 @@ def update_model(
                 ethics=ethics,
                 add_tags=add_tags,
                 remove_tags=remove_tags,
+                save_models_to_registry=save_models_to_registry,
             ),
         )
 

diff --git a/src/zenml/models/v2/core/model.py b/src/zenml/models/v2/core/model.py
@@ -105,6 +105,7 @@ class ModelUpdate(BaseModel):
     ethics: Optional[str] = None
     add_tags: Optional[List[str]] = None
     remove_tags: Optional[List[str]] = None
+    save_models_to_registry: Optional[bool] = None
 
 
 # ------------------ Response Model ------------------

diff --git a/tests/integration/functional/cli/test_model.py b/tests/integration/functional/cli/test_model.py
@@ -62,6 +62,8 @@ def test_model_create_short_names(clean_client_with_models: "Client"):
             "j",
             "-t",
             "k",
+            "-s",
+            "true",
         ],
     )
     assert result.exit_code == 0, result.stderr
@@ -75,6 +77,7 @@ def test_model_create_short_names(clean_client_with_models: "Client"):
     assert model.trade_offs == "f"
     assert model.ethics == "g"
     assert model.limitations == "e"
+    assert model.save_models_to_registry
     assert {t.name for t in model.tags} == {"i", "j", "k"}
 
 
@@ -108,6 +111,8 @@ def test_model_create_full_names(clean_client_with_models: "Client"):
             "j",
             "--tag",
             "k",
+            "--save-models-to-registry",
+            "false",
         ],
     )
     assert result.exit_code == 0, result.stderr
@@ -121,6 +126,7 @@ def test_model_create_full_names(clean_client_with_models: "Client"):
     assert model.trade_offs == "f"
     assert model.ethics == "g"
     assert model.limitations == "e"
+    assert not model.save_models_to_registry
     assert {t.name for t in model.tags} == {"i", "j", "k"}
 
 
@@ -144,6 +150,7 @@ def test_model_create_only_required(clean_client_with_models: "Client"):
     assert model.trade_offs is None
     assert model.ethics is None
     assert model.limitations is None
+    assert model.save_models_to_registry
     assert len(model.tags) == 0
 
 
@@ -164,14 +171,15 @@ def test_model_update(clean_client_with_models: "Client"):
 
     result = runner.invoke(
         update_command,
-        args=[NAME, "-d", "bar", "-r", "a", "-t", "b"],
+        args=[NAME, "-d", "bar", "-r", "a", "-t", "b", "-s", "false"],
     )
     assert result.exit_code == 0, result.stderr
 
     model = clean_client_with_models.get_model(NAME)
     assert model.trade_offs == "foo"
     assert {t.name for t in model.tags} == {"b"}
     assert model.description == "bar"
+    assert not model.save_models_to_registry
 
 
 def test_model_create_without_required_fails(