zenml-nightly 0.70.0.dev20241119__py3-none-any.whl → 0.70.0.dev20241121__py3-none-any.whl

zenml/models/v2/core/model_version_pipeline_run.py

@@ -23,23 +23,19 @@ from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
     BaseDatedResponseBody,
     BaseIdentifiedResponse,
+    BaseRequest,
     BaseResponseMetadata,
     BaseResponseResources,
 )
-from zenml.models.v2.base.filter import StrFilter
-from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
-    WorkspaceScopedRequest,
-)
+from zenml.models.v2.base.filter import BaseFilter, StrFilter
 from zenml.models.v2.core.pipeline_run import PipelineRunResponse
 
 # ------------------ Request Model ------------------
 
 
-class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
+class ModelVersionPipelineRunRequest(BaseRequest):
     """Request model for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
     pipeline_run: UUID
 
@@ -62,7 +58,6 @@ class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
 class ModelVersionPipelineRunResponseBody(BaseDatedResponseBody):
     """Response body for links between model versions and pipeline runs."""
 
-    model: UUID
     model_version: UUID
     pipeline_run: PipelineRunResponse
 
@@ -88,16 +83,6 @@ class ModelVersionPipelineRunResponse(
 ):
     """Response model for links between model versions and pipeline runs."""
 
-    # Body and metadata properties
-    @property
-    def model(self) -> UUID:
-        """The `model` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().model
-
     @property
     def model_version(self) -> UUID:
         """The `model_version` property.
@@ -120,39 +105,21 @@ class ModelVersionPipelineRunResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
+class ModelVersionPipelineRunFilter(BaseFilter):
     """Model version pipeline run links filter model."""
 
     FILTER_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *BaseFilter.FILTER_EXCLUDE_FIELDS,
         "pipeline_run_name",
         "user",
     ]
     CLI_EXCLUDE_FIELDS = [
-        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
-        "model_id",
+        *BaseFilter.CLI_EXCLUDE_FIELDS,
         "model_version_id",
-        "user_id",
-        "workspace_id",
         "updated",
         "id",
     ]
 
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The workspace of the Model Version",
-        union_mode="left_to_right",
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="The user of the Model Version",
-        union_mode="left_to_right",
-    )
-    model_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Filter by model ID",
-        union_mode="left_to_right",
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Filter by model version ID",

zenml/orchestrators/step_launcher.py

@@ -294,7 +294,6 @@ class StepLauncher:
                     ):
                         step_run_utils.link_output_artifacts_to_model_version(
                             artifacts=step_run.outputs,
-                            output_configurations=step_run.config.outputs,
                             model_version=model_version,
                         )
 

zenml/orchestrators/step_run_utils.py

@@ -14,12 +14,12 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
-from zenml.config.step_configurations import ArtifactConfiguration, Step
+from zenml.config.step_configurations import Step
 from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ArtifactSaveType, ExecutionStatus
+from zenml.enums import ExecutionStatus
 from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
@@ -344,7 +344,6 @@ def create_cached_step_runs(
             if model_version := step_model_version or pipeline_model_version:
                 link_output_artifacts_to_model_version(
                     artifacts=step_run.outputs,
-                    output_configurations=step_run.config.outputs,
                     model_version=model_version,
                 )
 
@@ -542,10 +541,7 @@ def link_pipeline_run_to_model_version(
     client = Client()
     client.zen_store.create_model_version_pipeline_run_link(
         ModelVersionPipelineRunRequest(
-            user=client.active_user.id,
-            workspace=client.active_workspace.id,
             pipeline_run=pipeline_run.id,
-            model=model_version.model.id,
             model_version=model_version.id,
         )
     )
@@ -553,26 +549,17 @@ def link_pipeline_run_to_model_version(
 
 def link_output_artifacts_to_model_version(
     artifacts: Dict[str, List[ArtifactVersionResponse]],
-    output_configurations: Mapping[str, ArtifactConfiguration],
     model_version: ModelVersionResponse,
 ) -> None:
     """Link the outputs of a step run to a model version.
 
     Args:
         artifacts: The step output artifacts.
