zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/{new/steps → steps}/step_context.py

@@ -24,13 +24,12 @@ from typing import (
     Type,
 )
 
-from zenml.exceptions import EntityExistsError, StepContextError
+from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
     from zenml.artifacts.artifact_config import ArtifactConfig
-    from zenml.config.step_run_info import StepRunInfo
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.metadata.metadata_types import MetadataType
     from zenml.model.model import Model
@@ -95,8 +94,6 @@ class StepContext(metaclass=SingletonMetaClass):
         output_materializers: Mapping[str, Sequence[Type["BaseMaterializer"]]],
         output_artifact_uris: Mapping[str, str],
         output_artifact_configs: Mapping[str, Optional["ArtifactConfig"]],
-        step_run_info: "StepRunInfo",
-        cache_enabled: bool,
     ) -> None:
         """Initialize the context of the currently running step.
 
@@ -109,8 +106,6 @@ class StepContext(metaclass=SingletonMetaClass):
                 context is used in.
             output_artifact_configs: The outputs' ArtifactConfigs of the step that this
                 context is used in.
-            step_run_info: (Deprecated) info about the currently running step.
-            cache_enabled: (Deprecated) Whether caching is enabled for the step.
 
         Raises:
             StepContextError: If the keys of the output materializers and
@@ -128,11 +123,9 @@ class StepContext(metaclass=SingletonMetaClass):
         except KeyError:
             pass
         self.step_run = step_run
-        self._step_run_info = step_run_info
-        self._cache_enabled = cache_enabled
-
-        # Get the stack that we are running in
-        self._stack = Client().active_stack
+        self.model_version = (
+            step_run.model_version or pipeline_run.model_version
+        )
 
         self.step_name = self.step_run.name
 
@@ -176,45 +169,24 @@ class StepContext(metaclass=SingletonMetaClass):
         """Returns configured Model.
 
         Order of resolution to search for Model is:
-            1. Model from @step
-            2. Model from @pipeline
+            1. Model from the step context
+            2. Model from the pipeline context
 
         Returns:
             The `Model` object associated with the current step.
 
         Raises:
