zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/new/pipelines/run_utils.py

@@ -1,18 +1,18 @@
 """Utility functions for running pipelines."""
 
-import hashlib
 import time
-from collections import defaultdict
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
 from uuid import UUID
 
+from pydantic import BaseModel
+
 from zenml import constants
 from zenml.client import Client
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
-from zenml.config.source import SourceType
+from zenml.config.source import Source, SourceType
 from zenml.config.step_configurations import StepConfigurationUpdate
-from zenml.enums import ExecutionStatus, ModelStages
+from zenml.enums import ExecutionStatus
 from zenml.logger import get_logger
 from zenml.models import (
     FlavorFilter,
@@ -22,15 +22,12 @@ from zenml.models import (
     PipelineRunResponse,
     StackResponse,
 )
-from zenml.new.pipelines.model_utils import NewModelRequest
 from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Flavor, Stack
-from zenml.utils import cloud_utils, code_utils, notebook_utils
+from zenml.utils import code_utils, notebook_utils, source_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
-    from zenml.model.model import Model
-
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
@@ -158,116 +155,6 @@ def deploy_pipeline(
         constants.SHOULD_PREVENT_PIPELINE_EXECUTION = previous_value
 
 
-def _update_new_requesters(
-    requester_name: str,
-    model: "Model",
-    new_versions_requested: Dict[Tuple[str, Optional[str]], NewModelRequest],
-    other_models: Set["Model"],
-) -> None:
-    key = (
-        model.name,
-        str(model.version) if model.version else None,
-    )
-    if model.version is None:
-        version_existed = False
-    else:
-        try:
-            model._get_model_version()
-            version_existed = key not in new_versions_requested
-        except KeyError as e:
-            if model.version in ModelStages.values():
-                raise KeyError(
-                    f"Unable to get model `{model.name}` using stage "
-                    f"`{model.version}`, please check that the model "
-                    "version in given stage exists before running a pipeline."
-                ) from e
-            version_existed = False
-    if not version_existed:
-        model.was_created_in_this_run = True
-        new_versions_requested[key].update_request(
-            model,
-            NewModelRequest.Requester(source="step", name=requester_name),
-        )
-    else:
-        other_models.add(model)
-
-
-def prepare_model_versions(
-    deployment: Union["PipelineDeploymentBase", "PipelineDeploymentResponse"],
-) -> None:
-    """Create model versions which are missing and validate existing ones that are used in the pipeline run.
-
-    Args:
-        deployment: The pipeline deployment configuration.
-    """
-    new_versions_requested: Dict[
-        Tuple[str, Optional[str]], NewModelRequest
-    ] = defaultdict(NewModelRequest)
-    other_models: Set["Model"] = set()
-    all_steps_have_own_config = True
-    for step in deployment.step_configurations.values():
-        step_model = step.config.model
-        all_steps_have_own_config = (
-            all_steps_have_own_config and step.config.model is not None
-        )
-        if step_model:
-            _update_new_requesters(
-                model=step_model,
-                requester_name=step.config.name,
-                new_versions_requested=new_versions_requested,
-                other_models=other_models,
-            )
-    if not all_steps_have_own_config:
-        pipeline_model = deployment.pipeline_configuration.model
-        if pipeline_model:
-            _update_new_requesters(
-                model=pipeline_model,
-                requester_name=deployment.pipeline_configuration.name,
-                new_versions_requested=new_versions_requested,
-                other_models=other_models,
-            )
-    elif deployment.pipeline_configuration.model is not None:
-        logger.warning(
-            f"ModelConfig of pipeline `{deployment.pipeline_configuration.name}` is overridden in all "
-            f"steps. "
-        )
-
-    _validate_new_version_requests(new_versions_requested)
-
-    for other_model in other_models:
-        other_model._validate_config_in_runtime()
-
-
-def _validate_new_version_requests(
-    new_versions_requested: Dict[Tuple[str, Optional[str]], NewModelRequest],
-) -> None:
-    """Validate the model version that are used in the pipeline run.
-
-    Args:
-        new_versions_requested: A dict of new model version request objects.
-
-    """
-    is_cloud_model = True
-    for key, data in new_versions_requested.items():
-        model_name, model_version = key
-        if len(data.requesters) > 1:
-            logger.warning(
-                f"New version of model version `{model_name}::{model_version or 'NEW'}` "
-                f"requested in multiple decorators:\n{data.requesters}\n We recommend "
-                "that `Model` requesting new version is configured only in one "
-                "place of the pipeline."
-            )
-        model_version_response = data.model._validate_config_in_runtime()
-        is_cloud_model &= cloud_utils.is_cloud_model_version(
-            model_version_response
-        )
-    if not is_cloud_model:
-        logger.info(
-            "Models can be viewed in the dashboard using ZenML Pro. Sign up "
-            "for a free trial at https://www.zenml.io/pro/"
-        )
-
-
 def wait_for_pipeline_run_to_finish(run_id: UUID) -> "PipelineRunResponse":
     """Waits until a pipeline run is finished.
 
@@ -383,9 +270,8 @@ def upload_notebook_cell_code_if_necessary(
         RuntimeError: If the code for one of the steps that will run out of
             process cannot be extracted into a python file.
     """
-    code_archive = code_utils.CodeArchive(root=None)
     should_upload = False
-    sources_that_require_upload = []
+    resolved_notebook_sources = source_utils.get_resolved_notebook_sources()
 
     for step in deployment.step_configurations.values():
         source = step.spec.source
@@ -396,7 +282,9 @@ def upload_notebook_cell_code_if_necessary(
                 or step.config.step_operator
             ):
                 should_upload = True
-                cell_code = getattr(step.spec.source, "_cell_code", None)
+                cell_code = resolved_notebook_sources.get(
+                    source.import_path, None
+                )
 
                 # Code does not run in-process, which means we need to
                 # extract the step code into a python file
@@ -410,20 +298,53 @@ def upload_notebook_cell_code_if_necessary(
                         "of a notebook."
                     )
 
-                notebook_utils.warn_about_notebook_cell_magic_commands(
-                    cell_code=cell_code
-                )
+    if should_upload:
+        logger.info("Uploading notebook code...")
 
-                code_hash = hashlib.sha1(cell_code.encode()).hexdigest()  # nosec
-                module_name = f"extracted_notebook_code_{code_hash}"
-                file_name = f"{module_name}.py"
-                code_archive.add_file(source=cell_code, destination=file_name)
+        for _, cell_code in resolved_notebook_sources.items():
+            notebook_utils.warn_about_notebook_cell_magic_commands(
+                cell_code=cell_code
+            )
+            module_name = notebook_utils.compute_cell_replacement_module_name(
+                cell_code=cell_code
+            )
+            file_name = f"{module_name}.py"
 
-                setattr(step.spec.source, "replacement_module", module_name)
-                sources_that_require_upload.append(source)
+            code_utils.upload_notebook_code(
+                artifact_store=stack.artifact_store,
+                cell_code=cell_code,
+                file_name=file_name,
+            )
 
-    if should_upload:
-        logger.info("Archiving notebook code...")
-        code_path = code_utils.upload_code_if_necessary(code_archive)
-        for source in sources_that_require_upload:
-            setattr(source, "code_path", code_path)
+        all_deployment_sources = get_all_sources_from_value(deployment)
+
+        for source in all_deployment_sources:
+            if source.type == SourceType.NOTEBOOK:
+                setattr(source, "artifact_store_id", stack.artifact_store.id)
+
+        logger.info("Upload finished.")
+
+
+def get_all_sources_from_value(value: Any) -> List[Source]:
+    """Get all source objects from a value.
+
+    Args:
+        value: The value from which to get all the source objects.
+
+    Returns:
+        List of source objects for the given value.
+    """
+    sources = []
+    if isinstance(value, Source):
+        sources.append(value)
+    elif isinstance(value, BaseModel):
+        for v in value.__dict__.values():
+            sources.extend(get_all_sources_from_value(v))
+    elif isinstance(value, Dict):
+        for v in value.values():
+            sources.extend(get_all_sources_from_value(v))
+    elif isinstance(value, (List, Set, tuple)):
+        for v in value:
+            sources.extend(get_all_sources_from_value(v))
+
+    return sources

zenml/new/steps/step_context.py

@@ -118,7 +118,15 @@ class StepContext(metaclass=SingletonMetaClass):
         """
         from zenml.client import Client
 
+        try:
+            pipeline_run = Client().get_pipeline_run(pipeline_run.id)
+        except KeyError:
+            pass
         self.pipeline_run = pipeline_run
+        try:
+            step_run = Client().get_run_step(step_run.id)
+        except KeyError:
+            pass
         self.step_run = step_run
         self._step_run_info = step_run_info
         self._cache_enabled = cache_enabled
@@ -177,19 +185,22 @@ class StepContext(metaclass=SingletonMetaClass):
         Raises:
             StepContextError: If the `Model` object is not set in `@step` or `@pipeline`.
         """
-        if self.step_run.config.model is not None:
-            model = self.step_run.config.model
+        if (
+            self.step_run.config.model is not None
+            and self.step_run.model_version is not None
+        ):
+            model = self.step_run.model_version.to_model_class()
         elif self.pipeline_run.config.model is not None:
-            model = self.pipeline_run.config.model
+            if self.pipeline_run.model_version:
+                model = self.pipeline_run.model_version.to_model_class()
+            else:
+                model = self.pipeline_run.config.model
         else:
             raise StepContextError(
                 f"Unable to get Model in step `{self.step_name}` of pipeline "
                 f"run '{self.pipeline_run.id}': it was not set in `@step` or `@pipeline`."
             )
 
-        # warm-up the model version
-        model._get_or_create_model_version()
-
         return model
 
     # TODO: deprecate me

zenml/orchestrators/base_orchestrator.py

@@ -68,6 +68,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
         """
         return False
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return False
+
 
 class BaseOrchestrator(StackComponent, ABC):
     """Base class for all orchestrators.

zenml/orchestrators/step_launcher.py

@@ -225,6 +225,30 @@ class StepLauncher:
                         step_run
                     )
 
+                    # warm-up and register model version
+                    _step_run = None
+                    model = (
+                        self._deployment.step_configurations[
+                            step_run.name
+                        ].config.model
+                        or self._deployment.pipeline_configuration.model
+                    )
+                    if self._deployment.step_configurations[
+                        step_run.name
+                    ].config.model:
+                        _step_run = step_run_response
+
+                    if model:
+                        prep_logs_to_show = (
+                            model._prepare_model_version_before_step_launch(
+                                pipeline_run=pipeline_run,
+                                step_run=_step_run,
+                                return_logs=True,
+                            )
+                        )
+                        if prep_logs_to_show:
+                            logger.info(prep_logs_to_show)
+
                 logger.info(f"Step `{self._step_name}` has started.")
                 if execution_needed:
                     retries = 0
@@ -296,6 +320,17 @@ class StepLauncher:
                                     step_run_response.id
                                 )
                                 raise
