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

zenml/VERSION

@@ -1 +1 @@
-0.70.0.dev20241128
+0.70.0.dev20241129

zenml/artifacts/artifact_config.py

@@ -21,6 +21,7 @@ from zenml.enums import ArtifactType
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.string_utils import format_name_template
 
 logger = get_logger(__name__)
 
@@ -45,7 +46,13 @@ class ArtifactConfig(BaseModel):
     ```
 
     Attributes:
-        name: The name of the artifact.
+        name: The name of the artifact:
+            - static string e.g. "name"
+            - dynamic string e.g. "name_{date}_{time}_{custom_placeholder}"
+            If you use any placeholders besides `date` and `time`,
+            you need to provide the values for them in the `substitutions`
+            argument of the step decorator or the `substitutions` argument
+            of `with_options` of the step.
         version: The version of the artifact.
         tags: The tags of the artifact.
         run_metadata: Metadata to add to the artifact.
@@ -111,3 +118,16 @@ class ArtifactConfig(BaseModel):
             data.setdefault("artifact_type", ArtifactType.SERVICE)
 
         return data
+
+    def _evaluated_name(self, substitutions: Dict[str, str]) -> Optional[str]:
+        """Evaluated name of the artifact.
+
+        Args:
+            substitutions: Extra placeholders to use in the name template.
+
+        Returns:
+            The evaluated name of the artifact.
+        """
+        if self.name:
+            return format_name_template(self.name, substitutions=substitutions)
+        return self.name

zenml/artifacts/utils.py

@@ -689,7 +689,11 @@ def _link_artifact_version_to_the_step_and_model(
         client.zen_store.update_run_step(
             step_run_id=step_run.id,
             step_run_update=StepRunUpdate(
-                outputs={artifact_version.artifact.name: artifact_version.id}
+                outputs={
+                    artifact_version.artifact.name: [
+                        artifact_version.id,
+                    ]
+                }
             ),
         )
         error_message = "model"

zenml/cli/pipeline.py

@@ -315,6 +315,86 @@ def run_pipeline(
         pipeline_instance()
 
 
+@pipeline.command(
+    "create-run-template",
+    help="Create a run template for a pipeline. The SOURCE argument needs to "
+    "be an importable source path resolving to a ZenML pipeline instance, e.g. "
+    "`my_module.my_pipeline_instance`.",
+)
+@click.argument("source")
+@click.option(
+    "--name",
+    "-n",
+    type=str,
+    required=True,
+    help="Name for the template",
+)
+@click.option(
+    "--config",
+    "-c",
+    "config_path",
+    type=click.Path(exists=True, dir_okay=False),
+    required=False,
+    help="Path to configuration file for the build.",
+)
+@click.option(
+    "--stack",
+    "-s",
+    "stack_name_or_id",
+    type=str,
+    required=False,
+    help="Name or ID of the stack to use for the build.",
+)
+def create_run_template(
+    source: str,
+    name: str,
+    config_path: Optional[str] = None,
+    stack_name_or_id: Optional[str] = None,
+) -> None:
+    """Create a run template for a pipeline.
+
+    Args:
+        source: Importable source resolving to a pipeline instance.
+        name: Name of the run template.
+        config_path: Path to pipeline configuration file.
+        stack_name_or_id: Name or ID of the stack for which the template should
+            be created.
+    """
+    if not Client().root:
+        cli_utils.warning(
+            "You're running the `zenml pipeline create-run-template` command "
+            "without a ZenML repository. Your current working directory will "
+            "be used as the source root relative to which the registered step "
+            "classes will be resolved. To silence this warning, run `zenml "
+            "init` at your source code root."
+        )
+
+    try:
+        pipeline_instance = source_utils.load(source)
+    except ModuleNotFoundError as e:
+        source_root = source_utils.get_source_root()
+        cli_utils.error(
+            f"Unable to import module `{e.name}`. Make sure the source path is "
+            f"relative to your source root `{source_root}`."
+        )
+    except AttributeError as e:
+        cli_utils.error("Unable to load attribute from module: " + str(e))
+
+    if not isinstance(pipeline_instance, Pipeline):
+        cli_utils.error(
+            f"The given source path `{source}` does not resolve to a pipeline "
+            "object."
+        )
+
+    with cli_utils.temporary_active_stack(stack_name_or_id=stack_name_or_id):
+        pipeline_instance = pipeline_instance.with_options(
+            config_path=config_path
+        )
+        template = pipeline_instance.create_run_template(name=name)
+
+    cli_utils.declare(f"Created run template `{template.id}`.")
+
+
 @pipeline.command("list", help="List all registered pipelines.")
 @list_options(PipelineFilter)
 def list_pipelines(**kwargs: Any) -> None:

