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

zenml/{new/pipelines/pipeline.py → pipelines/pipeline_definition.py}

@@ -37,6 +37,7 @@ from uuid import UUID
 
 import yaml
 from pydantic import ConfigDict, ValidationError
+from typing_extensions import Self
 
 from zenml import constants
 from zenml.analytics.enums import AnalyticsEvent
@@ -65,10 +66,11 @@ from zenml.models import (
     PipelineRequest,
     PipelineResponse,
     PipelineRunResponse,
+    RunTemplateResponse,
     ScheduleRequest,
 )
-from zenml.new.pipelines import build_utils
-from zenml.new.pipelines.run_utils import (
+from zenml.pipelines import build_utils
+from zenml.pipelines.run_utils import (
     create_placeholder_run,
     deploy_pipeline,
     upload_notebook_cell_code_if_necessary,
@@ -99,6 +101,7 @@ if TYPE_CHECKING:
     from zenml.config.source import Source
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
     from zenml.model.model import Model
+    from zenml.models import ArtifactVersionResponse
     from zenml.types import HookSpecification
 
     StepConfigurationUpdateOrDict = Union[
@@ -107,7 +110,6 @@ if TYPE_CHECKING:
 
 logger = get_logger(__name__)
 
-T = TypeVar("T", bound="Pipeline")
 F = TypeVar("F", bound=Callable[..., None])
 
 
@@ -128,6 +130,7 @@ class Pipeline:
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
         settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
+        tags: Optional[List[str]] = None,
         extra: Optional[Dict[str, Any]] = None,
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
@@ -144,7 +147,8 @@ class Pipeline:
             enable_artifact_visualization: If artifact visualization should be
                 enabled for this pipeline.
             enable_step_logs: If step logs should be enabled for this pipeline.
-            settings: settings for this pipeline.
+            settings: Settings for this pipeline.
+            tags: Tags to apply to runs of this pipeline.
             extra: Extra configurations for this pipeline.
             on_failure: Callback function in event of failure of the step. Can
                 be a function with a single argument of type `BaseException`, or
@@ -168,6 +172,7 @@ class Pipeline:
                 enable_artifact_visualization=enable_artifact_visualization,
                 enable_step_logs=enable_step_logs,
                 settings=settings,
+                tags=tags,
                 extra=extra,
                 on_failure=on_failure,
                 on_success=on_success,
@@ -279,19 +284,20 @@ class Pipeline:
         self.__suppress_warnings_flag__ = False
 
     def configure(
-        self: T,
+        self,
         enable_cache: Optional[bool] = None,
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
         settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
+        tags: Optional[List[str]] = None,
         extra: Optional[Dict[str, Any]] = None,
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
         parameters: Optional[Dict[str, Any]] = None,
         merge: bool = True,
-    ) -> T:
+    ) -> Self:
         """Configures the pipeline.
 
         Configuration merging example:
@@ -312,6 +318,7 @@ class Pipeline:
                 enabled for this pipeline.
             enable_step_logs: If step logs should be enabled for this pipeline.
             settings: settings for this pipeline.
+            tags: Tags to apply to runs of this pipeline.
             extra: Extra configurations for this pipeline.
             on_failure: Callback function in event of failure of the step. Can
                 be a function with a single argument of type `BaseException`, or
@@ -340,6 +347,11 @@ class Pipeline:
             # string of on_success hook function to be used for this pipeline
             success_hook_source = resolve_and_validate_hook(on_success)
 
+        if merge and tags and self._configuration.tags:
+            # Merge tags explicitly here as the recursive update later only
+            # merges dicts
+            tags = self._configuration.tags + tags
+
         values = dict_utils.remove_none_values(
             {
                 "enable_cache": enable_cache,
@@ -347,6 +359,7 @@ class Pipeline:
                 "enable_artifact_visualization": enable_artifact_visualization,
                 "enable_step_logs": enable_step_logs,
                 "settings": settings,
+                "tags": tags,
                 "extra": extra,
                 "failure_hook_source": failure_hook_source,
                 "success_hook_source": success_hook_source,
@@ -539,7 +552,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 code_repository=code_repository,
             )
 
-    def _run(
+    def _create_deployment(
         self,
         *,
         run_name: Optional[str] = None,
@@ -557,8 +570,8 @@ 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,
-    ) -> Optional[PipelineRunResponse]:
-        """Runs the pipeline on the active stack.
+    ) -> PipelineDeploymentResponse:
+        """Create a pipeline deployment.
 
         Args:
             run_name: Name of the pipeline run.
@@ -585,175 +598,182 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 `DockerSettings.prevent_build_reuse` instead.
 
         Returns:
-            Model of the pipeline run if running without a schedule, `None` if
-            running with a schedule.
+            The pipeline deployment.
 
         Raises:
-            ValueError: if the orchestrator doesn't support scheduling, but schedule was given
+            ValueError: If the orchestrator doesn't support scheduling, but a
+                schedule was given
         """
-        if constants.SHOULD_PREVENT_PIPELINE_EXECUTION:
-            # An environment variable was set to stop the execution of
-            # pipelines. This is done to prevent execution of module-level
-            # pipeline.run() calls when importing modules needed to run a step.
-            logger.info(
-                "Preventing execution of pipeline '%s'. If this is not "
-                "intended behavior, make sure to unset the environment "
-                "variable '%s'.",
-                self.name,
-                constants.ENV_ZENML_PREVENT_PIPELINE_EXECUTION,
-            )
-            return None
+        deployment, schedule, build = self._compile(
+            config_path=config_path,
+            run_name=run_name,
+            enable_cache=enable_cache,
+            enable_artifact_metadata=enable_artifact_metadata,
+            enable_artifact_visualization=enable_artifact_visualization,
+            enable_step_logs=enable_step_logs,
+            steps=step_configurations,
+            settings=settings,
+            schedule=schedule,
+            build=build,
+            extra=extra,
+        )
 
-        logger.info(f"Initiating a new run for the pipeline: `{self.name}`.")
+        skip_pipeline_registration = constants.handle_bool_env_var(
+            constants.ENV_ZENML_SKIP_PIPELINE_REGISTRATION,
+            default=False,
+        )
 
-        with track_handler(AnalyticsEvent.RUN_PIPELINE) as analytics_handler:
-            deployment, schedule, build = self._compile(
-                config_path=config_path,
-                run_name=run_name,
-                enable_cache=enable_cache,
-                enable_artifact_metadata=enable_artifact_metadata,
-                enable_artifact_visualization=enable_artifact_visualization,
-                enable_step_logs=enable_step_logs,
-                steps=step_configurations,
-                settings=settings,
-                schedule=schedule,
-                build=build,
-                extra=extra,
-            )
+        register_pipeline = not (skip_pipeline_registration or unlisted)
 
-            skip_pipeline_registration = constants.handle_bool_env_var(
-                constants.ENV_ZENML_SKIP_PIPELINE_REGISTRATION,
-                default=False,
-            )
+        pipeline_id = None
+        if register_pipeline:
+            pipeline_id = self._register().id
 
-            register_pipeline = not (skip_pipeline_registration or unlisted)
+        else:
+            logger.debug(f"Pipeline {self.name} is unlisted.")
+
+        stack = Client().active_stack
+        stack.validate()
+
+        schedule_id = None
+        if schedule:
+            if not stack.orchestrator.config.is_schedulable:
+                raise ValueError(
+                    f"Stack {stack.name} does not support scheduling. "
+                    "Not all orchestrator types support scheduling, "
+                    "kindly consult with "
+                    "https://docs.zenml.io/how-to/build-pipelines/schedule-a-pipeline "
+                    "for details."
+                )
+            if schedule.name:
+                schedule_name = schedule.name
+            else:
+                schedule_name = format_name_template(
+                    deployment.run_name_template
+                )
+            components = Client().active_stack_model.components
+            orchestrator = components[StackComponentType.ORCHESTRATOR][0]
+            schedule_model = ScheduleRequest(
+                workspace=Client().active_workspace.id,
+                user=Client().active_user.id,
+                pipeline_id=pipeline_id,
+                orchestrator_id=orchestrator.id,
+                name=schedule_name,
+                active=True,
+                cron_expression=schedule.cron_expression,
+                start_time=schedule.start_time,
+                end_time=schedule.end_time,
+                interval_second=schedule.interval_second,
+                catchup=schedule.catchup,
+                run_once_start_time=schedule.run_once_start_time,
+            )
+            schedule_id = Client().zen_store.create_schedule(schedule_model).id
+            logger.info(
+                f"Created schedule `{schedule_name}` for pipeline "
+                f"`{deployment.pipeline_configuration.name}`."
+            )
 
-            pipeline_id = None
-            if register_pipeline:
-                pipeline_id = self._register().id
+        stack = Client().active_stack
+        stack.validate()
+        upload_notebook_cell_code_if_necessary(
+            deployment=deployment, stack=stack
+        )
 
-            else:
-                logger.debug(f"Pipeline {self.name} is unlisted.")
+        local_repo_context = (
+            code_repository_utils.find_active_code_repository()
+        )
+        code_repository = build_utils.verify_local_repository_context(
+            deployment=deployment, local_repo_context=local_repo_context
+        )
 
-            stack = Client().active_stack
-            stack.validate()
-
-            schedule_id = None
-            if schedule:
-                if not stack.orchestrator.config.is_schedulable:
-                    raise ValueError(
-                        f"Stack {stack.name} does not support scheduling. "
-                        "Not all orchestrator types support scheduling, "
-                        "kindly consult with "
-                        "https://docs.zenml.io/how-to/build-pipelines/schedule-a-pipeline "
-                        "for details."
-                    )
-                if schedule.name:
-                    schedule_name = schedule.name
-                else:
-                    schedule_name = format_name_template(
-                        deployment.run_name_template
-                    )
-                components = Client().active_stack_model.components
-                orchestrator = components[StackComponentType.ORCHESTRATOR][0]
-                schedule_model = ScheduleRequest(
-                    workspace=Client().active_workspace.id,
-                    user=Client().active_user.id,
-                    pipeline_id=pipeline_id,
-                    orchestrator_id=orchestrator.id,
-                    name=schedule_name,
-                    active=True,
-                    cron_expression=schedule.cron_expression,
-                    start_time=schedule.start_time,
-                    end_time=schedule.end_time,
-                    interval_second=schedule.interval_second,
-                    catchup=schedule.catchup,
-                    run_once_start_time=schedule.run_once_start_time,
-                )
-                schedule_id = (
-                    Client().zen_store.create_schedule(schedule_model).id
-                )
-                logger.info(
-                    f"Created schedule `{schedule_name}` for pipeline "
-                    f"`{deployment.pipeline_configuration.name}`."
-                )
+        if prevent_build_reuse:
+            logger.warning(
+                "Passing `prevent_build_reuse=True` to "
+                "`pipeline.with_opitions(...)` is deprecated. Use "
+                "`DockerSettings.prevent_build_reuse` instead."
+            )
 
-            stack = Client().active_stack
-            stack.validate()
-            upload_notebook_cell_code_if_necessary(
-                deployment=deployment, stack=stack
+        build_model = build_utils.reuse_or_create_pipeline_build(
+            deployment=deployment,
+            pipeline_id=pipeline_id,
+            allow_build_reuse=not prevent_build_reuse,
+            build=build,
+            code_repository=code_repository,
+        )
+        build_id = build_model.id if build_model else None
+
+        code_reference = None
+        if local_repo_context and not local_repo_context.is_dirty:
+            source_root = source_utils.get_source_root()
+            subdirectory = (
+                Path(source_root)
+                .resolve()
+                .relative_to(local_repo_context.root)
             )
 
-            local_repo_context = (
-                code_repository_utils.find_active_code_repository()
+            code_reference = CodeReferenceRequest(
+                commit=local_repo_context.current_commit,
+                subdirectory=subdirectory.as_posix(),
+                code_repository=local_repo_context.code_repository_id,
             )
-            code_repository = build_utils.verify_local_repository_context(
-                deployment=deployment, local_repo_context=local_repo_context
+
+        code_path = None
+        if build_utils.should_upload_code(
+            deployment=deployment,
+            build=build_model,
+            code_reference=code_reference,
+        ):
+            code_archive = code_utils.CodeArchive(
+                root=source_utils.get_source_root()
             )
+            logger.info("Archiving pipeline code...")
+            code_path = code_utils.upload_code_if_necessary(code_archive)
+
+        request = PipelineDeploymentRequest(
+            user=Client().active_user.id,
+            workspace=Client().active_workspace.id,
+            stack=stack.id,
+            pipeline=pipeline_id,
+            build=build_id,
+            schedule=schedule_id,
+            code_reference=code_reference,
+            code_path=code_path,
+            **deployment.model_dump(),
+        )
+        return Client().zen_store.create_deployment(deployment=request)
 
-            if prevent_build_reuse:
-                logger.warning(
-                    "Passing `prevent_build_reuse=True` to "
-                    "`pipeline.with_opitions(...)` is deprecated. Use "
-                    "`DockerSettings.prevent_build_reuse` instead."
-                )
+    def _run(
+        self,
+    ) -> Optional[PipelineRunResponse]:
+        """Runs the pipeline on the active stack.
 
-            build_model = build_utils.reuse_or_create_pipeline_build(
-                deployment=deployment,
-                pipeline_id=pipeline_id,
-                allow_build_reuse=not prevent_build_reuse,
-                build=build,
-                code_repository=code_repository,
+        Returns:
+            The pipeline run or `None` if running with a schedule.
+        """
+        if constants.SHOULD_PREVENT_PIPELINE_EXECUTION:
+            # An environment variable was set to stop the execution of
+            # pipelines. This is done to prevent execution of module-level
+            # pipeline.run() calls when importing modules needed to run a step.
+            logger.info(
+                "Preventing execution of pipeline '%s'. If this is not "
+                "intended behavior, make sure to unset the environment "
+                "variable '%s'.",
+                self.name,
+                constants.ENV_ZENML_PREVENT_PIPELINE_EXECUTION,
             )
-            build_id = build_model.id if build_model else None
-
-            code_reference = None
-            if local_repo_context and not local_repo_context.is_dirty:
-                source_root = source_utils.get_source_root()
-                subdirectory = (
-                    Path(source_root)
-                    .resolve()
-                    .relative_to(local_repo_context.root)
-                )
-
-                code_reference = CodeReferenceRequest(
-                    commit=local_repo_context.current_commit,
-                    subdirectory=subdirectory.as_posix(),
-                    code_repository=local_repo_context.code_repository_id,
-                )
+            return None
 
-            code_path = None
-            if build_utils.should_upload_code(
-                deployment=deployment,
-                build=build_model,
-                code_reference=code_reference,
-            ):
-                code_archive = code_utils.CodeArchive(
-                    root=source_utils.get_source_root()
-                )
-                logger.info("Archiving pipeline code...")
-                code_path = code_utils.upload_code_if_necessary(code_archive)
+        logger.info(f"Initiating a new run for the pipeline: `{self.name}`.")
 
-            deployment_request = PipelineDeploymentRequest(
-                user=Client().active_user.id,
-                workspace=Client().active_workspace.id,
-                stack=stack.id,
-                pipeline=pipeline_id,
-                build=build_id,
-                schedule=schedule_id,
-                code_reference=code_reference,
-                code_path=code_path,
-                **deployment.model_dump(),
-            )
-            deployment_model = Client().zen_store.create_deployment(
-                deployment=deployment_request
-            )
+        with track_handler(AnalyticsEvent.RUN_PIPELINE) as analytics_handler:
+            stack = Client().active_stack
+            deployment = self._create_deployment(**self._run_args)
 
-            self.log_pipeline_deployment_metadata(deployment_model)
-            run = create_placeholder_run(deployment=deployment_model)
+            self.log_pipeline_deployment_metadata(deployment)
+            run = create_placeholder_run(deployment=deployment)
 
             analytics_handler.metadata = self._get_pipeline_analytics_metadata(
-                deployment=deployment_model,
+                deployment=deployment,
                 stack=stack,
                 run_id=run.id if run else None,
             )
@@ -766,11 +786,11 @@ To avoid this consider setting pipeline parameters only in one place (config or
                     logger.info(
                         "You can visualize your pipeline runs in the `ZenML "
                         "Dashboard`. In order to try it locally, please run "
-                        "`zenml up`."
+                        "`zenml login --local`."
                     )
 
             deploy_pipeline(
-                deployment=deployment_model, stack=stack, placeholder_run=run
+                deployment=deployment, stack=stack, placeholder_run=run
             )
             if run:
                 return Client().get_pipeline_run(run.id)
@@ -786,15 +806,6 @@ To avoid this consider setting pipeline parameters only in one place (config or
             deployment_model: The model for the pipeline deployment
         """
         try:
-            # Log about the schedule/run
-            if deployment_model.schedule:
-                logger.info(
-                    "Scheduling a run with the schedule: "
-                    f"`{deployment_model.schedule.name}`"
-                )
-            else:
-                logger.info("Executing a new run.")
-
             # Log about the caching status
             if deployment_model.pipeline_configuration.enable_cache is False:
                 logger.info(
@@ -847,22 +858,6 @@ To avoid this consider setting pipeline parameters only in one place (config or
         except Exception as e:
             logger.debug(f"Logging pipeline deployment metadata failed: {e}")
 
-    def get_runs(self, **kwargs: Any) -> List["PipelineRunResponse"]:
-        """(Deprecated) Get runs of this pipeline.
-
-        Args:
-            **kwargs: Further arguments for filtering or pagination that are
-                passed to `client.list_pipeline_runs()`.
-
-        Returns:
-            List of runs of this pipeline.
-        """
-        logger.warning(
-            "`Pipeline.get_runs()` is deprecated and will be removed in a "
-            "future version. Please use `Pipeline.model.get_runs()` instead."
-        )
-        return self.model.get_runs(**kwargs)
-
     def write_run_configuration_template(
         self, path: str, stack: Optional["Stack"] = None
     ) -> None:
@@ -897,19 +892,12 @@ To avoid this consider setting pipeline parameters only in one place (config or
         steps = {}
         for step_name, invocation in self.invocations.items():
             step = invocation.step
-            parameters = (
-                pydantic_utils.TemplateGenerator(
-                    step.entrypoint_definition.legacy_params.annotation
-                ).run()
-                if step.entrypoint_definition.legacy_params
-                else {}
-            )
             outputs = {
                 name: PartialArtifactConfiguration()
                 for name in step.entrypoint_definition.outputs
             }
             step_template = StepConfigurationUpdate(
-                parameters=parameters,
+                parameters={},
                 settings=step_settings,
                 outputs=outputs,
             )
@@ -1112,7 +1100,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
         self,
         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],
@@ -1213,7 +1203,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
         raise RuntimeError("Unable to find step ID")
 
-    def __enter__(self: T) -> T:
+    def __enter__(self) -> Self:
         """Activate the pipeline context.
 
         Raises:
@@ -1260,14 +1250,6 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 {k: v for k, v in _from_config_file.items() if k in matcher}
             )
 
-            # TODO: deprecate me
-            if "model_version" in _from_config_file:
-                logger.warning(
-                    "YAML config option `model_version` is deprecated. Please use `model`."
-                )
-                _from_config_file["model"] = _from_config_file["model_version"]
-                del _from_config_file["model_version"]
-
             if "model" in _from_config_file:
                 if "model" in self._from_config_file:
                     _from_config_file["model"] = self._from_config_file[
@@ -1356,7 +1338,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
         """
         return copy.deepcopy(self)
 
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+    def __call__(
+        self, *args: Any, **kwargs: Any
+    ) -> Optional[PipelineRunResponse]:
         """Handle a call of the pipeline.
 
         This method does one of two things:
@@ -1371,7 +1355,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
             **kwargs: Entrypoint function keyword arguments.
 
         Returns:
-            The outputs of the entrypoint function call.
+            If called within another pipeline, returns the outputs of the
+            `entrypoint` method. Otherwise, returns the pipeline run or `None`
+            if running with a schedule.
         """
         if Pipeline.ACTIVE_PIPELINE:
             # Calling a pipeline inside a pipeline, we return the potential
@@ -1383,7 +1369,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
             return self.entrypoint(*args, **kwargs)
 
         self.prepare(*args, **kwargs)
-        return self._run(**self._run_args)
+        return self._run()
 
     def _call_entrypoint(self, *args: Any, **kwargs: Any) -> None:
         """Calls the pipeline entrypoint function with the given arguments.
@@ -1431,6 +1417,26 @@ To avoid this consider setting pipeline parameters only in one place (config or
             else:
                 self.prepare()
 
+    def create_run_template(
+        self, name: str, **kwargs: Any
+    ) -> RunTemplateResponse:
+        """Create a run template for the pipeline.
+
+        Args:
+            name: The name of the run template.
+            **kwargs: Keyword arguments for the client method to create a run
+                template.
+
+        Returns:
+            The created run template.
+        """
+        self._prepare_if_possible()
+        deployment = self._create_deployment(**self._run_args)
+
+        return Client().create_run_template(
+            name=name, deployment_id=deployment.id, **kwargs
+        )
+
     def _reconfigure_from_file_with_overrides(
         self,
         config_path: Optional[str] = None,
@@ -1454,10 +1460,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         if config_path:
             self._from_config_file = self._parse_config_file(
                 config_path=config_path,
-                matcher=inspect.getfullargspec(self.configure)[0]
-                + [
-                    "model_version"
-                ],  # TODO: deprecate `model_version` later on
+                matcher=inspect.getfullargspec(self.configure)[0],
             )
 
         _from_config_file = dict_utils.recursive_update(