+                else:
+                    orchestrator_utils._link_cached_artifacts_to_model(
+                        model_from_context=model,
+                        step_run=step_run,
+                        step_source=self._step.spec.source,
+                    )
+                    if model:
+                        orchestrator_utils._link_pipeline_run_to_model_from_context(
+                            pipeline_run_id=step_run.pipeline_run_id,
+                            model=model,
+                        )
 
         except:  # noqa: E722
             logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
@@ -368,10 +403,6 @@ class StepLauncher:
             Tuple that specifies whether the step needs to be executed as
             well as the response model of the registered step run.
         """
-        model = (
-            self._deployment.step_configurations[step_run.name].config.model
-            or self._deployment.pipeline_configuration.model
-        )
         input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
             step=self._step,
             run_id=step_run.pipeline_run_id,
@@ -414,20 +445,12 @@ class StepLauncher:
                 execution_needed = False
                 cached_outputs = cached_step_run.outputs
                 step_run.original_step_run_id = cached_step_run.id
+
                 step_run.outputs = {
                     output_name: artifact.id
                     for output_name, artifact in cached_outputs.items()
                 }
-                orchestrator_utils._link_cached_artifacts_to_model(
-                    model_from_context=model,
-                    step_run=step_run,
-                    step_source=self._step.spec.source,
-                )
-                if self._step.config.model:
-                    orchestrator_utils._link_pipeline_run_to_model_from_context(
-                        pipeline_run_id=step_run.pipeline_run_id,
-                        model=self._step.config.model,
-                    )
+
                 step_run.status = ExecutionStatus.CACHED
                 step_run.end_time = step_run.start_time
 

zenml/orchestrators/step_runner.py

@@ -37,7 +37,7 @@ from zenml.constants import (
     ENV_ZENML_IGNORE_FAILURE_HOOK,
     handle_bool_env_var,
 )
-from zenml.exceptions import StepContextError, StepInterfaceError
+from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.logging.step_logging import StepLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
@@ -181,8 +181,6 @@ class StepRunner:
                         for k, v in output_annotations.items()
                     },
                 )
-                # Prepare Model Context
-                self._prepare_model_context_for_step()
 
                 # Parse the inputs for the entrypoint function.
                 function_params = self._parse_inputs(
@@ -447,14 +445,24 @@ class StepRunner:
             # we use the datatype of the stored artifact
             data_type = source_utils.load(artifact.data_type)
 
+        from zenml.orchestrators.utils import (
+            register_artifact_store_filesystem,
+        )
+
         materializer_class: Type[BaseMaterializer] = (
             source_utils.load_and_validate_class(
                 artifact.materializer, expected_class=BaseMaterializer
             )
         )
-        materializer: BaseMaterializer = materializer_class(artifact.uri)
-        materializer.validate_type_compatibility(data_type)
-        return materializer.load(data_type=data_type)
+
+        with register_artifact_store_filesystem(
+            artifact.artifact_store_id
+        ) as target_artifact_store:
+            materializer: BaseMaterializer = materializer_class(
+                uri=artifact.uri, artifact_store=target_artifact_store
+            )
+            materializer.validate_type_compatibility(data_type)
+            return materializer.load(data_type=data_type)
 
     def _validate_outputs(
         self,
@@ -638,13 +646,6 @@ class StepRunner:
 
         return output_artifacts
 
-    def _prepare_model_context_for_step(self) -> None:
-        try:
-            model = get_step_context().model
-            model._get_or_create_model_version()
-        except StepContextError:
-            return
-
     def load_and_run_hook(
         self,
         hook_source: "Source",

zenml/orchestrators/utils.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Utility functions for the orchestrator."""
 
