zenml-nightly 0.70.0.dev20241120__tar.gz → 0.70.0.dev20241125__tar.gz

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.70.0.dev20241120
+Version: 0.70.0.dev20241125
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.70.0.dev20241120"
+version = "0.70.0.dev20241125"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.70.0.dev20241125/src/zenml/VERSION

@@ -0,0 +1 @@
+0.70.0.dev20241125

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/artifacts/artifact_config.py

@@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field, model_validator
 
+from zenml.enums import ArtifactType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.utils.pydantic_utils import before_validator_handler
@@ -36,6 +37,7 @@ class ArtifactConfig(BaseModel):
         int, ArtifactConfig(
             name="my_artifact",  # override the default artifact name
             version=42,  # set a custom version
+            artifact_type=ArtifactType.MODEL,  # Specify the artifact type
             tags=["tag1", "tag2"],  # set custom tags
         )
     ]:
@@ -47,8 +49,9 @@ class ArtifactConfig(BaseModel):
         version: The version of the artifact.
         tags: The tags of the artifact.
         run_metadata: Metadata to add to the artifact.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
+        artifact_type: Optional type of the artifact. If not given, the type
+            specified by the materializer that is used to save this artifact
+            is used.
     """
 
     name: Optional[str] = None
@@ -58,8 +61,7 @@ class ArtifactConfig(BaseModel):
     tags: Optional[List[str]] = None
     run_metadata: Optional[Dict[str, MetadataType]] = None
 
-    is_model_artifact: bool = False
-    is_deployment_artifact: bool = False
+    artifact_type: Optional[ArtifactType] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -70,6 +72,10 @@ class ArtifactConfig(BaseModel):
         Args:
             data: The model data.
 
+        Raises:
+            ValueError: If the artifact is configured to be
+                both a model and a deployment artifact.
+
         Returns:
             Model data without the removed attributes.
         """
@@ -82,4 +88,26 @@ class ArtifactConfig(BaseModel):
                 "artifact is not supported anymore."
             )
 
+        is_model_artifact = data.pop("is_model_artifact", None)
+        is_deployment_artifact = data.pop("is_deployment_artifact", None)
+
+        if is_model_artifact and is_deployment_artifact:
+            raise ValueError(
+                "An artifact can only be a model artifact or deployment "
+                "artifact."
+            )
+        elif is_model_artifact:
+            logger.warning(
+                "`ArtifactConfig.is_model_artifact` is deprecated and will be "
+                "removed soon. Use `ArtifactConfig.artifact_type` instead."
+            )
+            data.setdefault("artifact_type", ArtifactType.MODEL)
+        elif is_deployment_artifact:
+            logger.warning(
+                "`ArtifactConfig.is_deployment_artifact` is deprecated and "
+                "will be removed soon. Use `ArtifactConfig.artifact_type` "
+                "instead."
+            )
+            data.setdefault("artifact_type", ArtifactType.SERVICE)
+
         return data

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/artifacts/utils.py

@@ -30,7 +30,6 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.preexisting_data_materializer import (
     PreexistingDataMaterializer,
 )
