zenml-nightly 0.83.0.dev20250623__py3-none-any.whl → 0.83.1.dev20250625__py3-none-any.whl

zenml/integrations/databricks/orchestrators/databricks_orchestrator.py

@@ -15,7 +15,7 @@
 
 import itertools
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, cast
 from uuid import UUID
 
 from databricks.sdk import WorkspaceClient as DatabricksClient
@@ -48,10 +48,12 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType, Uri
 from zenml.models.v2.core.schedule import ScheduleResponse
+from zenml.orchestrators import (
+    SubmissionResult,
+    WheeledOrchestrator,
+)
 from zenml.orchestrators.utils import get_orchestrator_run_name
-from zenml.orchestrators.wheeled_orchestrator import WheeledOrchestrator
 from zenml.stack import StackValidator
-from zenml.utils import io_utils
 from zenml.utils.package_utils import clean_requirements
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
@@ -67,20 +69,13 @@ logger = get_logger(__name__)
 ZENML_STEP_DEFAULT_ENTRYPOINT_COMMAND = "entrypoint.main"
 DATABRICKS_WHEELS_DIRECTORY_PREFIX = "dbfs:/FileStore/zenml"
 DATABRICKS_LOCAL_FILESYSTEM_PREFIX = "file:/"
-DATABRICKS_CLUSTER_DEFAULT_NAME = "zenml-databricks-cluster"
 DATABRICKS_SPARK_DEFAULT_VERSION = "15.3.x-scala2.12"
 DATABRICKS_JOB_ID_PARAMETER_REFERENCE = "{{job.id}}"
 DATABRICKS_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH = "."
 
 
 class DatabricksOrchestrator(WheeledOrchestrator):
-    """Base class for Orchestrator responsible for running pipelines remotely in a VM.
-
-    This orchestrator does not support running on a schedule.
-    """
-
-    # The default instance type to use if none is specified in settings
-    DEFAULT_INSTANCE_TYPE: Optional[str] = None
+    """Databricks orchestrator."""
 
     @property
     def validator(self) -> Optional[StackValidator]:
@@ -168,69 +163,39 @@ class DatabricksOrchestrator(WheeledOrchestrator):
                 f"{ENV_ZENML_DATABRICKS_ORCHESTRATOR_RUN_ID}."
             )
 
-    @property
-    def root_directory(self) -> str:
-        """Path to the root directory for all files concerning this orchestrator.
-
-        Returns:
-            Path to the root directory.
-        """
-        return os.path.join(
-            io_utils.get_global_config_directory(),
-            "databricks",
-            str(self.id),
-        )
-
-    @property
-    def pipeline_directory(self) -> str:
-        """Returns path to a directory in which the kubeflow pipeline files are stored.
-
-        Returns:
-            Path to the pipeline directory.
-        """
-        return os.path.join(self.root_directory, "pipelines")
-
     def setup_credentials(self) -> None:
         """Set up credentials for the orchestrator."""
         connector = self.get_connector()
         assert connector is not None
         connector.configure_local_client()
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Creates a wheel and uploads the pipeline to Databricks.
-
-        This functions as an intermediary representation of the pipeline which
-        is then deployed to the kubeflow pipelines instance.
-
-        How it works:
-        -------------
-        Before this method is called the `prepare_pipeline_deployment()`
-        method builds a docker image that contains the code for the
-        pipeline, all steps the context around these files.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
 
-        Based on this docker image a callable is created which builds
-        task for each step (`_construct_databricks_pipeline`).
-        To do this the entrypoint of the docker image is configured to
-        run the correct step within the docker image. The dependencies
-        between these task are then also configured onto each
-        task by pointing at the downstream steps.
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             ValueError: If the schedule is not set or if the cron expression
                 is not set.