+import os
 import random
-from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, cast
 from uuid import UUID
 
 from zenml.client import Client
@@ -26,16 +26,21 @@ from zenml.config.source import Source
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
+    ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
     PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
-from zenml.enums import StoreType
+from zenml.enums import StackComponentType, StoreType
 from zenml.exceptions import StepContextError
+from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_config_to_model
 from zenml.models.v2.core.step_run import StepRunRequest
 from zenml.new.steps.step_context import get_step_context
+from zenml.stack import StackComponent
+from zenml.utils.string_utils import format_name_template
 
 if TYPE_CHECKING:
+    from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
     from zenml.artifacts.external_artifact_config import (
         ExternalArtifactConfiguration,
     )
@@ -148,10 +153,7 @@ def get_run_name(run_name_template: str) -> str:
     Returns:
         The run name derived from the template.
     """
-    date = datetime.utcnow().strftime("%Y_%m_%d")
-    time = datetime.utcnow().strftime("%H_%M_%S_%f")
-
-    run_name = run_name_template.format(date=date, time=time)
+    run_name = format_name_template(run_name_template)
 
     if run_name == "":
         raise ValueError("Empty run names are not allowed.")
@@ -305,3 +307,101 @@ def _get_model_versions_from_artifacts(
             else:
                 break
     return models
+
+
+class register_artifact_store_filesystem:
+    """Context manager for the artifact_store/filesystem_registry dependency.
+
+    Even though it is rare, sometimes we bump into cases where we are trying to
+    load artifacts that belong to an artifact store which is different from
+    the active artifact store.
+
+    In cases like this, we will try to instantiate the target artifact store
+    by creating the corresponding artifact store Python object, which ends up
+    registering the right filesystem in the filesystem registry.
+
+    The problem is, the keys in the filesystem registry are schemes (such as
+    "s3://" or "gcs://"). If we have two artifact stores with the same set of
+    supported schemes, we might end up overwriting the filesystem that belongs
+    to the active artifact store (and its authentication). That's why we have
+    to re-instantiate the active artifact store again, so the correct filesystem
+    will be restored.
+    """
+
+    def __init__(self, target_artifact_store_id: Optional[UUID]) -> None:
+        """Initialization of the context manager.
+
+        Args:
+            target_artifact_store_id: the ID of the artifact store to load.
+        """
+        self.target_artifact_store_id = target_artifact_store_id
+
+    def __enter__(self) -> "BaseArtifactStore":
+        """Entering the context manager.
+
+        It creates an instance of the target artifact store to register the
+        correct filesystem in the registry.
+
+        Returns:
+            The target artifact store object.
+
+        Raises:
+            RuntimeError: If the target artifact store can not be fetched or
+                initiated due to missing dependencies.
+        """
+        try:
+            if self.target_artifact_store_id is not None:
+                if (
+                    Client().active_stack.artifact_store.id
+                    != self.target_artifact_store_id
+                ):
+                    get_logger(__name__).debug(
+                        f"Trying to use the artifact store with ID:"
+                        f"'{self.target_artifact_store_id}'"
+                        f"which is currently not the active artifact store."
+                    )
+
+                artifact_store_model_response = Client().get_stack_component(
+                    component_type=StackComponentType.ARTIFACT_STORE,
+                    name_id_or_prefix=self.target_artifact_store_id,
+                )
+                return cast(
+                    "BaseArtifactStore",
+                    StackComponent.from_model(artifact_store_model_response),
+                )
+            else:
+                return Client().active_stack.artifact_store
+
+        except KeyError:
+            raise RuntimeError(
+                "Unable to fetch the artifact store with id: "
+                f"'{self.target_artifact_store_id}'. Check whether the "
+                "artifact store still exists and you have the right "
+                "permissions to access it."
+            )
+        except ImportError:
+            raise RuntimeError(
+                "Unable to load the implementation of the artifact store with"
+                f"id: '{self.target_artifact_store_id}'. Please make sure that "
+                "the environment that you are loading this artifact from "
+                "has the right dependencies."
+            )
+
+    def __exit__(
+        self,
+        exc_type: Optional[Any],
+        exc_value: Optional[Any],
+        traceback: Optional[Any],
+    ) -> None:
+        """Set it back to the original state.
+
+        Args:
+            exc_type: The class of the exception
+            exc_value: The instance of the exception
+            traceback: The traceback of the exception
+        """
+        if ENV_ZENML_SERVER not in os.environ:
+            # As we exit the handler, we have to re-register the filesystem
+            # that belongs to the active artifact store as it may have been
+            # overwritten.
+            Client().active_stack.artifact_store._register()

zenml/service_connectors/service_connector_utils.py

@@ -18,10 +18,10 @@ from uuid import UUID
 
 from zenml.client import Client
 from zenml.enums import StackComponentType
-from zenml.models.v2.core.service_connector import ServiceConnectorRequest
-from zenml.models.v2.misc.full_stack import (
+from zenml.models import (
     ResourcesInfo,
     ServiceConnectorInfo,
+    ServiceConnectorRequest,
     ServiceConnectorResourcesInfo,
 )
 from zenml.utils.pagination_utils import depaginate

zenml/stack/utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Util functions for handling stacks, components, and flavors."""
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Type
 
 from zenml.client import Client
 from zenml.enums import StackComponentType, StoreType