@@ -118,6 +117,7 @@ def _store_artifact_data_and_prepare_request(
     materializer_class: Type["BaseMaterializer"],
     save_type: ArtifactSaveType,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     store_metadata: bool = True,
     store_visualizations: bool = True,
@@ -134,6 +134,8 @@ def _store_artifact_data_and_prepare_request(
             artifact data.
         save_type: Save type of the artifact version.
         version: The artifact version.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
         tags: Tags for the artifact version.
         store_metadata: Whether to store metadata for the artifact version.
         store_visualizations: Whether to store visualizations for the artifact
@@ -176,7 +178,7 @@ def _store_artifact_data_and_prepare_request(
         artifact_name=name,
         version=version,
         tags=tags,
-        type=materializer.ASSOCIATED_ARTIFACT_TYPE,
+        type=artifact_type or materializer.ASSOCIATED_ARTIFACT_TYPE,
         uri=materializer.uri,
         materializer=source_utils.resolve(materializer.__class__),
         data_type=source_utils.resolve(data_type),
@@ -198,14 +200,13 @@ def save_artifact(
     data: Any,
     name: str,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     extract_metadata: bool = True,
     include_visualizations: bool = True,
     user_metadata: Optional[Dict[str, "MetadataType"]] = None,
     materializer: Optional["MaterializerClassOrSource"] = None,
     uri: Optional[str] = None,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
     # TODO: remove these once external artifact does not use this function anymore
     save_type: ArtifactSaveType = ArtifactSaveType.MANUAL,
     has_custom_name: bool = True,
@@ -218,6 +219,8 @@ def save_artifact(
         version: The version of the artifact. If not provided, a new
             auto-incremented version will be used.
         tags: Tags to associate with the artifact.
+        artifact_type: The artifact type. If not given, the type will be defined
+            by the materializer that is used to save the artifact.
         extract_metadata: If artifact metadata should be extracted and returned.
         include_visualizations: If artifact visualizations should be generated.
         user_metadata: User-provided metadata to store with the artifact.
@@ -226,8 +229,6 @@ def save_artifact(
         uri: The URI within the artifact store to upload the artifact
             to. If not provided, the artifact will be uploaded to
             `custom_artifacts/{name}/{version}`.
-        is_model_artifact: If the artifact is a model artifact.
-        is_deployment_artifact: If the artifact is a deployment artifact.
         save_type: The type of save operation that created the artifact version.
         has_custom_name: If the artifact name is custom and should be listed in
             the dashboard "Artifacts" tab.
@@ -273,6 +274,7 @@ def save_artifact(
         materializer_class=materializer_class,
         save_type=save_type,
         version=version,
+        artifact_type=artifact_type,
         tags=tags,
         store_metadata=extract_metadata,
         store_visualizations=include_visualizations,
@@ -286,8 +288,6 @@ def save_artifact(
     if save_type == ArtifactSaveType.MANUAL:
         _link_artifact_version_to_the_step_and_model(
             artifact_version=artifact_version,
-            is_model_artifact=is_model_artifact,
-            is_deployment_artifact=is_deployment_artifact,
         )
 
     return artifact_version
@@ -297,10 +297,9 @@ def register_artifact(
     folder_or_file_uri: str,
     name: str,
     version: Optional[Union[int, str]] = None,
+    artifact_type: Optional[ArtifactType] = None,
     tags: Optional[List[str]] = None,
     has_custom_name: bool = True,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
     artifact_metadata: Dict[str, "MetadataType"] = {},
 ) -> "ArtifactVersionResponse":
     """Register existing data stored in the artifact store as a ZenML Artifact.
@@ -311,11 +310,11 @@ def register_artifact(
         name: The name of the artifact.
         version: The version of the artifact. If not provided, a new
             auto-incremented version will be used.
+        artifact_type: The artifact type. If not given, the type will default
+            to `data`.
         tags: Tags to associate with the artifact.
         has_custom_name: If the artifact name is custom and should be listed in
             the dashboard "Artifacts" tab.
-        is_model_artifact: If the artifact is a model artifact.
-        is_deployment_artifact: If the artifact is a deployment artifact.
         artifact_metadata: Metadata dictionary to attach to the artifact version.
 
     Returns:
@@ -346,7 +345,7 @@ def register_artifact(
         artifact_name=name,
         version=version,
         tags=tags,
-        type=ArtifactType.DATA,
+        type=artifact_type or ArtifactType.DATA,
         save_type=ArtifactSaveType.PREEXISTING,
         uri=folder_or_file_uri,
         materializer=source_utils.resolve(PreexistingDataMaterializer),
@@ -365,8 +364,6 @@ def register_artifact(
 
     _link_artifact_version_to_the_step_and_model(
         artifact_version=artifact_version,
-        is_model_artifact=is_model_artifact,
-        is_deployment_artifact=is_deployment_artifact,
     )
 
     return artifact_version
@@ -674,8 +671,6 @@ def _check_if_artifact_with_given_uri_already_registered(
 
 def _link_artifact_version_to_the_step_and_model(
     artifact_version: ArtifactVersionResponse,
-    is_model_artifact: bool,
-    is_deployment_artifact: bool,
 ) -> None:
     """Link an artifact version to the step and its' context model.
 
@@ -685,8 +680,6 @@ def _link_artifact_version_to_the_step_and_model(
 
     Args:
         artifact_version: The artifact version to link.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
     """
     client = Client()
     try:
@@ -706,14 +699,9 @@ def _link_artifact_version_to_the_step_and_model(
                 link_artifact_version_to_model_version,
             )
 
-            artifact_config = ArtifactConfig(
-                is_model_artifact=is_model_artifact,
-                is_deployment_artifact=is_deployment_artifact,
-            )
             link_artifact_version_to_model_version(
                 artifact_version=artifact_version,
                 model_version=step_context.model_version,
-                artifact_config=artifact_config,
             )
     except (RuntimeError, StepContextError):
         logger.debug(f"Unable to link saved artifact to {error_message}.")

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/cli/base.py

@@ -79,7 +79,7 @@ class ZenMLProjectTemplateLocation(BaseModel):
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.10.30",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.11.20",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/client.py

@@ -4220,6 +4220,7 @@ class Client(metaclass=ClientMetaClass):
         materializer: Optional[str] = None,
         workspace_id: Optional[Union[str, UUID]] = None,
         user_id: Optional[Union[str, UUID]] = None,
+        model_version_id: Optional[Union[str, UUID]] = None,
         only_unused: Optional[bool] = False,
         has_custom_name: Optional[bool] = None,
         user: Optional[Union[UUID, str]] = None,
@@ -4249,7 +4250,8 @@ class Client(metaclass=ClientMetaClass):
             uri: The uri of the artifact to filter by.
             materializer: The materializer of the artifact to filter by.
             workspace_id: The id of the workspace to filter by.
-            user_id: The  id of the user to filter by.
+            user_id: The id of the user to filter by.
+            model_version_id: Filter by model version ID.
             only_unused: Only return artifact versions that are not used in
                 any pipeline runs.
             has_custom_name: Filter artifacts with/without custom names.
@@ -4283,6 +4285,7 @@ class Client(metaclass=ClientMetaClass):
             materializer=materializer,
             workspace_id=workspace_id,
             user_id=user_id,
+            model_version_id=model_version_id,
             only_unused=only_unused,
             has_custom_name=has_custom_name,
             tag=tag,
@@ -6435,9 +6438,6 @@ class Client(metaclass=ClientMetaClass):
         logical_operator: LogicalOperators = LogicalOperators.AND,
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[Union[UUID, str]] = None,
-        user_id: Optional[Union[UUID, str]] = None,
-        model_id: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[UUID, str]] = None,
         artifact_version_id: Optional[Union[UUID, str]] = None,
         artifact_name: Optional[str] = None,
@@ -6457,9 +6457,6 @@ class Client(metaclass=ClientMetaClass):
             logical_operator: Which logical operator to use [and, or]
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
-            workspace_id: Use the workspace id for filtering
-            user_id: Use the user id for filtering
-            model_id: Use the model id for filtering
             model_version_id: Use the model version id for filtering
             artifact_version_id: Use the artifact id for filtering
             artifact_name: Use the artifact name for filtering
@@ -6482,9 +6479,6 @@ class Client(metaclass=ClientMetaClass):
                 size=size,
                 created=created,
                 updated=updated,
-                workspace_id=workspace_id,
-                user_id=user_id,
-                model_id=model_id,
                 model_version_id=model_version_id,
                 artifact_version_id=artifact_version_id,
                 artifact_name=artifact_name,
@@ -6556,9 +6550,6 @@ class Client(metaclass=ClientMetaClass):
         logical_operator: LogicalOperators = LogicalOperators.AND,
         created: Optional[Union[datetime, str]] = None,
         updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[Union[UUID, str]] = None,
-        user_id: Optional[Union[UUID, str]] = None,
-        model_id: Optional[Union[UUID, str]] = None,
         model_version_id: Optional[Union[UUID, str]] = None,
         pipeline_run_id: Optional[Union[UUID, str]] = None,
         pipeline_run_name: Optional[str] = None,
@@ -6574,9 +6565,6 @@ class Client(metaclass=ClientMetaClass):
             logical_operator: Which logical operator to use [and, or]
             created: Use to filter by time of creation
             updated: Use the last updated date for filtering
-            workspace_id: Use the workspace id for filtering
-            user_id: Use the user id for filtering
-            model_id: Use the model id for filtering
             model_version_id: Use the model version id for filtering
             pipeline_run_id: Use the pipeline run id for filtering
             pipeline_run_name: Use the pipeline run name for filtering
@@ -6595,9 +6583,6 @@ class Client(metaclass=ClientMetaClass):
                 size=size,
                 created=created,
                 updated=updated,
-                workspace_id=workspace_id,
-                user_id=user_id,
-                model_id=model_id,
                 model_version_id=model_version_id,
                 pipeline_run_id=pipeline_run_id,
                 pipeline_run_name=pipeline_run_name,

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/constants.py

@@ -353,6 +353,7 @@ FLAVORS = "/flavors"
 GET_OR_CREATE = "/get-or-create"
 HEALTH = "/health"
 INFO = "/info"
+LOAD_INFO = "/load-info"
 LOGIN = "/login"
 LOGOUT = "/logout"
 LOGS = "/logs"

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -42,8 +42,8 @@ from kubernetes import config as k8s_config
 from kubernetes.client.rest import ApiException
 
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
-    build_cluster_role_binding_manifest_for_service_account,
     build_namespace_manifest,
+    build_role_binding_manifest_for_service_account,
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
@@ -288,7 +288,7 @@ def create_edit_service_account(
     rbac_api: k8s_client.RbacAuthorizationV1Api,
     service_account_name: str,
     namespace: str,
-    cluster_role_binding_name: str = "zenml-edit",
+    role_binding_name: str = "zenml-edit",
 ) -> None:
     """Create a new Kubernetes service account with "edit" rights.
 
@@ -297,16 +297,17 @@ def create_edit_service_account(
         rbac_api: Client of Rbac Authorization V1 API of Kubernetes API.
         service_account_name: Name of the service account.
         namespace: Kubernetes namespace. Defaults to "default".
-        cluster_role_binding_name: Name of the cluster role binding.
-            Defaults to "zenml-edit".
+        role_binding_name: Name of the role binding. Defaults to "zenml-edit".
     """
-    crb_manifest = build_cluster_role_binding_manifest_for_service_account(
-        name=cluster_role_binding_name,
+    rb_manifest = build_role_binding_manifest_for_service_account(
+        name=role_binding_name,
         role_name="edit",
         service_account_name=service_account_name,
         namespace=namespace,
     )
-    _if_not_exists(rbac_api.create_cluster_role_binding)(body=crb_manifest)
+    _if_not_exists(rbac_api.create_namespaced_role_binding)(
+        namespace=namespace, body=rb_manifest
+    )
 
     sa_manifest = build_service_account_manifest(
         name=service_account_name, namespace=namespace

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -59,6 +59,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_cron_job_manifest,
     build_pod_manifest,
 )
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
@@ -328,6 +329,45 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             custom_validation_function=_validate_local_requirements,
         )
 
+    @classmethod
+    def apply_default_resource_requests(
+        cls,
+        memory: str,
+        cpu: Optional[str] = None,
+        pod_settings: Optional[KubernetesPodSettings] = None,
+    ) -> KubernetesPodSettings:
+        """Applies default resource requests to a pod settings object.
+
+        Args:
+            memory: The memory resource request.
+            cpu: The CPU resource request.
+            pod_settings: The pod settings to update. A new one will be created
+                if not provided.
+
+        Returns:
+            The new or updated pod settings.
+        """
+        resources = {
+            "requests": {"memory": memory},
+        }
+        if cpu:
+            resources["requests"]["cpu"] = cpu
+        if not pod_settings:
+            pod_settings = KubernetesPodSettings(resources=resources)
+        elif not pod_settings.resources:
+            # We can't update the pod settings in place (because it's a frozen
+            # pydantic model), so we have to create a new one.
+            pod_settings = KubernetesPodSettings(
+                **pod_settings.model_dump(exclude_unset=True),
+                resources=resources,
+            )
+        else:
+            set_requests = pod_settings.resources.get("requests", {})
+            resources["requests"].update(set_requests)
+            pod_settings.resources["requests"] = resources["requests"]
+
+        return pod_settings
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -422,6 +462,17 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             )
             return
 
+        # We set some default minimum resource requests for the orchestrator pod
+        # here if the user has not specified any, because the orchestrator pod
+        # takes up some memory resources itself and, if not specified, the pod
+        # will be scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        orchestrator_pod_settings = self.apply_default_resource_requests(
+            memory="400Mi",
+            cpu="100m",
+            pod_settings=settings.orchestrator_pod_settings,
+        )
+
         # Create and run the orchestrator pod.
         pod_manifest = build_pod_manifest(
             run_name=orchestrator_run_name,
@@ -431,7 +482,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             command=command,
             args=args,
             privileged=False,
-            pod_settings=settings.orchestrator_pod_settings,
+            pod_settings=orchestrator_pod_settings,
             service_account_name=service_account_name,
             env=environment,
             mount_local_stores=self.config.is_local,

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -116,6 +116,16 @@ def main() -> None:
         env = get_config_environment_vars()
         env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
 
+        # We set some default minimum memory resource requests for the step pod
+        # here if the user has not specified any, because the step pod takes up
+        # some memory resources itself and, if not specified, the pod will be
+        # scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        pod_settings = KubernetesOrchestrator.apply_default_resource_requests(
+            memory="400Mi",
+            pod_settings=settings.pod_settings,
+        )
+
         # Define Kubernetes pod manifest.
         pod_manifest = build_pod_manifest(
             pod_name=pod_name,
@@ -126,7 +136,7 @@ def main() -> None:
             args=step_args,
             env=env,
             privileged=settings.privileged,
-            pod_settings=settings.pod_settings,
+            pod_settings=pod_settings,
             service_account_name=settings.step_pod_service_account_name
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -304,13 +304,13 @@ def build_cron_job_manifest(
     return job_manifest
 
 
-def build_cluster_role_binding_manifest_for_service_account(
+def build_role_binding_manifest_for_service_account(
     name: str,
     role_name: str,
     service_account_name: str,
     namespace: str = "default",
 ) -> Dict[str, Any]:
-    """Build a manifest for a cluster role binding of a service account.
+    """Build a manifest for a role binding of a service account.
 
     Args:
         name: Name of the cluster role binding.
@@ -323,7 +323,7 @@ def build_cluster_role_binding_manifest_for_service_account(
     """
     return {
         "apiVersion": "rbac.authorization.k8s.io/v1",
-        "kind": "ClusterRoleBinding",
+        "kind": "RoleBinding",
         "metadata": {"name": name},
         "subjects": [
             {

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -26,6 +26,7 @@ from typing import Any, List, Optional
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 from pydantic import Field
+from urllib3.exceptions import HTTPError
 
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -572,7 +573,7 @@ class KubernetesServiceConnector(ServiceConnector):
                 auth_settings=["BearerToken"],
                 response_type="VersionInfo",
             )
-        except k8s_client.ApiException as err:
+        except (k8s_client.ApiException, HTTPError) as err:
             raise AuthorizationException(
                 f"failed to verify Kubernetes cluster access: {err}"
             ) from err

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/model/utils.py

@@ -16,7 +16,6 @@
 from typing import Dict, Optional, Union
 from uuid import UUID
 
-from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
 from zenml.enums import ModelStages
 from zenml.exceptions import StepContextError
@@ -81,32 +80,18 @@ def log_model_metadata(
 def link_artifact_version_to_model_version(
     artifact_version: ArtifactVersionResponse,
     model_version: ModelVersionResponse,
-    artifact_config: Optional[ArtifactConfig] = None,
 ) -> None:
     """Link an artifact version to a model version.
 
     Args:
         artifact_version: The artifact version to link.
         model_version: The model version to link.
-        artifact_config: Output artifact configuration.
     """
-    if artifact_config:
-        is_model_artifact = artifact_config.is_model_artifact
-        is_deployment_artifact = artifact_config.is_deployment_artifact
-    else:
-        is_model_artifact = False
-        is_deployment_artifact = False
-
     client = Client()
     client.zen_store.create_model_version_artifact_link(
         ModelVersionArtifactRequest(
-            user=client.active_user.id,
-            workspace=client.active_workspace.id,
             artifact_version=artifact_version.id,
-            model=model_version.model.id,
             model_version=model_version.id,
-            is_model_artifact=is_model_artifact,
-            is_deployment_artifact=is_deployment_artifact,
         )
     )
 
@@ -114,16 +99,12 @@ def link_artifact_version_to_model_version(
 def link_artifact_to_model(
     artifact_version: ArtifactVersionResponse,
     model: Optional["Model"] = None,
-    is_model_artifact: bool = False,
-    is_deployment_artifact: bool = False,
 ) -> None:
     """Link the artifact to the model.
 
     Args:
         artifact_version: The artifact version to link.
         model: The model to link to.
-        is_model_artifact: Whether the artifact is a model artifact.
-        is_deployment_artifact: Whether the artifact is a deployment artifact.
 
     Raises:
         RuntimeError: If called outside of a step.
@@ -145,14 +126,9 @@ def link_artifact_to_model(
             )
 
     model_version = model._get_or_create_model_version()
-    artifact_config = ArtifactConfig(
-        is_model_artifact=is_model_artifact,
-        is_deployment_artifact=is_deployment_artifact,
-    )
     link_artifact_version_to_model_version(
         artifact_version=artifact_version,
         model_version=model_version,
-        artifact_config=artifact_config,
     )
 
 

{zenml_nightly-0.70.0.dev20241120 → zenml_nightly-0.70.0.dev20241125}/src/zenml/models/__init__.py

@@ -328,7 +328,11 @@ from zenml.models.v2.misc.service_connector_type import (
     ServiceConnectorTypedResourcesModel,
     ResourceTypeModel,
 )
-from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel
+from zenml.models.v2.misc.server_models import (
+    ServerDatabaseType,
+    ServerLoadInfo,
+    ServerModel,
+)
 from zenml.models.v2.core.trigger import (
     TriggerRequest,
     TriggerFilter,
@@ -731,6 +735,7 @@ __all__ = [
     "ExternalUserModel",
     "BuildItem",
     "LoadedVisualization",
+    "ServerLoadInfo",
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",