zenml-nightly 0.70.0.dev20241128__py3-none-any.whl → 0.70.0.dev20241130__py3-none-any.whl

zenml/orchestrators/step_runner.py

@@ -56,7 +56,7 @@ from zenml.steps.utils import (
     parse_return_type_annotations,
     resolve_type_annotation,
 )
-from zenml.utils import materializer_utils, source_utils
+from zenml.utils import materializer_utils, source_utils, string_utils
 from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
@@ -152,6 +152,15 @@ class StepRunner:
                 func=step_instance.entrypoint
             )
 
+            self._evaluate_artifact_names_in_collections(
+                step_run,
+                output_annotations,
+                [
+                    output_artifact_uris,
+                    output_materializers,
+                ],
+            )
+
             self._stack.prepare_step_run(info=step_run_info)
 
             # Initialize the step context singleton
@@ -257,7 +266,9 @@ class StepRunner:
 
             # Update the status and output artifacts of the step run.
             output_artifact_ids = {
-                output_name: artifact.id
+                output_name: [
+                    artifact.id,
+                ]
                 for output_name, artifact in output_artifacts.items()
             }
             publish_successful_step_run(
@@ -265,6 +276,32 @@ class StepRunner:
                 output_artifact_ids=output_artifact_ids,
             )
 
+    def _evaluate_artifact_names_in_collections(
+        self,
+        step_run: "StepRunResponse",
+        output_annotations: Dict[str, OutputSignature],
+        collections: List[Dict[str, Any]],
+    ) -> None:
+        """Evaluates the artifact names in the collections.
+
+        Args:
+            step_run: The step run.
+            output_annotations: The output annotations of the step function
+                (also evaluated).
+            collections: The collections to evaluate.
+        """
+        collections.append(output_annotations)
+        for k, v in list(output_annotations.items()):
+            name = k
+            if v.artifact_config and v.artifact_config.name:
+                name = string_utils.format_name_template(
+                    v.artifact_config.name,
+                    substitutions=step_run.config.substitutions,
+                )
+
+            for d in collections:
+                d[name] = d.pop(k)
+
     def _load_step(self) -> "BaseStep":
         """Load the step instance.
 

zenml/orchestrators/utils.py

@@ -32,7 +32,6 @@ from zenml.constants import (
 from zenml.enums import AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
 from zenml.stack import StackComponent
-from zenml.utils.string_utils import format_name_template
 
 logger = get_logger(__name__)
 
@@ -196,26 +195,6 @@ def get_config_environment_vars(
     return environment_vars
 
 
-def get_run_name(run_name_template: str) -> str:
-    """Fill out the run name template to get a complete run name.
-
-    Args:
-        run_name_template: The run name template to fill out.
-
-    Raises:
-        ValueError: If the run name is empty.
-
-    Returns:
-        The run name derived from the template.
-    """
-    run_name = format_name_template(run_name_template)
-
-    if run_name == "":
-        raise ValueError("Empty run names are not allowed.")
-
-    return run_name
-
-
 class register_artifact_store_filesystem:
     """Context manager for the artifact_store/filesystem_registry dependency.
 

zenml/pipelines/pipeline_decorator.py

@@ -55,6 +55,7 @@ def pipeline(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
+    substitutions: Optional[Dict[str, str]] = None,
 ) -> Callable[["F"], "Pipeline"]: ...
 
 
@@ -71,6 +72,7 @@ def pipeline(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
+    substitutions: Optional[Dict[str, str]] = None,
 ) -> Union["Pipeline", Callable[["F"], "Pipeline"]]:
     """Decorator to create a pipeline.
 
@@ -91,6 +93,7 @@ def pipeline(
             function with no arguments, or a source path to such a function
             (e.g. `module.my_function`).
         model: configuration of the model in the Model Control Plane.
+        substitutions: Extra placeholders to use in the name templates.
 
     Returns:
         A pipeline instance.
@@ -111,6 +114,7 @@ def pipeline(
             on_success=on_success,
             model=model,
             entrypoint=func,
+            substitutions=substitutions,
         )
 
         p.__doc__ = func.__doc__

zenml/pipelines/pipeline_definition.py

@@ -135,6 +135,7 @@ class Pipeline:
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
+        substitutions: Optional[Dict[str, str]] = None,
     ) -> None:
         """Initializes a pipeline.
 
@@ -157,6 +158,7 @@ class Pipeline:
                 be a function with no arguments, or a source path to such a
                 function (e.g. `module.my_function`).
             model: configuration of the model in the Model Control Plane.
+            substitutions: Extra placeholders to use in the name templates.
         """
         self._invocations: Dict[str, StepInvocation] = {}
         self._run_args: Dict[str, Any] = {}
@@ -177,6 +179,7 @@ class Pipeline:
                 on_failure=on_failure,
                 on_success=on_success,
                 model=model,
+                substitutions=substitutions,
             )
         self.entrypoint = entrypoint
         self._parameters: Dict[str, Any] = {}
@@ -297,6 +300,7 @@ class Pipeline:
         model: Optional["Model"] = None,
         parameters: Optional[Dict[str, Any]] = None,
         merge: bool = True,
+        substitutions: Optional[Dict[str, str]] = None,
     ) -> Self:
         """Configures the pipeline.
 