+
+        Returns:
+            Optional submission result.
         """
         settings = cast(
             DatabricksOrchestratorSettings, self.get_settings(deployment)
@@ -339,11 +304,6 @@ class DatabricksOrchestrator(WheeledOrchestrator):
         orchestrator_run_name = get_orchestrator_run_name(
             pipeline_name=deployment.pipeline_configuration.name
         )
-        # Get a filepath to use to save the finished yaml to
-        fileio.makedirs(self.pipeline_directory)
-        pipeline_file_path = os.path.join(
-            self.pipeline_directory, f"{orchestrator_run_name}.yaml"
-        )
 
         # Copy the repository to a temporary directory and add a setup.py file
         repository_temp_dir = (
@@ -382,11 +342,6 @@ class DatabricksOrchestrator(WheeledOrchestrator):
 
         fileio.rmtree(repository_temp_dir)
 
-        logger.info(
-            "Writing Databricks workflow definition to `%s`.",
-            pipeline_file_path,
-        )
-
         # using the databricks client uploads the pipeline to databricks
         job_cluster_key = self.sanitize_name(f"{deployment_id}")
         self._upload_and_run_pipeline(
@@ -399,6 +354,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
             job_cluster_key=job_cluster_key,
             schedule=deployment.schedule,
         )
+        return None
 
     def _upload_and_run_pipeline(
         self,

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -37,7 +37,6 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
-    Iterator,
     List,
     Optional,
     Tuple,
@@ -89,7 +88,7 @@ from zenml.integrations.gcp.vertex_custom_job_parameters import (
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType, Uri
-from zenml.orchestrators import ContainerizedOrchestrator
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack.stack_validator import StackValidator
 from zenml.utils.io_utils import get_global_config_directory
@@ -402,57 +401,29 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
 
         return custom_job_component
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Iterator[Dict[str, MetadataType]]:
-        """Creates a KFP JSON pipeline.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
 
-        # noqa: DAR402
-
-        This is an intermediary representation of the pipeline which is then
-        deployed to Vertex AI Pipelines service.
-
-        How it works:
-        -------------
-        Before this method is called the `prepare_pipeline_deployment()` method
-        builds a Docker image that contains the code for the pipeline, all steps
-        the context around these files.
-
-        Based on this Docker image a callable is created which builds
-        container_ops for each step (`_construct_kfp_pipeline`). The function
-        `kfp.components.load_component_from_text` is used to create the
-        `ContainerOp`, because using the `dsl.ContainerOp` class directly is
-        deprecated when using the Kubeflow SDK v2. The step entrypoint command
-        with the entrypoint arguments is the command that will be executed by
-        the container created using the previously created Docker image.
-
-        This callable is then compiled into a JSON file that is used as the
-        intermediary representation of the Kubeflow pipeline.
-
-        This file then is submitted to the Vertex AI Pipelines service for
-        execution.
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
-        Raises:
-            ValueError: If the attribute `pipeline_root` is not set, and it
-                can be not generated using the path of the artifact store in the
-                stack because it is not a
-                `zenml.integrations.gcp.artifact_store.GCPArtifactStore`. Also gets
-                raised if attempting to schedule pipeline run without using the
-                `zenml.integrations.gcp.artifact_store.GCPArtifactStore`.
-
-        Yields:
-            A dictionary of metadata related to the pipeline run.
+        Returns:
+            Optional submission result.
         """
         orchestrator_run_name = get_orchestrator_run_name(
             pipeline_name=deployment.pipeline_configuration.name
@@ -651,16 +622,13 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             VertexOrchestratorSettings, self.get_settings(deployment)
         )
 
-        # Using the Google Cloud AIPlatform client, upload and execute the
-        # pipeline on the Vertex AI Pipelines service.
-        if metadata := self._upload_and_run_pipeline(
+        return self._upload_and_run_pipeline(
             pipeline_name=deployment.pipeline_configuration.name,
             pipeline_file_path=pipeline_file_path,
             run_name=orchestrator_run_name,
             settings=settings,
             schedule=deployment.schedule,
-        ):
-            yield from metadata
+        )
 
     def _upload_and_run_pipeline(
         self,
@@ -669,7 +637,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         run_name: str,
         settings: VertexOrchestratorSettings,
         schedule: Optional["ScheduleResponse"] = None,
-    ) -> Iterator[Dict[str, MetadataType]]:
+    ) -> Optional[SubmissionResult]:
         """Uploads and run the pipeline on the Vertex AI Pipelines service.
 
         Args:
@@ -684,8 +652,8 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             RuntimeError: If the Vertex Orchestrator fails to provision or any
                 other Runtime errors.
 
-        Yields:
-            A dictionary of metadata related to the pipeline run.
+        Returns:
+            Optional submission result.
         """
         # We have to replace the hyphens in the run name with underscores
         # and lower case the string, because the Vertex AI Pipelines service
@@ -772,6 +740,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     service_account=self.config.workload_service_account,
                     network=self.config.network,
                 )
+                return None
 
             else:
                 logger.info(
@@ -793,17 +762,23 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     run._dashboard_uri(),
                 )
 
-                # Yield metadata based on the generated job object
-                yield from self.compute_metadata(run)
+                _wait_for_completion = None
 
                 if settings.synchronous:
-                    logger.info(
-                        "Waiting for the Vertex AI Pipelines job to finish..."
-                    )
-                    run.wait()
-                    logger.info(
-                        "Vertex AI Pipelines job completed successfully."
-                    )
+
+                    def _wait_for_completion() -> None:
+                        logger.info(
+                            "Waiting for the Vertex AI Pipelines job to finish..."
+                        )
+                        run.wait()
+                        logger.info(
+                            "Vertex AI Pipelines job completed successfully."
+                        )
+
+                return SubmissionResult(
+                    metadata=self.compute_metadata(run),
+                    wait_for_completion=_wait_for_completion,
+                )
 
         except google_exceptions.ClientError as e:
             logger.error("Failed to create the Vertex AI Pipelines job: %s", e)
@@ -993,13 +968,13 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
 
     def compute_metadata(
         self, job: aiplatform.PipelineJob
-    ) -> Iterator[Dict[str, MetadataType]]:
+    ) -> Dict[str, MetadataType]:
         """Generate run metadata based on the corresponding Vertex PipelineJob.
 
         Args:
             job: The corresponding PipelineJob object.
 
-        Yields:
+        Returns:
             A dictionary of metadata related to the pipeline run.
         """
         metadata: Dict[str, MetadataType] = {}
@@ -1016,7 +991,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         if logs_url := self._compute_orchestrator_logs_url(job):
             metadata[METADATA_ORCHESTRATOR_LOGS_URL] = Uri(logs_url)
 
-        yield metadata
+        return metadata
 
     @staticmethod
     def _compute_orchestrator_url(

zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py

@@ -30,9 +30,7 @@ from zenml.integrations.hyperai.flavors.hyperai_orchestrator_flavor import (
     HyperAIOrchestratorSettings,
 )
 from zenml.logger import get_logger
-from zenml.orchestrators import (
-    ContainerizedOrchestrator,
-)
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.stack import Stack, StackValidator
 
 if TYPE_CHECKING:
@@ -159,14 +157,19 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
                 f"Failed to write {description} to HyperAI instance. Does the user have permissions to write?"
             )
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Sequentially runs all pipeline steps in Docker containers.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Assumes that:
         - A HyperAI (hyperai.ai) instance is running on the configured IP address.
@@ -179,26 +182,25 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
             orchestrator.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
-            RuntimeError: If a step fails.
+            RuntimeError: If running the pipeline fails.
+
+        Returns:
+            Optional submission result.
         """
         from zenml.integrations.hyperai.service_connectors.hyperai_service_connector import (
             HyperAIServiceConnector,
         )
 
-        # Basic Docker Compose definition
         compose_definition: Dict[str, Any] = {"version": "3", "services": {}}
-
-        # Get deployment id
         deployment_id = deployment.id
 
-        # Set environment
         os.environ[ENV_ZENML_HYPERAI_RUN_ID] = str(deployment_id)
         environment[ENV_ZENML_HYPERAI_RUN_ID] = str(deployment_id)
 
@@ -208,12 +210,9 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
             # Get image
             image = self.get_image(deployment=deployment, step_name=step_name)
 
-            # Get settings
             step_settings = cast(
                 HyperAIOrchestratorSettings, self.get_settings(step)
             )
-
-            # Define container name as combination between deployment id and step name
             container_name = f"{deployment_id}-{step_name}"
 
             # Make Compose service definition for step
@@ -246,10 +245,9 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
                     }
                 }
 