zenml/config/compiler.py

@@ -99,7 +99,10 @@ class Compiler:
 
         self._apply_stack_default_settings(pipeline=pipeline, stack=stack)
         if run_configuration.run_name:
-            self._verify_run_name(run_configuration.run_name)
+            self._verify_run_name(
+                run_configuration.run_name,
+                pipeline.configuration.substitutions,
+            )
 
         pipeline_settings = self._filter_and_validate_settings(
             settings=pipeline.configuration.settings,
@@ -305,16 +308,22 @@ class Compiler:
         return default_settings
 
     @staticmethod
-    def _verify_run_name(run_name: str) -> None:
+    def _verify_run_name(
+        run_name: str,
+        substitutions: Dict[str, str],
+    ) -> None:
         """Verifies that the run name contains only valid placeholders.
 
         Args:
             run_name: The run name to verify.
+            substitutions: The substitutions to be used in the run name.
 
         Raises:
             ValueError: If the run name contains invalid placeholders.
         """
-        valid_placeholder_names = {"date", "time"}
+        valid_placeholder_names = {"date", "time"}.union(
+            set(substitutions.keys())
+        )
         placeholders = {
             v[1] for v in string.Formatter().parse(run_name) if v[1]
         }

zenml/config/pipeline_configurations.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from pydantic import SerializeAsAny, field_validator
@@ -46,6 +47,25 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
     retry: Optional[StepRetryConfig] = None
+    substitutions: Dict[str, str] = {}
+
+    def _get_full_substitutions(
+        self, start_time: Optional[datetime]
+    ) -> Dict[str, str]:
+        """Returns the full substitutions dict.
+
+        Args:
+            start_time: Start time of the pipeline run.
+
+        Returns:
+            The full substitutions dict including date and time.
+        """
+        if start_time is None:
+            start_time = datetime.utcnow()
+        ret = self.substitutions.copy()
+        ret.setdefault("date", start_time.strftime("%Y_%m_%d"))
+        ret.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
+        return ret
 
 
 class PipelineConfiguration(PipelineConfigurationUpdate):

zenml/config/pipeline_run_configuration.py

@@ -52,3 +52,4 @@ class PipelineRunConfiguration(
     retry: Optional[StepRetryConfig] = None
     failure_hook_source: Optional[SourceWithValidator] = None
     success_hook_source: Optional[SourceWithValidator] = None
+    substitutions: Dict[str, str] = {}

zenml/config/step_configurations.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
+from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -49,6 +50,7 @@ from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings, ResourceSettings
+    from zenml.config.pipeline_configurations import PipelineConfiguration
 
 logger = get_logger(__name__)
 
@@ -152,6 +154,7 @@ class StepConfigurationUpdate(StrictBaseModel):
     success_hook_source: Optional[SourceWithValidator] = None
     model: Optional[Model] = None
     retry: Optional[StepRetryConfig] = None
+    substitutions: Dict[str, str] = {}
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
@@ -237,6 +240,24 @@ class StepConfiguration(PartialStepConfiguration):
             model_or_dict = model_or_dict.model_dump()
         return DockerSettings.model_validate(model_or_dict)
 
+    def _get_full_substitutions(
+        self,
+        pipeline_config: "PipelineConfiguration",
+        start_time: Optional[datetime],
+    ) -> Dict[str, str]:
+        """Get the full set of substitutions for this step configuration.
+
+        Args:
+            pipeline_config: The pipeline configuration.
+            start_time: The start time of the pipeline run.
+
+        Returns:
+            The full set of substitutions for this step configuration.
+        """
+        ret = pipeline_config._get_full_substitutions(start_time)
+        ret.update(self.substitutions)
+        return ret
+
 
 class InputSpec(StrictBaseModel):
     """Step input specification."""

zenml/enums.py

@@ -60,6 +60,7 @@ class VisualizationType(StrEnum):
     HTML = "html"
     IMAGE = "image"
     MARKDOWN = "markdown"
+    JSON = "json"
 
 
 class ZenMLServiceType(StrEnum):

zenml/integrations/feast/__init__.py

@@ -31,7 +31,7 @@ class FeastIntegration(Integration):
 
     NAME = FEAST
     # click is added to keep the feast click version in sync with ZenML's click
-    REQUIREMENTS = ["feast", "click>=8.0.1,<8.1.4"]
+    REQUIREMENTS = ["feast>=0.12.0", "click>=8.0.1,<8.1.4"]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["click", "pandas"]
 
     @classmethod

zenml/integrations/feast/feature_stores/feast_feature_store.py

@@ -16,7 +16,7 @@
 from typing import Any, Dict, List, Union, cast
 
 import pandas as pd
-from feast import FeatureStore  # type: ignore
+from feast import FeatureService, FeatureStore  # type: ignore
 from feast.infra.registry.base_registry import BaseRegistry  # type: ignore
 
 from zenml.feature_stores.base_feature_store import BaseFeatureStore
@@ -43,14 +43,14 @@ class FeastFeatureStore(BaseFeatureStore):
     def get_historical_features(
         self,
         entity_df: Union[pd.DataFrame, str],
-        features: List[str],
+        features: Union[List[str], FeatureService],
         full_feature_names: bool = False,
     ) -> pd.DataFrame:
         """Returns the historical features for training or batch scoring.
 
         Args:
             entity_df: The entity DataFrame or entity name.
-            features: The features to retrieve.
+            features: The features to retrieve or a FeatureService.
             full_feature_names: Whether to return the full feature names.
 
         Raise:
@@ -70,14 +70,14 @@ class FeastFeatureStore(BaseFeatureStore):
     def get_online_features(
         self,
         entity_rows: List[Dict[str, Any]],
-        features: List[str],
+        features: Union[List[str], FeatureService],
         full_feature_names: bool = False,
     ) -> Dict[str, Any]:
         """Returns the latest online feature data.
 
         Args:
             entity_rows: The entity rows to retrieve.
-            features: The features to retrieve.
+            features: The features to retrieve or a FeatureService.
             full_feature_names: Whether to return the full feature names.
 
         Raise:
@@ -118,17 +118,21 @@ class FeastFeatureStore(BaseFeatureStore):
         fs = FeatureStore(repo_path=self.config.feast_repo)
         return [ds.name for ds in fs.list_entities()]
 
-    def get_feature_services(self) -> List[str]:
-        """Returns the feature service names.
+    def get_feature_services(self) -> List[FeatureService]:
+        """Returns the feature services.
 
         Raise:
             ConnectionError: If the online component (Redis) is not available.
 
         Returns:
-            The feature service names.
+            The feature services.
         """
         fs = FeatureStore(repo_path=self.config.feast_repo)
-        return [ds.name for ds in fs.list_feature_services()]
+        feature_services: List[FeatureService] = list(
+            fs.list_feature_services()
+        )
+
+        return feature_services
 
     def get_feature_views(self) -> List[str]:
         """Returns the feature view names.

zenml/materializers/built_in_materializer.py

@@ -28,7 +28,7 @@ from typing import (
 )
 
 from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-from zenml.enums import ArtifactType
+from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
@@ -415,6 +415,23 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 self.artifact_store.rmtree(entry["path"])
             raise e
 
+    # save dict type objects to JSON file with JSON visualization type
+    def save_visualizations(self, data: Any) -> Dict[str, "VisualizationType"]:
+        """Save visualizations for the given data.
+
+        Args:
+            data: The data to save visualizations for.
+
+        Returns:
+            A dictionary of visualization URIs and their types.
+        """
+        # dict/list type objects are always saved as JSON files
+        # doesn't work for non-serializable types as they
+        # are saved as list of lists in different files
+        if _is_serializable(data):
+            return {self.data_path: VisualizationType.JSON}
+        return {}
+
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given built-in container object.
 

zenml/materializers/structured_string_materializer.py

@@ -19,22 +19,23 @@ from typing import Dict, Type, Union
 from zenml.enums import ArtifactType, VisualizationType
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
-from zenml.types import CSVString, HTMLString, MarkdownString
+from zenml.types import CSVString, HTMLString, JSONString, MarkdownString
 
 logger = get_logger(__name__)
 
 
-STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString]
+STRUCTURED_STRINGS = Union[CSVString, HTMLString, MarkdownString, JSONString]
 
 HTML_FILENAME = "output.html"
 MARKDOWN_FILENAME = "output.md"
 CSV_FILENAME = "output.csv"
+JSON_FILENAME = "output.json"
 
 
 class StructuredStringMaterializer(BaseMaterializer):
     """Materializer for HTML or Markdown strings."""
 
-    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString)
+    ASSOCIATED_TYPES = (CSVString, HTMLString, MarkdownString, JSONString)
     ASSOCIATED_ARTIFACT_TYPE = ArtifactType.DATA_ANALYSIS
 
     def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