-            StepContextError: If the `Model` object is not set in `@step` or `@pipeline`.
+            StepContextError: If no `Model` object was specified for the step
+                or pipeline.
         """
-        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:
-            if self.pipeline_run.model_version:
-                model = self.pipeline_run.model_version.to_model_class()
-            else:
-                model = self.pipeline_run.config.model
-        else:
+        if not self.model_version:
             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`."
+                f"run '{self.pipeline_run.id}': No model has been specified "
+                "the step or pipeline."
             )
 
-        return model
-
-    # TODO: deprecate me
-    @property
-    def model_version(self) -> "Model":
-        """DEPRECATED, use `model` instead.
-
-        Returns:
-            The `Model` object associated with the current step.
-        """
-        logger.warning(
-            "Step context `model_version` is deprecated. Please use `model` instead."
-        )
-        return self.model
+        return self.model_version.to_model_class()
 
     @property
     def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
@@ -409,32 +381,6 @@ class StepContext(metaclass=SingletonMetaClass):
             output.tags = []
         output.tags += tags
 
-    def _set_artifact_config(
-        self,
-        artifact_config: "ArtifactConfig",
-        output_name: Optional[str] = None,
-    ) -> None:
-        """Adds artifact config for a given step output.
-
-        Args:
-            artifact_config: The artifact config of the output to set.
-            output_name: Optional name of the output for which to set the
-                output signature. If no name is given and the step only has a single
-                output, the metadata of this output will be added. If the
-                step has multiple outputs, an exception will be raised.
-
-        Raises:
-            EntityExistsError: If the output already has an output signature.
-        """
-        output = self._get_output(output_name)
-
-        if output.artifact_config is None:
-            output.artifact_config = artifact_config
-        else:
-            raise EntityExistsError(
-                f"Output with name '{output_name}' already has artifact config."
-            )
-
 
 class StepContextOutput:
     """Represents a step output in the step context."""

zenml/steps/step_decorator.py

@@ -17,7 +17,6 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
-    ClassVar,
     Dict,
     Mapping,
     Optional,
@@ -29,66 +28,32 @@ from typing import (
 )
 
 from zenml.logger import get_logger
-from zenml.steps import BaseStep
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.config.retry_config import StepRetryConfig
     from zenml.config.source import Source
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.model.model import Model
+    from zenml.steps import BaseStep
     from zenml.types import HookSpecification
 
-    MaterializerClassOrSource = Union[str, "Source", Type["BaseMaterializer"]]
+    MaterializerClassOrSource = Union[str, Source, Type[BaseMaterializer]]
+
     OutputMaterializersSpecification = Union[
-        "MaterializerClassOrSource",
-        Sequence["MaterializerClassOrSource"],
-        Mapping[str, "MaterializerClassOrSource"],
-        Mapping[str, Sequence["MaterializerClassOrSource"]],
+        MaterializerClassOrSource,
+        Sequence[MaterializerClassOrSource],
+        Mapping[str, MaterializerClassOrSource],
+        Mapping[str, Sequence[MaterializerClassOrSource]],
     ]
+    F = TypeVar("F", bound=Callable[..., Any])
 
 
-STEP_INNER_FUNC_NAME = "entrypoint"
-PARAM_STEP_NAME = "name"
-PARAM_ENABLE_CACHE = "enable_cache"
-PARAM_ENABLE_ARTIFACT_METADATA = "enable_artifact_metadata"
-PARAM_ENABLE_ARTIFACT_VISUALIZATION = "enable_artifact_visualization"
-PARAM_ENABLE_STEP_LOGS = "enable_step_logs"
-PARAM_STEP_OPERATOR = "step_operator"
-PARAM_EXPERIMENT_TRACKER = "experiment_tracker"
-CLASS_CONFIGURATION = "_CLASS_CONFIGURATION"
-PARAM_OUTPUT_ARTIFACTS = "output_artifacts"
-PARAM_OUTPUT_MATERIALIZERS = "output_materializers"
-PARAM_SETTINGS = "settings"
-PARAM_EXTRA_OPTIONS = "extra"
-PARAM_ON_FAILURE = "on_failure"
-PARAM_ON_SUCCESS = "on_success"
-PARAM_MODEL = "model"
-
 logger = get_logger(__name__)
 
-F = TypeVar("F", bound=Callable[..., Any])
-
-
-class _DecoratedStep(BaseStep):
-    _CLASS_CONFIGURATION: ClassVar[Optional[Dict[str, Any]]] = None
-
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        class_config = self._CLASS_CONFIGURATION or {}
-        kwargs = {**class_config, **kwargs}
-        super().__init__(*args, **kwargs)
-
-    @property
-    def source_object(self) -> Any:
-        """The source object of this step.
-
-        Returns:
-            The source object of this step.
-        """
-        return self.entrypoint
-
 
 @overload
-def step(_func: F) -> Type[BaseStep]: ...
+def step(_func: "F") -> "BaseStep": ...
 
 
 @overload
@@ -107,11 +72,12 @@ def step(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
-) -> Callable[[F], Type[BaseStep]]: ...
+    retry: Optional["StepRetryConfig"] = None,
+) -> Callable[["F"], "BaseStep"]: ...
 
 
 def step(
-    _func: Optional[F] = None,
+    _func: Optional["F"] = None,
     *,
     name: Optional[str] = None,
     enable_cache: Optional[bool] = None,
@@ -126,11 +92,9 @@ def step(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
-) -> Union[Type[BaseStep], Callable[[F], Type[BaseStep]]]:
-    """Outer decorator function for the creation of a ZenML step.
-
-    In order to be able to work with parameters such as `name`, it features a
-    nested decorator structure.
+    retry: Optional["StepRetryConfig"] = None,
+) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
+    """Decorator to create a ZenML step.
 
     Args:
         _func: The decorated function.
@@ -158,56 +122,45 @@ def step(
         on_success: Callback function in event of success of the step. Can be a
             function with no arguments, or a source path to such a function
             (e.g. `module.my_function`).
-        model: configuration of the model version in the Model Control Plane.
+        model: configuration of the model in the Model Control Plane.
+        retry: configuration of step retry in case of step failure.
 
     Returns:
-        The inner decorator which creates the step class based on the
-        ZenML BaseStep
+        The step instance.
     """
 
-    def inner_decorator(func: F) -> Type[BaseStep]:
-        """Inner decorator function for the creation of a ZenML Step.
-
-        Args:
-            func: types.FunctionType, this function will be used as the
-                "process" method of the generated Step.
-
-        Returns:
-            The class of a newly generated ZenML Step.
-        """
-        step_name = name or func.__name__
-        logger.warning(
-            f"The `@step` decorator that you used to define your {step_name} "
-            "step is deprecated. Check out the 0.40.0 migration guide for more "
-            "information on how to migrate your steps to the new syntax: "
-            "https://docs.zenml.io/reference/migration-guide/migration-zero-forty"
-        )
+    def inner_decorator(func: "F") -> "BaseStep":
+        from zenml.steps.decorated_step import _DecoratedStep
 