@@ -333,6 +337,7 @@ class Pipeline:
                 method for an example.
             model: configuration of the model version in the Model Control Plane.
             parameters: input parameters for the pipeline.
+            substitutions: Extra placeholders to use in the name templates.
 
         Returns:
             The pipeline instance that this method was called on.
@@ -365,6 +370,7 @@ class Pipeline:
                 "success_hook_source": success_hook_source,
                 "model": model,
                 "parameters": parameters,
+                "substitutions": substitutions,
             }
         )
         if not self.__suppress_warnings_flag__:
@@ -577,6 +583,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         config_path: Optional[str] = None,
         unlisted: bool = False,
         prevent_build_reuse: bool = False,
+        skip_schedule_registration: bool = False,
     ) -> PipelineDeploymentResponse:
         """Create a pipeline deployment.
 
@@ -603,6 +610,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 to any pipeline).
             prevent_build_reuse: DEPRECATED: Use
                 `DockerSettings.prevent_build_reuse` instead.
+            skip_schedule_registration: Whether to skip schedule registration.
 
         Returns:
             The pipeline deployment.
@@ -643,7 +651,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         stack.validate()
 
         schedule_id = None
-        if schedule:
+        if schedule and not skip_schedule_registration:
             if not stack.orchestrator.config.is_schedulable:
                 raise ValueError(
                     f"Stack {stack.name} does not support scheduling. "
@@ -656,7 +664,8 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 schedule_name = schedule.name
             else:
                 schedule_name = format_name_template(
-                    deployment.run_name_template
+                    deployment.run_name_template,
+                    substitutions=deployment.pipeline_configuration.substitutions,
                 )
             components = Client().active_stack_model.components
             orchestrator = components[StackComponentType.ORCHESTRATOR][0]
@@ -1438,7 +1447,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
             The created run template.
         """
         self._prepare_if_possible()
-        deployment = self._create_deployment(**self._run_args)
+        deployment = self._create_deployment(
+            **self._run_args, skip_schedule_registration=True
+        )
 
         return Client().create_run_template(
             name=name, deployment_id=deployment.id, **kwargs

zenml/pipelines/run_utils.py

@@ -23,9 +23,8 @@ from zenml.models import (
     StackResponse,
 )
 from zenml.orchestrators.publish_utils import publish_failed_pipeline_run
-from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Flavor, Stack
-from zenml.utils import code_utils, notebook_utils, source_utils
+from zenml.utils import code_utils, notebook_utils, source_utils, string_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
@@ -66,16 +65,21 @@ def create_placeholder_run(
 
     if deployment.schedule:
         return None
-
+    start_time = datetime.utcnow()
     run_request = PipelineRunRequest(
-        name=get_run_name(run_name_template=deployment.run_name_template),
+        name=string_utils.format_name_template(
+            name_template=deployment.run_name_template,
+            substitutions=deployment.pipeline_configuration._get_full_substitutions(
+                start_time
+            ),
+        ),
         # We set the start time on the placeholder run already to
         # make it consistent with the {time} placeholder in the
         # run name. This means the placeholder run will usually
         # have longer durations than scheduled runs, as for them
         # the start_time is only set once the first step starts
         # running.
-        start_time=datetime.utcnow(),
+        start_time=start_time,
         orchestrator_run_id=None,
         user=deployment.user.id,
         workspace=deployment.workspace.id,

zenml/steps/base_step.py

@@ -116,6 +116,7 @@ class BaseStep:
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
         retry: Optional[StepRetryConfig] = None,
+        substitutions: Optional[Dict[str, str]] = None,
     ) -> None:
         """Initializes a step.
 
@@ -144,11 +145,13 @@ class BaseStep:
                 function (e.g. `module.my_function`).
             model: configuration of the model version in the Model Control Plane.
             retry: Configuration for retrying the step in case of failure.
+            substitutions: Extra placeholders to use in the name template.
         """
         from zenml.config.step_configurations import PartialStepConfiguration
 
         self.entrypoint_definition = validate_entrypoint_function(
-            self.entrypoint, reserved_arguments=["after", "id"]
+            self.entrypoint,
+            reserved_arguments=["after", "id"],
         )
 
         name = name or self.__class__.__name__
@@ -203,6 +206,7 @@ class BaseStep:
             on_success=on_success,
             model=model,
             retry=retry,
+            substitutions=substitutions,
         )
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
@@ -595,6 +599,7 @@ class BaseStep:
         model: Optional["Model"] = None,
         merge: bool = True,
         retry: Optional[StepRetryConfig] = None,
+        substitutions: Optional[Dict[str, str]] = None,
     ) -> T:
         """Configures the step.
 
@@ -637,6 +642,7 @@ class BaseStep:
                 overwrite all existing ones. See the general description of this
                 method for an example.
             retry: Configuration for retrying the step in case of failure.
+            substitutions: Extra placeholders to use in the name template.
 
         Returns:
             The step instance that this method was called on.
@@ -701,6 +707,7 @@ class BaseStep:
                 "success_hook_source": success_hook_source,
                 "model": model,
                 "retry": retry,
+                "substitutions": substitutions,
             }
         )
         config = StepConfigurationUpdate(**values)
@@ -725,6 +732,7 @@ class BaseStep:
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
         merge: bool = True,
+        substitutions: Optional[Dict[str, str]] = None,
     ) -> "BaseStep":
         """Copies the step and applies the given configurations.
 