@@ -94,6 +95,8 @@ class StructuredStringMaterializer(BaseMaterializer):
             filename = HTML_FILENAME
         elif issubclass(data_type, MarkdownString):
             filename = MARKDOWN_FILENAME
+        elif issubclass(data_type, JSONString):
+            filename = JSON_FILENAME
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."
@@ -120,6 +123,8 @@ class StructuredStringMaterializer(BaseMaterializer):
             return VisualizationType.HTML
         elif issubclass(data_type, MarkdownString):
             return VisualizationType.MARKDOWN
+        elif issubclass(data_type, JSONString):
+            return VisualizationType.JSON
         else:
             raise ValueError(
                 f"Data type {data_type} is not supported by this materializer."

zenml/model/model.py

@@ -57,7 +57,9 @@ class Model(BaseModel):
     ethics: The ethical implications of the model.
     tags: Tags associated with the model.
     version: The version name, version number or stage is optional and points model context
-        to a specific version/stage. If skipped new version will be created.
+        to a specific version/stage. If skipped new version will be created. `version`
+        also supports placeholders: standard `{date}` and `{time}` and any custom placeholders
+        that are passed as substitutions in the pipeline or step decorators.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
     """
@@ -534,6 +536,8 @@ class Model(BaseModel):
         from zenml.models import ModelRequest
 
         zenml_client = Client()
+        # backup logic, if the Model class is used directly from the code
+        self.name = format_name_template(self.name, substitutions={})
         if self.model_version_id:
             mv = zenml_client.get_model_version(
                 model_version_name_or_number_or_id=self.model_version_id,
@@ -663,7 +667,7 @@ class Model(BaseModel):
 
         # backup logic, if the Model class is used directly from the code
         if isinstance(self.version, str):
-            self.version = format_name_template(self.version)
+            self.version = format_name_template(self.version, substitutions={})
 
         try:
             if self.version or self.model_version_id:

zenml/models/v2/core/pipeline_run.py

@@ -237,6 +237,10 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         default=False,
         description="Whether a template can be created from this run.",
     )
+    steps_substitutions: Dict[str, Dict[str, str]] = Field(
+        title="Substitutions used in the step runs of this pipeline run.",
+        default_factory=dict,
+    )
 
 
 class PipelineRunResponseResources(WorkspaceScopedResponseResources):

zenml/models/v2/core/step_run.py

@@ -142,7 +142,7 @@ class StepRunRequest(WorkspaceScopedRequest):
 class StepRunUpdate(BaseModel):
     """Update model for step runs."""
 
-    outputs: Dict[str, UUID] = Field(
+    outputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the output artifact versions of the step run.",
         default={},
     )

zenml/orchestrators/publish_utils.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
 
 
 def publish_successful_step_run(
-    step_run_id: "UUID", output_artifact_ids: Dict[str, "UUID"]
+    step_run_id: "UUID", output_artifact_ids: Dict[str, List["UUID"]]
 ) -> "StepRunResponse":
     """Publishes a successful step run.
 

zenml/orchestrators/step_launcher.py

@@ -309,8 +309,12 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
+        start_time = datetime.utcnow()
         run_name = orchestrator_utils.get_run_name(
-            run_name_template=self._deployment.run_name_template
+            run_name_template=self._deployment.run_name_template,
+            substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
+                start_time
+            ),
         )
 
         logger.debug("Creating pipeline run %s", run_name)