-        output_configurations: The output configurations for the step.
         model_version: The model version to link.
     """
-    for output_name, output_artifacts in artifacts.items():
+    for output_artifacts in artifacts.values():
         for output_artifact in output_artifacts:
-            artifact_config = None
-            if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT and (
-                output_config := output_configurations.get(output_name, None)
-            ):
-                artifact_config = output_config.artifact_config
-
             link_artifact_version_to_model_version(
                 artifact_version=output_artifact,
                 model_version=model_version,
-                artifact_config=artifact_config,
             )

zenml/orchestrators/step_runner.py

@@ -247,7 +247,6 @@ class StepRunner:
                                 artifacts={
                                     k: [v] for k, v in output_artifacts.items()
                                 },
-                                output_configurations=step_run.config.outputs,
                                 model_version=model_version,
                             )
                 finally:
@@ -577,9 +576,11 @@ class StepRunner:
             uri = output_artifact_uris[output_name]
             artifact_config = output_annotations[output_name].artifact_config
 
+            artifact_type = None
             if artifact_config is not None:
                 has_custom_name = bool(artifact_config.name)
                 version = artifact_config.version
+                artifact_type = artifact_config.artifact_type
             else:
                 has_custom_name, version = False, None
 
@@ -605,6 +606,7 @@ class StepRunner:
                 data=return_value,
                 materializer_class=materializer_class,
                 uri=uri,
+                artifact_type=artifact_type,
                 store_metadata=artifact_metadata_enabled,
                 store_visualizations=artifact_visualization_enabled,
                 has_custom_name=has_custom_name,

zenml/service_connectors/service_connector_registry.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of a service connector registry."""
 
+import threading
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from zenml.logger import get_logger
@@ -34,6 +35,7 @@ class ServiceConnectorRegistry:
         """Initialize the service connector registry."""
         self.service_connector_types: Dict[str, ServiceConnectorTypeModel] = {}
         self.initialized = False
