zenml-nightly 0.70.0.dev20241127__py3-none-any.whl → 0.70.0.dev20241129__py3-none-any.whl

zenml/orchestrators/utils.py

@@ -40,7 +40,9 @@ if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
 
 
-def get_orchestrator_run_name(pipeline_name: str) -> str:
+def get_orchestrator_run_name(
+    pipeline_name: str, max_length: Optional[int] = None
+) -> str:
     """Gets an orchestrator run name.
 
     This run name is not the same as the ZenML run name but can instead be
@@ -48,11 +50,31 @@ def get_orchestrator_run_name(pipeline_name: str) -> str:
 
     Args:
         pipeline_name: Name of the pipeline that will run.
+        max_length: Maximum length of the generated name.
+
+    Raises:
+        ValueError: If the max length is below 8 characters.
 
     Returns:
         The orchestrator run name.
     """
-    return f"{pipeline_name}_{random.Random().getrandbits(128):032x}"
+    suffix_length = 32
+    pipeline_name = f"{pipeline_name}_"
+
+    if max_length:
+        if max_length < 8:
+            raise ValueError(
+                "Maximum length for orchestrator run name must be 8 or above."
+            )
+
+        # Make sure we always have a certain suffix to guarantee no overlap
+        # with other runs
+        suffix_length = min(32, max(8, max_length - len(pipeline_name)))
+        pipeline_name = pipeline_name[: (max_length - suffix_length)]
+
+    suffix = "".join(random.choices("0123456789abcdef", k=suffix_length))
+
+    return f"{pipeline_name}{suffix}"
 
 
 def is_setting_enabled(
@@ -174,11 +196,12 @@ def get_config_environment_vars(
     return environment_vars
 
 
-def get_run_name(run_name_template: str) -> str:
+def get_run_name(run_name_template: str, substitutions: Dict[str, 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.
+        substitutions: The substitutions to use in the template.
 
     Raises:
         ValueError: If the run name is empty.
@@ -186,7 +209,9 @@ def get_run_name(run_name_template: str) -> str:
     Returns:
         The run name derived from the template.
     """
-    run_name = format_name_template(run_name_template)
+    run_name = format_name_template(
+        run_name_template, substitutions=substitutions
+    )
 
     if run_name == "":
         raise ValueError("Empty run names are not allowed.")

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

@@ -66,16 +66,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=get_run_name(
+            run_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
@@ -94,7 +102,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 +237,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:

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."""

zenml/utils/string_utils.py

@@ -14,10 +14,10 @@
 """Utils for strings."""
 
 import base64
-import datetime
 import functools
 import random
 import string
+from datetime import datetime
 from typing import Any, Callable, Dict, TypeVar, cast
 
 from pydantic import BaseModel
@@ -147,7 +147,7 @@ def validate_name(model: BaseModel) -> None:
 
 def format_name_template(
     name_template: str,
-    **kwargs: str,
+    substitutions: Dict[str, str],
 ) -> str:
     """Formats a name template with the given arguments.
 
@@ -157,20 +157,38 @@ def format_name_template(
 
     Args:
         name_template: The name template to format.
-        **kwargs: The arguments to replace in the template.
+        substitutions: A dictionary of substitutions to use in the template.
 
     Returns:
         The formatted name template.
+
+    Raises:
+        KeyError: If a key in template is missing in the kwargs.
     """
-    kwargs["date"] = kwargs.get(
-        "date",
-        datetime.datetime.now(datetime.timezone.utc).strftime("%Y_%m_%d"),
-    )
-    kwargs["time"] = kwargs.get(
-        "time",
-        datetime.datetime.now(datetime.timezone.utc).strftime("%H_%M_%S_%f"),
-    )
-    return name_template.format(**kwargs)
+    if ("date" not in substitutions and "{date}" in name_template) or (
+        "time" not in substitutions and "{time}" in name_template
+    ):
+        from zenml import get_step_context
+
+        try:
+            pr = get_step_context().pipeline_run
+            start_time = pr.start_time
+            substitutions.update(pr.config.substitutions)
+        except RuntimeError:
+            start_time = None
+
+        if start_time is None:
+            start_time = datetime.utcnow()
+        substitutions.setdefault("date", start_time.strftime("%Y_%m_%d"))
+        substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
+
+    try:
+        return name_template.format(**substitutions)
+    except KeyError as e:
+        raise KeyError(
+            f"Could not format the name template `{name_template}`. "
+            f"Missing key: {e}"
+        )
 
 
 def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:

zenml/utils/visualization_utils.py

@@ -13,9 +13,10 @@
 #  permissions and limitations under the License.
 """Utility functions for dashboard visualizations."""
 
+import json
 from typing import TYPE_CHECKING, Optional
 
-from IPython.core.display import HTML, Image, Markdown, display
+from IPython.core.display import HTML, JSON, Image, Markdown, display
 
 from zenml.artifacts.utils import load_artifact_visualization
 from zenml.enums import VisualizationType
@@ -63,6 +64,8 @@ def visualize_artifact(
             assert isinstance(visualization.value, str)
             table = format_csv_visualization_as_html(visualization.value)
             display(HTML(table))
+        elif visualization.type == VisualizationType.JSON:
+            display(JSON(json.loads(visualization.value)))
         else:
             display(visualization.value)
 

zenml/zen_server/template_execution/utils.py

@@ -363,6 +363,7 @@ def deployment_request_from_template(
                     "external_input_artifacts",
                     "model_artifacts_or_metadata",
                     "client_lazy_loaders",
+                    "substitutions",
                     "outputs",
                 }
             ),

zenml/zen_stores/schemas/artifact_schemas.py

@@ -352,8 +352,9 @@ class ArtifactVersionSchema(BaseSchema, table=True):
                 producer_step_run_id = step_run.original_step_run_id
 
         # Create the body of the model
+        artifact = self.artifact.to_model()
         body = ArtifactVersionResponseBody(
-            artifact=self.artifact.to_model(),
+            artifact=artifact,
             version=self.version or str(self.version_number),
             user=self.user.to_model() if self.user else None,
             uri=self.uri,

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -284,6 +284,10 @@ class PipelineRunSchema(NamedSchema, table=True):
             deployment = self.deployment.to_model()
 
             config = deployment.pipeline_configuration
+            new_substitutions = config._get_full_substitutions(self.start_time)
+            config = config.model_copy(
+                update={"substitutions": new_substitutions}
+            )
             client_environment = deployment.client_environment
 
             stack = deployment.stack
@@ -323,9 +327,11 @@ class PipelineRunSchema(NamedSchema, table=True):
             build=build,
             schedule=schedule,
             code_reference=code_reference,
-            trigger_execution=self.trigger_execution.to_model()
-            if self.trigger_execution
-            else None,
+            trigger_execution=(
+                self.trigger_execution.to_model()
+                if self.trigger_execution
+                else None
+            ),
             created=self.created,
             updated=self.updated,
             deployment_id=self.deployment_id,
@@ -344,6 +350,10 @@ class PipelineRunSchema(NamedSchema, table=True):
 
             steps = {step.name: step.to_model() for step in self.step_runs}
 
+            steps_substitutions = {
+                step_name: step.config.substitutions
+                for step_name, step in steps.items()
+            }
             metadata = PipelineRunResponseMetadata(
                 workspace=self.workspace.to_model(),
                 run_metadata=run_metadata,
@@ -361,6 +371,7 @@ class PipelineRunSchema(NamedSchema, table=True):
                 if self.deployment
                 else None,
                 is_templatable=is_templatable,
+                steps_substitutions=steps_substitutions,
             )
 
         resources = None

zenml/zen_stores/schemas/step_run_schemas.py

@@ -23,6 +23,7 @@ from sqlalchemy import TEXT, Column, String
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
 from sqlmodel import Field, Relationship, SQLModel
 
+from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.step_configurations import Step
 from zenml.constants import MEDIUMTEXT_MAX_LENGTH
 from zenml.enums import (
@@ -163,6 +164,9 @@ class StepRunSchema(NamedSchema, table=True):
             "primaryjoin": "StepRunParentsSchema.child_id == StepRunSchema.id",
         },
     )
+    pipeline_run: "PipelineRunSchema" = Relationship(
+        back_populates="step_runs"
+    )
     model_version: "ModelVersionSchema" = Relationship(
         back_populates="step_runs",
     )
@@ -248,6 +252,21 @@ class StepRunSchema(NamedSchema, table=True):
                 full_step_config = Step.model_validate(
                     step_configuration[self.name]
                 )
+                new_substitutions = (
+                    full_step_config.config._get_full_substitutions(
+                        PipelineConfiguration.model_validate_json(
+                            self.deployment.pipeline_configuration
+                        ),
+                        self.pipeline_run.start_time,
+                    )
+                )
+                full_step_config = full_step_config.model_copy(
+                    update={
+                        "config": full_step_config.config.model_copy(
+                            update={"substitutions": new_substitutions}
+                        )
+                    }
+                )
             elif not self.step_configuration:
                 raise ValueError(
                     f"Unable to load the configuration for step `{self.name}` from the"

zenml/zen_stores/sql_zen_store.py

@@ -2726,7 +2726,9 @@ class SqlZenStore(BaseZenStore):
     # -------------------- Artifact Versions --------------------
 
     def _get_or_create_artifact_for_name(
-        self, name: str, has_custom_name: bool
+        self,
+        name: str,
+        has_custom_name: bool,
     ) -> ArtifactSchema:
         """Get or create an artifact with a specific name.
 
@@ -2747,7 +2749,8 @@ class SqlZenStore(BaseZenStore):
                 try:
                     with session.begin_nested():
                         artifact_request = ArtifactRequest(
-                            name=name, has_custom_name=has_custom_name
+                            name=name,
+                            has_custom_name=has_custom_name,
                         )
                         artifact = ArtifactSchema.from_request(
                             artifact_request
@@ -8182,12 +8185,12 @@ class SqlZenStore(BaseZenStore):
                 )
 
             # Save output artifact IDs into the database.
-            for output_name, artifact_version_ids in step_run.outputs.items():
+            for name, artifact_version_ids in step_run.outputs.items():
                 for artifact_version_id in artifact_version_ids:
                     self._set_run_step_output_artifact(
                         step_run_id=step_schema.id,
                         artifact_version_id=artifact_version_id,
-                        name=output_name,
+                        name=name,
                         session=session,
                     )
 
@@ -8291,13 +8294,14 @@ class SqlZenStore(BaseZenStore):
             session.add(existing_step_run)
 
             # Update the artifacts.
-            for name, artifact_version_id in step_run_update.outputs.items():
-                self._set_run_step_output_artifact(
-                    step_run_id=step_run_id,
-                    artifact_version_id=artifact_version_id,
-                    name=name,
-                    session=session,
-                )
+            for name, artifact_version_ids in step_run_update.outputs.items():
+                for artifact_version_id in artifact_version_ids:
+                    self._set_run_step_output_artifact(
+                        step_run_id=step_run_id,
+                        artifact_version_id=artifact_version_id,
+                        name=name,
+                        session=session,
+                    )
 
             # Update loaded artifacts.
             for (

{zenml_nightly-0.70.0.dev20241127.dist-info → zenml_nightly-0.70.0.dev20241129.dist-info}/METADATA

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