@@ -329,7 +333,7 @@ class StepLauncher:
             ),
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
-            start_time=datetime.utcnow(),
+            start_time=start_time,
             tags=self._deployment.pipeline_configuration.tags,
         )
         return client.zen_store.get_or_create_run(pipeline_run)

zenml/orchestrators/step_run_utils.py

@@ -354,13 +354,16 @@ def create_cached_step_runs(
 
 
 def get_or_create_model_version_for_pipeline_run(
-    model: "Model", pipeline_run: PipelineRunResponse
+    model: "Model",
+    pipeline_run: PipelineRunResponse,
+    substitutions: Dict[str, str],
 ) -> Tuple[ModelVersionResponse, bool]:
     """Get or create a model version as part of a pipeline run.
 
     Args:
         model: The model to get or create.
         pipeline_run: The pipeline run for which the model should be created.
+        substitutions: Substitutions to apply to the model version name.
 
     Returns:
         The model version and a boolean indicating whether it was newly created
@@ -374,12 +377,14 @@ def get_or_create_model_version_for_pipeline_run(
         return model._get_model_version(), False
     elif model.version:
         if isinstance(model.version, str):
-            start_time = pipeline_run.start_time or datetime.utcnow()
             model.version = string_utils.format_name_template(
                 model.version,
-                date=start_time.strftime("%Y_%m_%d"),
-                time=start_time.strftime("%H_%M_%S_%f"),
+                substitutions=substitutions,
             )
+        model.name = string_utils.format_name_template(
+            model.name,
+            substitutions=substitutions,
+        )
 
         return (
             model._get_or_create_model_version(),
@@ -460,7 +465,9 @@ def prepare_pipeline_run_model_version(
         model_version = pipeline_run.model_version
     elif config_model := pipeline_run.config.model:
         model_version, _ = get_or_create_model_version_for_pipeline_run(
-            model=config_model, pipeline_run=pipeline_run
+            model=config_model,
+            pipeline_run=pipeline_run,
+            substitutions=pipeline_run.config.substitutions,
         )
         pipeline_run = Client().zen_store.update_run(
             run_id=pipeline_run.id,
@@ -492,7 +499,9 @@ def prepare_step_run_model_version(
         model_version = step_run.model_version
     elif config_model := step_run.config.model:
         model_version, created = get_or_create_model_version_for_pipeline_run(
-            model=config_model, pipeline_run=pipeline_run
+            model=config_model,
+            pipeline_run=pipeline_run,
+            substitutions=step_run.config.substitutions,
         )
         step_run = Client().zen_store.update_run_step(
             step_run_id=step_run.id,

zenml/orchestrators/step_runner.py

@@ -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,33 @@ 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()):
+            _evaluated_name = None
+            if v.artifact_config:
+                _evaluated_name = v.artifact_config._evaluated_name(
+                    step_run.config.substitutions
+                )
+            if _evaluated_name is None:
+                _evaluated_name = k
+
+            for d in collections:
+                d[_evaluated_name] = d.pop(k)
+
     def _load_step(self) -> "BaseStep":
         """Load the step instance.
 

zenml/orchestrators/utils.py

@@ -196,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.
@@ -208,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.")