@@ -756,6 +764,7 @@ class BaseStep:
                 configurations. If `False` the given configurations will
                 overwrite all existing ones. See the general description of this
                 method for an example.
+            substitutions: Extra placeholders for the step name.
 
         Returns:
             The copied step instance.
@@ -776,6 +785,7 @@ class BaseStep:
             on_success=on_success,
             model=model,
             merge=merge,
+            substitutions=substitutions,
         )
         return step_copy
 

zenml/steps/entrypoint_function_utils.py

@@ -211,7 +211,8 @@ class EntrypointFunctionDefinition(NamedTuple):
 
 
 def validate_entrypoint_function(
-    func: Callable[..., Any], reserved_arguments: Sequence[str] = ()
+    func: Callable[..., Any],
+    reserved_arguments: Sequence[str] = (),
 ) -> EntrypointFunctionDefinition:
     """Validates a step entrypoint function.
 
@@ -258,7 +259,8 @@ def validate_entrypoint_function(
         inputs[key] = parameter
 
     outputs = parse_return_type_annotations(
-        func=func, enforce_type_annotations=ENFORCE_TYPE_ANNOTATIONS
+        func=func,
+        enforce_type_annotations=ENFORCE_TYPE_ANNOTATIONS,
     )
 
     return EntrypointFunctionDefinition(

zenml/steps/step_decorator.py

@@ -73,6 +73,7 @@ def step(
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
     retry: Optional["StepRetryConfig"] = None,
+    substitutions: Optional[Dict[str, str]] = None,
 ) -> Callable[["F"], "BaseStep"]: ...
 
 
@@ -93,6 +94,7 @@ def step(
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
     retry: Optional["StepRetryConfig"] = None,
+    substitutions: Optional[Dict[str, str]] = None,
 ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
     """Decorator to create a ZenML step.
 
@@ -124,6 +126,7 @@ def step(
             (e.g. `module.my_function`).
         model: configuration of the model in the Model Control Plane.
         retry: configuration of step retry in case of step failure.
+        substitutions: Extra placeholders for the step name.
 
     Returns:
         The step instance.
@@ -157,6 +160,7 @@ def step(
             on_success=on_success,
             model=model,
             retry=retry,
+            substitutions=substitutions,
         )
 
         return step_instance

zenml/steps/utils.py

@@ -18,7 +18,15 @@ import ast
 import contextlib
 import inspect
 import textwrap
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import BaseModel
@@ -34,6 +42,7 @@ from zenml.enums import (
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
+from zenml.models import RunMetadataResource
 from zenml.steps.step_context import get_step_context
 from zenml.utils import settings_utils, source_code_utils, typing_utils
 
@@ -94,7 +103,8 @@ def get_args(obj: Any) -> Tuple[Any, ...]:
 
 
 def parse_return_type_annotations(
-    func: Callable[..., Any], enforce_type_annotations: bool = False
+    func: Callable[..., Any],
+    enforce_type_annotations: bool = False,
 ) -> Dict[str, OutputSignature]:
     """Parse the return type annotation of a step function.
 
@@ -228,9 +238,12 @@ def get_artifact_config_from_annotation_metadata(
 
     error_message = (
         "Artifact annotation should only contain two elements: the artifact "
-        "type, and either an output name or an `ArtifactConfig`, e.g.: "
-        "`Annotated[int, 'output_name']` or "
-        "`Annotated[int, ArtifactConfig(name='output_name'), ...]`."
+        "type, and one of the following: { a static or dynamic name || "
+        "an `ArtifactConfig` }, e.g.: "
+        "`Annotated[int, 'output_name']` || "
+        "`Annotated[int, 'output_{placeholder}']` || "
+        "`Annotated[int, ArtifactConfig(name='output_{placeholder}')]` ||"
+        "`Annotated[int, ArtifactConfig(name='output_name')]`."
     )
 
     if len(metadata) > 2:
@@ -477,8 +490,11 @@ def log_step_metadata(
         step_run_id = pipeline_run.steps[step_name].id
     client.create_run_metadata(
         metadata=metadata,
-        resource_id=step_run_id,
-        resource_type=MetadataResourceTypes.STEP_RUN,
+        resources=[
+            RunMetadataResource(
+                id=step_run_id, type=MetadataResourceTypes.STEP_RUN
+            )
+        ],
     )
 
 

zenml/types.py

@@ -33,3 +33,7 @@ class MarkdownString(str):
 
 class CSVString(str):
     """Special string class to indicate a CSV string."""
+
+
+class JSONString(str):
+    """Special string class to indicate a JSON string."""