Optimizing model hydration (zenml-io#1971)

* first checkpoint * second checkpoint * adding constants and exceptions * third checkpoint * more pipeline models * schedules and stacks and fixes * formatting and typos * components, flavors, and code references * code repository and reference * pipeline builds * formatting * service connectors * removed readme for now * plugins and servers * secret * service connectors * new naming scheme * adding filter and page models and renaming * core init * adding models * hydrated calls * started working on the schemas * formatting * adding missing models * working on schemas * base bodies added * formatting * first two models * progressing with the models * more fixes to core models * adjusting models * slight formatting * adjusted all the to_model methods to include bodies * better explanation * adding filter models * renaming and init * removed unused models for now * breaking the store * client changes * new model for loaded visualizations * base zen model in use now * builds and deployments have base classes now * formatting * formattingq * making nice progress * adjusting more utils * base zen store is also handled * rest zen store changes * endpoints * formatting * auth for server * fixing the CLI * more cli and integration fixes * fixing orchestration * fixing more references * formatting * added component base and moved constants * moved constants * moved constants * moving imports and removing models * on the way * fixing some imports * models and schemas refactored * test imports * endpoints fixec * managing integrations again * making sure that the cli works * fixing many imports * imports fixed * parameterizing the interface * fix hydration calls now * parameterized the sql zen store * running pipeline <3 * moved the external id * formatting * refactoring the rest zen store * missing import * adding the hydrate parameter to the call * modifying the endpoints * Fix forward ref issue * resolved the base issue * fixing the type hints for server * minor fixes * fixing tests * minor bugs * some linting issues * merged page models * handling the page model changes * resolving mypy issues * adding generic models in * fixed the created and updated fields * solving auth and linting * added pydantic base model * adding missing interface methods * more linting issues * solving more issues * more mypy issues and fixes * more linting issues * more issues fixed * resolving schedules * moving version to the body * linting * more linting with service connectors and users * service connector setter * minor fixes * solving the access to the metadata * fixing tests * linting * input types * response type hints * more linting * fixing input * fixing cast * unused import * fixing typehints * fixing the tests * the create methods return hyrated now * run metadata body and metadata changes * formatting * type hints gettings fixed by removing the any * fix the extra field in the service connector * ignore mypy type arg until hydration is completed * fixing docstring * fixing import * resolved conflicts * minor fixes * minor fixes * new misc module * moved server models now * removed redundant metadata * adding devices, api keys and service accounts * removed old models * moved external user * moved the other models * Fix service connector CLI * fixing imports * proper import * klinting * spell checking * fix references * name and id fixed in cli fixed functions * adding missing service account method * typos * renamed utils to update per suggestion * fixing describe functions * Add hydrate arg to where it was missing * linting * removed unnecessary methods * fixing the call * fixing the unit tests * adding missing flags * Remove code duplication in service connector models and add missing API key to connect CLI * Fix sql zen store functional tests * Fix mypy error * Add hydrate arg to `to_model` in model related schemas * flag mistake * fixed linting * Fixed zen server unit test failures and bugs * Exclude dashboard files from spell checker * update returns hydrated values now * Fix functional tests to allow hydreated update responses * missing return statements * access raises docstrings * excess returns * missing docstring parameter * missing exception * fixing docstrings * formatting * cleaned up the model creation * typo in the docstring * Fix more unit tests * Fixed more unit tests * Fix some unit tests * new validation logic implemented * fixing analytics * Fix remaining functional test failures * Fix linter * rearranging the properties for users * changed the devices as well * Fix last failing unit test * Fix the last integration test case failure --------- Co-authored-by: Michael Schuster <[email protected]> Co-authored-by: Stefan Nica <[email protected]>
Dev-Khant · Nov 16, 2023 · e17f4d3 · e17f4d3
1 parent 4ab2892
commit e17f4d3
Show file tree

Hide file tree

Showing 254 changed files with 20,440 additions and 15,433 deletions.
diff --git a/.typos.toml b/.typos.toml
@@ -1,5 +1,5 @@
 [files]
-extend-exclude = ["*.json", "*.js", "*.ipynb", "src/zenml/zen_stores/migrations/versions/", "tests/unit/materializers/test_built_in_materializer.py", "tests/integration/functional/cli/test_pipeline.py"]
+extend-exclude = ["*.json", "*.js", "*.ipynb", "src/zenml/zen_stores/migrations/versions/", "tests/unit/materializers/test_built_in_materializer.py", "tests/integration/functional/cli/test_pipeline.py", "src/zenml/zen_server/dashboard/"]
 
 [default.extend-identifiers]
 HashiCorp = "HashiCorp"

diff --git a/src/zenml/_hub/client.py b/src/zenml/_hub/client.py
@@ -28,7 +28,7 @@
 from zenml.client import Client
 from zenml.constants import ENV_ZENML_HUB_URL, IS_DEBUG_ENV
 from zenml.logger import get_logger
-from zenml.models.hub_plugin_models import (
+from zenml.models import (
     HubPluginRequestModel,
     HubPluginResponseModel,
     HubUserResponseModel,

diff --git a/src/zenml/analytics/context.py b/src/zenml/analytics/context.py
@@ -28,7 +28,7 @@
 
 if TYPE_CHECKING:
     from zenml.analytics.enums import AnalyticsEvent
-    from zenml.models.server_models import (
+    from zenml.models import (
         ServerDatabaseType,
         ServerDeploymentType,
     )

diff --git a/src/zenml/artifacts/external_artifact.py b/src/zenml/artifacts/external_artifact.py
@@ -41,28 +41,28 @@ class ExternalArtifact(ExternalArtifactConfiguration):
     write an additional step that returns this value.
 
     This class can be configured using the following parameters:
-    - value: The artifact value (any python object), that will be uploaded to the
-        artifact store.
+    - value: The artifact value (any python object), that will be uploaded to
+        the artifact store.
     - id: The ID of an artifact that is already registered in ZenML.
-    - pipeline_name & artifact_name: Name of a pipeline and artifact to search in
-        latest run.
-    - model_name & model_version & model_artifact_name & model_artifact_version: Name of a
-        model, model version, model artifact and artifact version to search.
+    - pipeline_name & artifact_name: Name of a pipeline and artifact to search
+        in latest run.
+    - model_name & model_version & model_artifact_name & model_artifact_version:
+        Name of a model, model version, model artifact and artifact version to
+        search.
 
     Args:
         value: The artifact value.
-
         id: The ID of an artifact that should be referenced by this external
             artifact.
-
         pipeline_name: Name of a pipeline to search for artifact in latest run.
-        artifact_name: Name of an artifact to be searched in latest pipeline run.
-
-        model_name: Name of a model to search for artifact in (if None - derived from step context).
-        model_version: Version of a model to search for artifact in (if None - derived from step context).
+        artifact_name: Name of an artifact to be searched in latest pipeline
+            run.
+        model_name: Name of a model to search for artifact in (if None -
+            derived from step context).
+        model_version: Version of a model to search for artifact in (if None -
+            derived from step context).
         model_artifact_name: Name of a model artifact to search for.
         model_artifact_version: Version of a model artifact to search for.
-
         materializer: The materializer to use for saving the artifact value
             to the artifact store. Only used when `value` is provided.
         store_artifact_metadata: Whether metadata for the artifact should

diff --git a/src/zenml/artifacts/external_artifact_config.py b/src/zenml/artifacts/external_artifact_config.py
@@ -26,7 +26,7 @@
 
 if TYPE_CHECKING:
     from zenml.model.model_config import ModelConfig
-    from zenml.models.artifact_models import ArtifactResponseModel
+    from zenml.models import ArtifactResponse
 
 
 logger = get_logger(__name__)
@@ -48,7 +48,7 @@ class ExternalArtifactConfiguration(BaseModel):
     model_artifact_pipeline_name: Optional[str] = None
     model_artifact_step_name: Optional[str] = None
 
-    def _get_artifact_from_pipeline_run(self) -> "ArtifactResponseModel":
+    def _get_artifact_from_pipeline_run(self) -> "ArtifactResponse":
         """Get artifact from pipeline run.
 
         Returns:
@@ -62,7 +62,9 @@ def _get_artifact_from_pipeline_run(self) -> "ArtifactResponseModel":
         client = Client()
 
         response = None
-        pipeline = client.get_pipeline(self.pipeline_name)  # type: ignore [arg-type]
+        pipeline = client.get_pipeline(
+            self.pipeline_name  # type:ignore[arg-type]
+        )
         for artifact in pipeline.last_successful_run.artifacts:
             if artifact.name == self.artifact_name:
                 response = artifact
@@ -79,7 +81,7 @@ def _get_artifact_from_pipeline_run(self) -> "ArtifactResponseModel":
 
     def _get_artifact_from_model(
         self, model_config: Optional["ModelConfig"] = None
-    ) -> "ArtifactResponseModel":
+    ) -> "ArtifactResponse":
         """Get artifact from Model Control Plane.
 
         Args:
@@ -90,8 +92,9 @@ def _get_artifact_from_model(
 
         Raises:
             RuntimeError: If artifact was not found in model version
-            RuntimeError: If `model_artifact_name` is set, but `model_name` is empty and
-                model configuration is missing in @step and @pipeline.
+            RuntimeError: If `model_artifact_name` is set, but `model_name`
+                is empty and model configuration is missing in @step and
+                @pipeline.
         """
         if self.model_name is None:
             if model_config is None:

diff --git a/src/zenml/artifacts/unmaterialized_artifact.py b/src/zenml/artifacts/unmaterialized_artifact.py
@@ -13,15 +13,16 @@
 #  permissions and limitations under the License.
 """Unmaterialized artifact class."""
 
+
 from zenml.models import (
-    ArtifactResponseModel,
-    RunMetadataResponseModel,
-    UserResponseModel,
-    WorkspaceResponseModel,
+    ArtifactResponse,
+    RunMetadataResponse,
+    UserResponse,
+    WorkspaceResponse,
 )
 
 
-class UnmaterializedArtifact(ArtifactResponseModel):
+class UnmaterializedArtifact(ArtifactResponse):
     """Unmaterialized artifact class.
 
     Typing a step input to have this type will cause ZenML to not materialize
@@ -42,7 +43,7 @@ def my_step(input_artifact: UnmaterializedArtifact):
 
 
 UnmaterializedArtifact.update_forward_refs(
-    UserResponseModel=UserResponseModel,
-    WorkspaceResponseModel=WorkspaceResponseModel,
-    RunMetadataResponseModel=RunMetadataResponseModel,
+    UserResponseModel=UserResponse,
+    WorkspaceResponseModel=WorkspaceResponse,
+    RunMetadataResponseModel=RunMetadataResponse,
 )
diff --git a/src/zenml/cli/artifact.py b/src/zenml/cli/artifact.py
@@ -23,7 +23,7 @@
 from zenml.client import Client
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
-from zenml.models.artifact_models import ArtifactFilterModel
+from zenml.models import ArtifactFilter
 from zenml.utils.pagination_utils import depaginate
 
 logger = get_logger(__name__)
@@ -34,7 +34,7 @@ def artifact() -> None:
     """List or delete artifacts."""
 
 
-@cli_utils.list_options(ArtifactFilterModel)
+@cli_utils.list_options(ArtifactFilter)
 @artifact.command("list", help="List all artifacts.")
 def list_artifacts(**kwargs: Any) -> None:
     """List all artifacts.

diff --git a/src/zenml/cli/authorized_device.py b/src/zenml/cli/authorized_device.py
@@ -23,7 +23,7 @@
 from zenml.console import console
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
-from zenml.models import OAuthDeviceFilterModel
+from zenml.models import OAuthDeviceFilter
 
 logger = get_logger(__name__)
 
@@ -58,7 +58,7 @@ def describe_authorized_device(id_or_prefix: str) -> None:
 @authorized_device.command(
     "list", help="List all authorized devices for the current user."
 )
-@list_options(OAuthDeviceFilterModel)
+@list_options(OAuthDeviceFilter)
 def list_authorized_devices(**kwargs: Any) -> None:
     """List all authorized devices.
 

diff --git a/src/zenml/cli/code_repository.py b/src/zenml/cli/code_repository.py
@@ -25,7 +25,7 @@
 from zenml.console import console
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
-from zenml.models import CodeRepositoryFilterModel
+from zenml.models import CodeRepositoryFilter
 from zenml.utils import source_utils
 
 logger = get_logger(__name__)
@@ -162,7 +162,7 @@ def register_code_repository(
 
 
 @code_repository.command("list", help="List all connected code repositories.")
-@list_options(CodeRepositoryFilterModel)
+@list_options(CodeRepositoryFilter)
 def list_code_repositories(**kwargs: Any) -> None:
     """List all connected code repositories.
 

diff --git a/src/zenml/cli/hub.py b/src/zenml/cli/hub.py
@@ -35,7 +35,7 @@
 from zenml.cli.utils import declare, error, print_table, warning
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
-from zenml.models.hub_plugin_models import (
+from zenml.models import (
     HubPluginRequestModel,
     HubPluginResponseModel,
     PluginStatus,

diff --git a/src/zenml/cli/pipeline.py b/src/zenml/cli/pipeline.py
@@ -26,14 +26,15 @@
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
 from zenml.models import (
-    PipelineBuildFilterModel,
-    PipelineFilterModel,
-    PipelineRunFilterModel,
+    PipelineBuildBase,
+    PipelineBuildFilter,
+    PipelineFilter,
+    PipelineRunFilter,
+    ScheduleFilter,
 )
-from zenml.models.pipeline_build_models import PipelineBuildBaseModel
-from zenml.models.schedule_model import ScheduleFilterModel
 from zenml.new.pipelines.pipeline import Pipeline
 from zenml.utils import source_utils, uuid_utils
+from zenml.utils.yaml_utils import write_yaml
 
 logger = get_logger(__name__)
 
@@ -193,19 +194,7 @@ def build_pipeline(
             cli_utils.declare(
                 f"Writing pipeline build output to `{output_path}`."
             )
-            with open(output_path, "w") as f:
-                f.write(
-                    build.yaml(
-                        exclude={
-                            "pipeline",
-                            "stack",
-                            "workspace",
-                            "user",
-                            "created",
-                            "updated",
-                        }
-                    )
-                )
+            write_yaml(output_path, build.to_yaml())
     else:
         cli_utils.declare("No docker builds required.")
 
@@ -284,12 +273,12 @@ def run_pipeline(
         name_id_or_prefix=pipeline_name_or_id, version=version
     )
 
-    build: Union[str, "PipelineBuildBaseModel", None] = None
+    build: Union[str, PipelineBuildBase, None] = None
     if build_path_or_id:
         if uuid_utils.is_valid_uuid(build_path_or_id):
             build = build_path_or_id
         elif os.path.exists(build_path_or_id):
-            build = PipelineBuildBaseModel.from_yaml(build_path_or_id)
+            build = PipelineBuildBase.from_yaml(build_path_or_id)
         else:
             cli_utils.error(
                 f"The specified build {build_path_or_id} is not a valid UUID "
@@ -307,7 +296,7 @@ def run_pipeline(
 
 
 @pipeline.command("list", help="List all registered pipelines.")
-@list_options(PipelineFilterModel)
+@list_options(PipelineFilter)
 def list_pipelines(**kwargs: Any) -> None:
     """List all registered pipelines.
 
@@ -383,7 +372,7 @@ def schedule() -> None:
 
 
 @schedule.command("list", help="List all pipeline schedules.")
-@list_options(ScheduleFilterModel)
+@list_options(ScheduleFilter)
 def list_schedules(**kwargs: Any) -> None:
     """List all pipeline schedules.
 
@@ -442,7 +431,7 @@ def runs() -> None:
 
 
 @runs.command("list", help="List all registered pipeline runs.")
-@list_options(PipelineRunFilterModel)
+@list_options(PipelineRunFilter)
 def list_pipeline_runs(**kwargs: Any) -> None:
     """List all registered pipeline runs for the filter.
 
@@ -508,7 +497,7 @@ def builds() -> None:
 
 
 @builds.command("list", help="List all pipeline builds.")
-@list_options(PipelineBuildFilterModel)
+@list_options(PipelineBuildFilter)
 def list_pipeline_builds(**kwargs: Any) -> None:
     """List all pipeline builds for the filter.
 

diff --git a/src/zenml/cli/role.py b/src/zenml/cli/role.py
@@ -24,7 +24,7 @@
 from zenml.console import console
 from zenml.enums import CliCategories, PermissionType
 from zenml.exceptions import EntityExistsError, IllegalOperationError
-from zenml.models import RoleFilterModel, UserRoleAssignmentFilterModel
+from zenml.models import RoleFilter, UserRoleAssignmentFilter
 
 
 @cli.group(cls=TagGroup, tag=CliCategories.IDENTITY_AND_SECURITY)
@@ -33,7 +33,7 @@ def role() -> None:
 
 
 @role.command("list")
-@list_options(RoleFilterModel)
+@list_options(RoleFilter)
 def list_roles(**kwargs: Any) -> None:
     """List all roles that fulfill the filter requirements.
 
@@ -279,7 +279,7 @@ def assignment() -> None:
 
 
 @assignment.command("list")
-@list_options(UserRoleAssignmentFilterModel)
+@list_options(UserRoleAssignmentFilter)
 def list_role_assignments(**kwargs: Any) -> None:
     """List all user role assignments that fulfill the filter requirements.
 

diff --git a/src/zenml/cli/server.py b/src/zenml/cli/server.py
@@ -613,6 +613,13 @@ def status() -> None:
     required=False,
     type=str,
 )
+@click.option(
+    "--api-key",
+    help="Use an API key to authenticate with a ZenML server. If "
+    "omitted, the web login will be used.",
+    required=False,
+    type=str,
+)
 @click.option(
     "--workspace",
     help="The workspace to use when connecting to the ZenML server.",
@@ -649,6 +656,7 @@ def connect(
     url: Optional[str] = None,
     username: Optional[str] = None,
     password: Optional[str] = None,
+    api_key: Optional[str] = None,
     workspace: Optional[str] = None,
     no_verify_ssl: bool = False,
     ssl_ca_cert: Optional[str] = None,
@@ -663,6 +671,8 @@ def connect(
             server.
         password: The password that is used to authenticate with the ZenML
             server.
+        api_key: The API key that is used to authenticate with the ZenML
+            server.
         workspace: The active workspace that is used to connect to the ZenML
             server.
         no_verify_ssl: Whether to verify the server's TLS certificate.
@@ -714,6 +724,7 @@ def connect(
         url = store_dict.get("url", url)
         username = username or store_dict.get("username")
         password = password or store_dict.get("password")
+        api_key = api_key or store_dict.get("api_key")
         verify_ssl = store_dict.get("verify_ssl", verify_ssl)
 
     elif url is None:
@@ -745,7 +756,7 @@ def connect(
     if store_type == StoreType.REST:
         store_dict["verify_ssl"] = verify_ssl
 
-    if not username:
+    if not username and not api_key:
         if store_type == StoreType.REST:
             store_dict["api_token"] = web_login(url=url, verify_ssl=verify_ssl)
         else:
@@ -761,6 +772,8 @@ def connect(
                 hide_input=True,
             )
         store_dict["password"] = password
+    elif api_key:
+        store_dict["api_key"] = api_key
 
     store_config_class = BaseZenStore.get_store_config_class(store_type)
     assert store_config_class is not None