-        return type(  # noqa
+        class_: Type["BaseStep"] = type(
             func.__name__,
             (_DecoratedStep,),
             {
-                STEP_INNER_FUNC_NAME: staticmethod(func),
-                CLASS_CONFIGURATION: {
-                    PARAM_STEP_NAME: name,
-                    PARAM_ENABLE_CACHE: enable_cache,
-                    PARAM_ENABLE_ARTIFACT_METADATA: enable_artifact_metadata,
-                    PARAM_ENABLE_ARTIFACT_VISUALIZATION: enable_artifact_visualization,
-                    PARAM_ENABLE_STEP_LOGS: enable_step_logs,
-                    PARAM_EXPERIMENT_TRACKER: experiment_tracker,
-                    PARAM_STEP_OPERATOR: step_operator,
-                    PARAM_OUTPUT_MATERIALIZERS: output_materializers,
-                    PARAM_SETTINGS: settings,
-                    PARAM_EXTRA_OPTIONS: extra,
-                    PARAM_ON_FAILURE: on_failure,
-                    PARAM_ON_SUCCESS: on_success,
-                    PARAM_MODEL: model,
-                },
+                "entrypoint": staticmethod(func),
                 "__module__": func.__module__,
                 "__doc__": func.__doc__,
             },
         )
 
+        step_instance = class_(
+            name=name or func.__name__,
+            enable_cache=enable_cache,
+            enable_artifact_metadata=enable_artifact_metadata,
+            enable_artifact_visualization=enable_artifact_visualization,
+            enable_step_logs=enable_step_logs,
+            experiment_tracker=experiment_tracker,
+            step_operator=step_operator,
+            output_materializers=output_materializers,
+            settings=settings,
+            extra=extra,
+            on_failure=on_failure,
+            on_success=on_success,
+            model=model,
+            retry=retry,
+        )
+
+        return step_instance
+
     if _func is None:
         return inner_decorator
     else:

zenml/steps/step_invocation.py

@@ -13,14 +13,16 @@
 #  permissions and limitations under the License.
 """Step invocation class definition."""
 
-from typing import TYPE_CHECKING, Any, Dict, Set
+from typing import TYPE_CHECKING, Any, Dict, Set, Union
+
+from zenml.models import ArtifactVersionResponse
 
 if TYPE_CHECKING:
     from zenml.artifacts.external_artifact import ExternalArtifact
     from zenml.client_lazy_loader import ClientLazyLoader
     from zenml.config.step_configurations import StepConfiguration
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
-    from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.steps import BaseStep
     from zenml.steps.entrypoint_function_utils import StepArtifact
 
@@ -33,7 +35,9 @@ class StepInvocation:
         id: str,
         step: "BaseStep",
         input_artifacts: Dict[str, "StepArtifact"],
-        external_artifacts: Dict[str, "ExternalArtifact"],
+        external_artifacts: Dict[
+            str, Union["ExternalArtifact", "ArtifactVersionResponse"]
+        ],
         model_artifacts_or_metadata: Dict[str, "ModelVersionDataLazyLoader"],
         client_lazy_loaders: Dict[str, "ClientLazyLoader"],
         parameters: Dict[str, Any],
@@ -64,59 +68,9 @@ class StepInvocation:
         self.client_lazy_loaders = client_lazy_loaders
         self.parameters = parameters
         self.default_parameters = default_parameters
-        self.invocation_upstream_steps = upstream_steps
+        self.upstream_steps = upstream_steps
         self.pipeline = pipeline
 
-    @property
-    def upstream_steps(self) -> Set[str]:
-        """The upstream steps of the invocation.
-
-        Returns:
-            The upstream steps of the invocation.
-        """
-        return self.invocation_upstream_steps.union(
-            self._get_and_validate_step_upstream_steps()
-        )
-
-    def _get_and_validate_step_upstream_steps(self) -> Set[str]:
-        """Validates the upstream steps defined on the step instance.
-
-        This is only allowed in legacy pipelines when calling `step.after(...)`
-        and we need to make sure that both the upstream and downstream steps
-        of such a relationship are only invoked once inside a pipeline.
-
-        Returns:
-            The upstream steps defined on the step instance.
-        """
-
-        def _verify_single_invocation(step: "BaseStep") -> str:
-            invocations = {
-                invocation
-                for invocation in self.pipeline.invocations.values()
-                if invocation.step is step
-            }
-            if len(invocations) > 1:
-                raise RuntimeError(
-                    "Setting upstream steps for a step using "
-                    "`step_1.after(step_2)` is not allowed in combination "
-                    "with calling one of the two steps multiple times."
-                )
-            return invocations.pop().id
-
-        if self.step.upstream_steps:
-            # If the step has upstream steps, make sure it only got invoked once
-            _verify_single_invocation(step=self.step)
-
-        upstream_steps = set()
-
-        for upstream_step in self.step.upstream_steps:
-            upstream_step_invocation_id = _verify_single_invocation(
-                step=upstream_step
-            )
-            upstream_steps.add(upstream_step_invocation_id)
-
-        return upstream_steps
-
     def finalize(self, parameters_to_ignore: Set[str]) -> "StepConfiguration":
         """Finalizes a step invocation.
 
@@ -134,9 +88,6 @@ class StepInvocation:
             ExternalArtifactConfiguration,
         )
 
-        # Validate the upstream steps for legacy .after() calls
-        self._get_and_validate_step_upstream_steps()
-
         parameters_to_apply = {
             key: value
             for key, value in self.parameters.items()
@@ -154,9 +105,14 @@ class StepInvocation:
 
         external_artifacts: Dict[str, ExternalArtifactConfiguration] = {}
         for key, artifact in self.external_artifacts.items():
-            if artifact.value is not None:
-                artifact.upload_by_value()
-            external_artifacts[key] = artifact.config
+            if isinstance(artifact, ArtifactVersionResponse):
+                external_artifacts[key] = ExternalArtifactConfiguration(
+                    id=artifact.id
+                )
+            else:
+                if artifact.value is not None:
+                    artifact.upload_by_value()
+                external_artifacts[key] = artifact.config
 
         return self.step._finalize_configuration(
             input_artifacts=self.input_artifacts,

zenml/steps/utils.py

@@ -30,8 +30,7 @@ from zenml.enums import ExecutionStatus, MetadataResourceTypes
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
-from zenml.new.steps.step_context import get_step_context
-from zenml.steps.step_output import Output
+from zenml.steps.step_context import get_step_context
 from zenml.utils import settings_utils, source_code_utils, typing_utils
 
 if TYPE_CHECKING:
@@ -129,26 +128,7 @@ def parse_return_type_annotations(
         else:
             return_annotation = Any
 
-    # Return type annotated using deprecated `Output(...)`
-    if isinstance(return_annotation, Output):
-        logger.warning(
-            "Using the `Output` class to define the outputs of your steps is "
-            "deprecated. You should instead use the standard Python way of "
-            "type annotating your functions. Check out our documentation "
-            "https://docs.zenml.io/how-to/build-pipelines/step-output-typing-and-annotation "
-            "for more information on how to assign custom names to your step "
-            "outputs."
-        )
-        return {
-            output_name: OutputSignature(
-                resolved_annotation=resolve_type_annotation(output_type),
-                artifact_config=None,
-                has_custom_name=True,
-            )
-            for output_name, output_type in return_annotation.items()
-        }
-
-    elif typing_utils.get_origin(return_annotation) is tuple:
+    if typing_utils.get_origin(return_annotation) is tuple:
         requires_multiple_artifacts = has_tuple_return(func)
         if requires_multiple_artifacts:
             output_signature: Dict[str, Any] = {}
@@ -517,7 +497,7 @@ def run_as_single_step_pipeline(
     """
     from zenml import ExternalArtifact, pipeline
     from zenml.config.base_settings import BaseSettings
-    from zenml.new.pipelines.run_utils import (
+    from zenml.pipelines.run_utils import (
         wait_for_pipeline_run_to_finish,
     )
 
@@ -558,6 +538,7 @@ def run_as_single_step_pipeline(
         __step(**inputs)
 
     run = single_step_pipeline.with_options(unlisted=True)()
+    assert run
     run = wait_for_pipeline_run_to_finish(run.id)
 
     if run.status != ExecutionStatus.COMPLETED:

zenml/utils/dashboard_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Utility class to help with interacting with the dashboard."""
 
-from typing import Optional, Tuple
+from typing import Optional
 from uuid import UUID
 
 from zenml import constants
@@ -24,28 +24,13 @@ from zenml.logger import get_logger
 from zenml.models import (
     ComponentResponse,
     PipelineRunResponse,
-    ServerModel,
+    ServerDeploymentType,
     StackResponse,
 )
 
 logger = get_logger(__name__)
 
 
-def is_cloud_server(server_info: ServerModel) -> bool:
-    """Checks whether the server info refers to a ZenML Pro server.
-
-    Args:
-        server_info: Info of the server.
-
-    Returns:
-        Whether the server info refers to a ZenML Pro server.
-    """
-    return (
-        "organization_id" in server_info.metadata
-        and "cloud.zenml.io" in server_info.dashboard_url
-    )
-
-
 def get_cloud_dashboard_url() -> Optional[str]:
     """Get the base url of the cloud dashboard if the server is a cloud tenant.
 
@@ -57,36 +42,31 @@ def get_cloud_dashboard_url() -> Optional[str]:
     if client.zen_store.type == StoreType.REST:
         server_info = client.zen_store.get_store_info()
 
-        if is_cloud_server(server_info):
+        if server_info.deployment_type == ServerDeploymentType.CLOUD:
             return server_info.dashboard_url
 
     return None
 
 
-def get_server_dashboard_url() -> Tuple[Optional[str], bool]:
+def get_server_dashboard_url() -> Optional[str]:
     """Get the base url of the dashboard deployed by the server.
 
     Returns:
-        The server dashboard url and whether the dashboard is the legacy
-        dashboard or not.
+        The server dashboard url.
     """
     client = Client()
 
     if client.zen_store.type == StoreType.REST:
         server_info = client.zen_store.get_store_info()
-        if server_info.use_legacy_dashboard:
-            suffix = f"{constants.WORKSPACES}/{client.active_workspace.name}"
-        else:
-            suffix = ""
 
         if server_info.server_url:
             url = server_info.server_url
         else:
             url = client.zen_store.url
 
-        return url + suffix, server_info.use_legacy_dashboard
+        return url
 
-    return None, False
+    return None
 
 
 def get_stack_url(stack: StackResponse) -> Optional[str]:
@@ -98,13 +78,10 @@ def get_stack_url(stack: StackResponse) -> Optional[str]:
     Returns:
         the URL to the stack if the dashboard is available, else None.
     """
-    base_url, is_legacy_dashboard = get_server_dashboard_url()
+    base_url = get_server_dashboard_url()
 
     if base_url:
-        if is_legacy_dashboard:
-            return base_url + f"{constants.STACKS}/{stack.id}/configuration"
-        else:
-            return base_url + constants.STACKS
+        return base_url + constants.STACKS
 
     return None
 
@@ -118,16 +95,10 @@ def get_component_url(component: ComponentResponse) -> Optional[str]:
     Returns:
         the URL to the component if the dashboard is available, else None.
     """
-    base_url, is_legacy_dashboard = get_server_dashboard_url()
+    base_url = get_server_dashboard_url()
 
     if base_url:
-        if is_legacy_dashboard:
-            return (
-                base_url
-                + f"{constants.STACK_COMPONENTS}/{component.type.value}/{component.id}/configuration"
-            )
-        else:
-            return base_url + constants.STACKS
+        return base_url + constants.STACKS
 
     return None
 
@@ -145,15 +116,9 @@ def get_run_url(run: PipelineRunResponse) -> Optional[str]:
     if cloud_url:
         return f"{cloud_url}{constants.RUNS}/{run.id}"
 
-    dashboard_url, is_legacy_dashboard = get_server_dashboard_url()
+    dashboard_url = get_server_dashboard_url()
     if dashboard_url:
-        if is_legacy_dashboard:
-            if run.pipeline:
-                return f"{dashboard_url}{constants.PIPELINES}/{run.pipeline.id}{constants.RUNS}/{run.id}/dag"
-            else:
-                return f"{dashboard_url}/all-runs/{run.id}/dag"
-        else:
-            return f"{dashboard_url}{constants.RUNS}/{run.id}"
+        return f"{dashboard_url}{constants.RUNS}/{run.id}"
 
     return None
 
@@ -167,9 +132,9 @@ def get_model_version_url(model_version_id: UUID) -> Optional[str]:
     Returns:
         the URL to the model version if the dashboard is available, else None.
     """
-    base_url = get_cloud_dashboard_url()
-    if base_url:
-        return f"{base_url}/model-versions/{str(model_version_id)}"
+    cloud_url = get_cloud_dashboard_url()
+    if cloud_url:
+        return f"{cloud_url}/model-versions/{str(model_version_id)}"
 
     return None
 

zenml/utils/materializer_utils.py

@@ -41,7 +41,7 @@ def select_materializer(
         for materializer_class in materializer_classes:
             if class_ in materializer_class.ASSOCIATED_TYPES:
                 return materializer_class
-            elif not fallback and materializer_class.can_handle_type(class_):
+            elif not fallback and materializer_class.can_save_type(class_):
                 fallback = materializer_class
 
     if fallback: