zenml-nightly 0.83.1.dev20250624__py3-none-any.whl → 0.83.1.dev20250626__py3-none-any.whl

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/integration.py

@@ -13,16 +13,14 @@
 #  permissions and limitations under the License.
 """Base and meta classes for ZenML integrations."""
 
-import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
 
-import pkg_resources
-from pkg_resources import Requirement
+from packaging.requirements import Requirement
 
 from zenml.integrations.registry import integration_registry
 from zenml.logger import get_logger
 from zenml.stack.flavor import Flavor
-from zenml.utils.integration_utils import parse_requirement
+from zenml.utils.package_utils import get_dependencies, requirement_installed
 
 if TYPE_CHECKING:
     from zenml.plugins.base_plugin_flavor import BasePluginFlavor
@@ -69,65 +67,32 @@ class Integration(metaclass=IntegrationMeta):
         Returns:
             True if all required packages are installed, False otherwise.
         """
-        for r in cls.get_requirements():
-            try:
-                # First check if the base package is installed
-                dist = pkg_resources.get_distribution(r)
-
-                # Next, check if the dependencies (including extras) are
-                # installed
-                deps: List[Requirement] = []
-
-                _, extras = parse_requirement(r)
-                if extras:
-                    extra_list = extras[1:-1].split(",")
-                    for extra in extra_list:
-                        try:
-                            requirements = dist.requires(extras=[extra])  # type: ignore[arg-type]
-                        except pkg_resources.UnknownExtra as e:
-                            logger.debug(f"Unknown extra: {str(e)}")
-                            return False
-                        deps.extend(requirements)
-                else:
-                    deps = dist.requires()
-
-                for ri in deps:
-                    try:
-                        # Remove the "extra == ..." part from the requirement string
-                        cleaned_req = re.sub(
-                            r"; extra == \"\w+\"", "", str(ri)
-                        )
-                        pkg_resources.get_distribution(cleaned_req)
-                    except pkg_resources.DistributionNotFound as e:
-                        logger.debug(
-                            f"Unable to find required dependency "
-                            f"'{e.req}' for requirement '{r}' "
-                            f"necessary for integration '{cls.NAME}'."
-                        )
-                        return False
-                    except pkg_resources.VersionConflict as e:
-                        logger.debug(
-                            f"Package version '{e.dist}' does not match "
-                            f"version '{e.req}' required by '{r}' "
-                            f"necessary for integration '{cls.NAME}'."
-                        )
-                        return False
-
-            except pkg_resources.DistributionNotFound as e:
-                logger.debug(
-                    f"Unable to find required package '{e.req}' for "
-                    f"integration {cls.NAME}."
-                )
-                return False
-            except pkg_resources.VersionConflict as e:
+        for requirement in cls.get_requirements():
+            parsed_requirement = Requirement(requirement)
+
+            if not requirement_installed(parsed_requirement):
                 logger.debug(
-                    f"Package version '{e.dist}' does not match version "
-                    f"'{e.req}' necessary for integration {cls.NAME}."
+                    "Requirement '%s' for integration '%s' is not installed "
+                    "or installed with the wrong version.",
+                    requirement,
+                    cls.NAME,
                 )
                 return False
 
+            dependencies = get_dependencies(parsed_requirement)
+
+            for dependency in dependencies:
+                if not requirement_installed(dependency):
+                    logger.debug(
+                        "Requirement '%s' for integration '%s' is not "
+                        "installed or installed with the wrong version.",
+                        dependency,
+                        cls.NAME,
+                    )
+                    return False
+
         logger.debug(
-            f"Integration {cls.NAME} is installed correctly with "
+            f"Integration '{cls.NAME}' is installed correctly with "
             f"requirements {cls.get_requirements()}."
         )
         return True

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.
 

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -33,7 +33,6 @@
 import os
 from typing import (
     TYPE_CHECKING,
-    Any,
     Dict,
     List,
     Optional,
@@ -62,7 +61,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
 )
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
-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
 
@@ -388,24 +387,32 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         """
         return f"zenml-token-{deployment_id}"
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Runs the pipeline in Kubernetes.
+    ) -> 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.
 
         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 the Kubernetes orchestrator is not configured.
+            RuntimeError: If a schedule without cron expression is given.
+
+        Returns:
+            Optional submission result.
         """
         for step_name, step in deployment.step_configurations.items():
             if self.requires_resources_in_orchestration_environment(step):
@@ -536,7 +543,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 f"Scheduling Kubernetes run `{pod_name}` with CRON expression "
                 f'`"{cron_expression}"`.'
             )
-            return
+            return None
         else:
             # Create and run the orchestrator pod.
             pod_manifest = build_pod_manifest(
@@ -565,18 +572,23 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 startup_timeout=settings.pod_startup_timeout,
             )
 
-            # Wait for the orchestrator pod to finish and stream logs.
             if settings.synchronous:
-                logger.info(
-                    "Waiting for Kubernetes orchestrator pod to finish..."
-                )
-                kube_utils.wait_pod(
-                    kube_client_fn=self.get_kube_client,
-                    pod_name=pod_name,
-                    namespace=self.config.kubernetes_namespace,
-                    exit_condition_lambda=kube_utils.pod_is_done,
-                    timeout_sec=settings.timeout,
-                    stream_logs=True,
+
+                def _wait_for_run_to_finish() -> None:
+                    logger.info(
+                        "Waiting for Kubernetes orchestrator pod to finish..."
+                    )
+                    kube_utils.wait_pod(
+                        kube_client_fn=self.get_kube_client,
+                        pod_name=pod_name,
+                        namespace=self.config.kubernetes_namespace,
+                        exit_condition_lambda=kube_utils.pod_is_done,
+                        timeout_sec=settings.timeout,
+                        stream_logs=True,
+                    )
+
+                return SubmissionResult(
+                    wait_for_completion=_wait_for_run_to_finish
                 )
             else:
                 logger.info(
@@ -585,6 +597,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                     f"Run the following command to inspect the logs: "
                     f"`kubectl logs {pod_name} -n {self.config.kubernetes_namespace}`."
                 )
+                return None
 
     def _get_service_account_name(
         self, settings: KubernetesOrchestratorSettings