-            # Depending on whether it is a scheduled or a realtime pipeline, add
-            # potential .env file to service definition for deployment ID override.
             if deployment.schedule:
-                # drop ZENML_HYPERAI_ORCHESTRATOR_RUN_ID from environment but only if it is set
+                # If running on a schedule, the run ID is set dynamically via
+                # the .env file.
                 if ENV_ZENML_HYPERAI_RUN_ID in environment:
                     del environment[ENV_ZENML_HYPERAI_RUN_ID]
                 compose_definition["services"][container_name]["env_file"] = [
@@ -282,15 +280,12 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
                         }
                     )
 
-        # Convert into yaml
-        logger.info("Finalizing Docker Compose definition.")
         compose_definition_yaml: str = yaml.dump(compose_definition)
 
         # Connect to configured HyperAI instance
         logger.info(
             "Connecting to HyperAI instance and placing Docker Compose file."
         )
-        paramiko_client: paramiko.SSHClient
         if connector := self.get_connector():
             paramiko_client = connector.connect()
             if paramiko_client is None:
@@ -510,3 +505,5 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
             raise RuntimeError(
                 "A cron expression or start time is required for scheduled pipelines."
             )
+
+        return None

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -69,7 +69,7 @@ from zenml.integrations.kubeflow.flavors.kubeflow_orchestrator_flavor import (
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType, Uri
-from zenml.orchestrators import ContainerizedOrchestrator
+from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils import io_utils, settings_utils, yaml_utils
@@ -466,47 +466,33 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
 
         return pipeline_task
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Creates a kfp yaml file.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
 
-        This functions as an intermediary representation of the pipeline which
-        is then deployed to the kubeflow pipelines instance.
-
-        How it works:
-        -------------
-        Before this method is called the `prepare_pipeline_deployment()`
-        method builds a docker image that contains the code for the
-        pipeline, all steps the context around these files.
-
-        Based on this docker image a callable is created which builds
-        container_ops for each step (`_construct_kfp_pipeline`).
-        To do this the entrypoint of the docker image is configured to
-        run the correct step within the docker image. The dependencies
-        between these container_ops are then also configured onto each
-        container_op by pointing at the downstream steps.
-
-        This callable is then compiled into a kfp yaml file that is used as
-        the intermediary representation of the kubeflow pipeline.
-
-        This file, together with some metadata, runtime configurations is
-        then uploaded into the kubeflow pipelines cluster for execution.
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             RuntimeError: If trying to run a pipeline in a notebook
                 environment.
+
+        Returns:
+            Optional submission result.
         """
         # First check whether the code running in a notebook
         if Environment.in_notebook():
@@ -672,7 +658,7 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
 
         # using the kfp client uploads the pipeline to kubeflow pipelines and
         # runs it there
-        self._upload_and_run_pipeline(
+        return self._upload_and_run_pipeline(
             deployment=deployment,
             pipeline_file_path=pipeline_file_path,
             run_name=orchestrator_run_name,
@@ -683,7 +669,7 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         pipeline_file_path: str,
         run_name: str,
-    ) -> None:
+    ) -> Optional[SubmissionResult]:
         """Tries to upload and run a KFP pipeline.
 
         Args:
@@ -693,6 +679,9 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
 
         Raises:
             RuntimeError: If Kubeflow API returns an error.
+
+        Returns:
+            Optional submission result.
         """
         pipeline_name = deployment.pipeline_configuration.name
         settings = cast(
@@ -788,8 +777,14 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
                 )
 
                 if settings.synchronous:
-                    client.wait_for_run_completion(
-                        run_id=result.run_id, timeout=settings.timeout
+
+                    def _wait_for_completion() -> None:
+                        client.wait_for_run_completion(
+                            run_id=result.run_id, timeout=settings.timeout
+                        )
+
+                    return SubmissionResult(
+                        wait_for_completion=_wait_for_completion
                     )
         except urllib3.exceptions.HTTPError as error:
             if kubernetes_context:
@@ -811,6 +806,8 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
                 f"Failed to upload Kubeflow pipeline: {error}. {msg}",
             )
 
+        return None
+
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.