+        self.lock = threading.RLock()
 
     def register_service_connector_type(
         self,
@@ -46,23 +48,25 @@ class ServiceConnectorRegistry:
             service_connector_type: Service connector type.
             overwrite: Whether to overwrite an existing service connector type.
         """
-        if (
-            service_connector_type.connector_type
-            not in self.service_connector_types
-            or overwrite
-        ):
-            self.service_connector_types[
+        with self.lock:
+            if (
                 service_connector_type.connector_type
-            ] = service_connector_type
-            logger.debug(
-                "Registered service connector type "
-                f"{service_connector_type.connector_type}."
-            )
-        else:
-            logger.debug(
-                f"Found existing service connector for type "
-                f"{service_connector_type.connector_type}: Skipping registration."
-            )
+                not in self.service_connector_types
+                or overwrite
+            ):
+                self.service_connector_types[
+                    service_connector_type.connector_type
+                ] = service_connector_type
+                logger.debug(
+                    "Registered service connector type "
+                    f"{service_connector_type.connector_type}."
+                )
+            else:
+                logger.debug(
+                    f"Found existing service connector for type "
+                    f"{service_connector_type.connector_type}: Skipping "
+                    "registration."
+                )
 
     def get_service_connector_type(
         self,
@@ -201,54 +205,61 @@ class ServiceConnectorRegistry:
     def register_builtin_service_connectors(self) -> None:
         """Registers the default built-in service connectors."""
         # Only register built-in service connectors once
-        if self.initialized:
-            return
+        with self.lock:
+            if self.initialized:
+                return
 
-        self.initialized = True
+            self.initialized = True
 
-        try:
-            from zenml.integrations.aws.service_connectors.aws_service_connector import (  # noqa
-                AWSServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(f"Could not import AWS service connector: {e}.")
+            try:
+                from zenml.integrations.aws.service_connectors.aws_service_connector import (  # noqa
+                    AWSServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(f"Could not import AWS service connector: {e}.")
 
-        try:
-            from zenml.integrations.gcp.service_connectors.gcp_service_connector import (  # noqa
-                GCPServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(f"Could not import GCP service connector: {e}.")
+            try:
+                from zenml.integrations.gcp.service_connectors.gcp_service_connector import (  # noqa
+                    GCPServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(f"Could not import GCP service connector: {e}.")
 
-        try:
-            from zenml.integrations.azure.service_connectors.azure_service_connector import (  # noqa
-                AzureServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(f"Could not import Azure service connector: {e}.")
+            try:
+                from zenml.integrations.azure.service_connectors.azure_service_connector import (  # noqa
+                    AzureServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(
+                    f"Could not import Azure service connector: {e}."
+                )
 
-        try:
-            from zenml.integrations.kubernetes.service_connectors.kubernetes_service_connector import (  # noqa
-                KubernetesServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(
-                f"Could not import Kubernetes service connector: {e}."
-            )
+            try:
+                from zenml.integrations.kubernetes.service_connectors.kubernetes_service_connector import (  # noqa
+                    KubernetesServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(
+                    f"Could not import Kubernetes service connector: {e}."
+                )
 
-        try:
-            from zenml.service_connectors.docker_service_connector import (  # noqa
-                DockerServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(f"Could not import Docker service connector: {e}.")
+            try:
+                from zenml.service_connectors.docker_service_connector import (  # noqa
+                    DockerServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(
+                    f"Could not import Docker service connector: {e}."
+                )
 
-        try:
-            from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import (  # noqa
-                HyperAIServiceConnector,
-            )
-        except ImportError as e:
-            logger.warning(f"Could not import HyperAI service connector: {e}.")
+            try:
+                from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import (  # noqa
+                    HyperAIServiceConnector,
+                )
+            except ImportError as e:
+                logger.warning(
+                    f"Could not import HyperAI service connector: {e}."
+                )
 
 
 service_connector_registry = ServiceConnectorRegistry()

zenml/zen_server/routers/model_versions_endpoints.py

@@ -29,9 +29,11 @@ from zenml.constants import (
 )
 from zenml.models import (
     ModelVersionArtifactFilter,
+    ModelVersionArtifactRequest,
     ModelVersionArtifactResponse,
     ModelVersionFilter,
     ModelVersionPipelineRunFilter,
+    ModelVersionPipelineRunRequest,
     ModelVersionPipelineRunResponse,
     ModelVersionResponse,
     ModelVersionUpdate,
@@ -198,6 +200,34 @@ model_version_artifacts_router = APIRouter(
 )
 
 
+@model_version_artifacts_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_artifact_link(
+    model_version_artifact_link: ModelVersionArtifactRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionArtifactResponse:
+    """Create a new model version to artifact link.
+
+    Args:
+        model_version_artifact_link: The model version to artifact link to create.
+
+    Returns:
+        The created model version to artifact link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_artifact_link.model_version
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_artifact_link(
+        model_version_artifact_link
+    )
+    return mv
+
+
 @model_version_artifacts_router.get(
     "",
     response_model=Page[ModelVersionArtifactResponse],
@@ -291,6 +321,35 @@ model_version_pipeline_runs_router = APIRouter(
 )
 
 
+@model_version_pipeline_runs_router.post(
+    "",
+    responses={401: error_response, 409: error_response, 422: error_response},
+)
+@handle_exceptions
+def create_model_version_pipeline_run_link(
+    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
+    _: AuthContext = Security(authorize),
+) -> ModelVersionPipelineRunResponse:
+    """Create a new model version to pipeline run link.
+
+    Args:
+        model_version_pipeline_run_link: The model version to pipeline run link to create.
+
+    Returns:
+        - If Model Version to Pipeline Run Link already exists - returns the existing link.
+        - Otherwise, returns the newly created model version to pipeline run link.
+    """
+    model_version = zen_store().get_model_version(
+        model_version_pipeline_run_link.model_version, hydrate=False
+    )
+    verify_permission_for_model(model_version, action=Action.UPDATE)
+
+    mv = zen_store().create_model_version_pipeline_run_link(
+        model_version_pipeline_run_link
+    )
+    return mv
+
+
 @model_version_pipeline_runs_router.get(
     "",
     response_model=Page[ModelVersionPipelineRunResponse],

zenml/zen_server/routers/workspaces_endpoints.py

@@ -20,7 +20,6 @@ from fastapi import APIRouter, Depends, Security
 
 from zenml.constants import (
     API,
-    ARTIFACTS,
     CODE_REPOSITORIES,
     GET_OR_CREATE,
     MODEL_VERSIONS,
@@ -54,10 +53,6 @@ from zenml.models import (
     ComponentResponse,
     ModelRequest,
     ModelResponse,
-    ModelVersionArtifactRequest,
-    ModelVersionArtifactResponse,
-    ModelVersionPipelineRunRequest,
-    ModelVersionPipelineRunResponse,
     ModelVersionRequest,
     ModelVersionResponse,
     Page,
@@ -1442,131 +1437,6 @@ def create_model_version(
     )
 
 
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + ARTIFACTS,
-    response_model=ModelVersionArtifactResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_artifact_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_artifact_link: ModelVersionArtifactRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionArtifactResponse:
-    """Create a new model version to artifact link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_artifact_link: The model version to artifact link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        The created model version to artifact link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(model_version_artifact_link.model_version):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_artifact_link.model_version}`"
-        )
-
-    if model_version_artifact_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model version to artifact links outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_artifact_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating model to artifact links for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_artifact_link(
-        model_version_artifact_link
-    )
-    return mv
-
-
-@router.post(
-    WORKSPACES
-    + "/{workspace_name_or_id}"
-    + MODEL_VERSIONS
-    + "/{model_version_id}"
-    + RUNS,
-    response_model=ModelVersionPipelineRunResponse,
-    responses={401: error_response, 409: error_response, 422: error_response},
-)
-@handle_exceptions
-def create_model_version_pipeline_run_link(
-    workspace_name_or_id: Union[str, UUID],
-    model_version_id: UUID,
-    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
-    auth_context: AuthContext = Security(authorize),
-) -> ModelVersionPipelineRunResponse:
-    """Create a new model version to pipeline run link.
-
-    Args:
-        workspace_name_or_id: Name or ID of the workspace.
-        model_version_id: ID of the model version.
-        model_version_pipeline_run_link: The model version to pipeline run link to create.
-        auth_context: Authentication context.
-
-    Returns:
-        - If Model Version to Pipeline Run Link already exists - returns the existing link.
-        - Otherwise, returns the newly created model version to pipeline run link.
-
-    Raises:
-        IllegalOperationError: If the workspace or user specified in the
-            model version does not match the current workspace or authenticated
-            user.
-    """
-    workspace = zen_store().get_workspace(workspace_name_or_id)
-    if str(model_version_id) != str(
-        model_version_pipeline_run_link.model_version
-    ):
-        raise IllegalOperationError(
-            f"The model version id in your path `{model_version_id}` does not "
-            f"match the model version specified in the request model "
-            f"`{model_version_pipeline_run_link.model_version}`"
-        )
-
-    if model_version_pipeline_run_link.workspace != workspace.id:
-        raise IllegalOperationError(
-            "Creating model versions outside of the workspace scope "
-            f"of this endpoint `{workspace_name_or_id}` is "
-            f"not supported."
-        )
-    if model_version_pipeline_run_link.user != auth_context.user.id:
-        raise IllegalOperationError(
-            "Creating models for a user other than yourself "
-            "is not supported."
-        )
-
-    model_version = zen_store().get_model_version(model_version_id)
-    verify_permission_for_model(model_version, action=Action.UPDATE)
-
-    mv = zen_store().create_model_version_pipeline_run_link(
-        model_version_pipeline_run_link
-    )
-    return mv
-
-
 @router.post(
     WORKSPACES + "/{workspace_name_or_id}" + SERVICES,
     response_model=ServiceResponse,

zenml/zen_stores/base_zen_store.py

@@ -42,6 +42,7 @@ from zenml.enums import (
     SecretsStoreType,
     StoreType,
 )
+from zenml.exceptions import IllegalOperationError
 from zenml.logger import get_logger
 from zenml.models import (
     ServerDatabaseType,
@@ -335,7 +336,7 @@ class BaseZenStore(
             # Ensure that the active stack is still valid
             try:
                 active_stack = self.get_stack(stack_id=active_stack_id)
-            except KeyError:
+            except (KeyError, IllegalOperationError):
                 logger.warning(
                     "The current %s active stack is no longer available. "
                     "Resetting the active stack to default.",