@@ -75,7 +75,16 @@ def validate_stack_component_config(
             return None
         raise
 
-    configuration = flavor_class.config_class(**configuration_dict)
+    config_class = flavor_class.config_class
+    # Make sure extras are forbidden for the config class. Due to inheritance
+    # order, some config classes allow extras by accident which we patch here.
+    validation_config_class: Type[StackComponentConfig] = type(
+        config_class.__name__,
+        (config_class,),
+        {"model_config": {"extra": "forbid"}},
+    )
+    configuration = validation_config_class(**configuration_dict)
+
     if not configuration.is_valid:
         raise ValueError(
             f"Invalid stack component configuration. Please verify "

zenml/stack_deployments/azure_stack_deployment.py

@@ -69,7 +69,8 @@ permissions and are aware of any potential costs:
 - An Azure Resource Group to contain all the resources required for the ZenML stack
 - An Azure Storage Account and Blob Storage Container registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/azure).
 - An Azure Container Registry registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/azure).
-- An AzureML Workspace registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/azureml) and used to run pipelines.
+- An AzureML Workspace registered as both a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/azureml) and a
+[ZenML step operator](https://docs.zenml.io/stack-components/step-operators/azureml) and used to run pipelines.
 A Key Vault and Application Insights instance will also be created in the same Resource Group and used to construct the AzureML Workspace.
 - An Azure Service Principal with the minimum necessary permissions to access
 the above resources.