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

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

zenml/integrations/lightning/orchestrators/lightning_orchestrator.py

@@ -16,14 +16,14 @@
 import os
 import tempfile
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
 from uuid import uuid4
 
 from lightning_sdk import Machine, Studio
 
+from zenml import __version__
 from zenml.constants import (
     ENV_ZENML_CUSTOM_SOURCE_ROOT,
-    ENV_ZENML_WHEEL_PACKAGE_NAME,
 )
 from zenml.entrypoints.step_entrypoint_configuration import (
     StepEntrypointConfiguration,
@@ -40,9 +40,8 @@ from zenml.integrations.lightning.orchestrators.utils import (
     sanitize_studio_name,
 )
 from zenml.logger import get_logger
+from zenml.orchestrators import BaseOrchestrator, SubmissionResult
 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 code_utils, io_utils, source_utils
 
 if TYPE_CHECKING:
@@ -52,48 +51,11 @@ if TYPE_CHECKING:
 
 logger = get_logger(__name__)
 ENV_ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID = "ZENML_LIGHTNING_ORCHESTRATOR_RUN_ID"
-ZENML_STEP_DEFAULT_ENTRYPOINT_COMMAND = "zenml.entrypoints.entrypoint"
 LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH = "."
 
 
-class LightningOrchestrator(WheeledOrchestrator):
-    """Base class for Orchestrator responsible for running pipelines remotely in a VM.
-
-    This orchestrator does not support running on a schedule.
-    """
-
-    @property
-    def validator(self) -> Optional[StackValidator]:
-        """Validates the stack.
-
-        In the remote case, checks that the stack contains a container registry,
-        image builder and only remote components.
-
-        Returns:
-            A `StackValidator` instance.
-        """
-
-        def _validate_remote_components(
-            stack: "Stack",
-        ) -> Tuple[bool, str]:
-            for component in stack.components.values():
-                if not component.config.is_local:
-                    continue
-
-                # return False, (
-                #    f"The Lightning orchestrator runs pipelines remotely, "
-                #    f"but the '{component.name}' {component.type.value} is "
-                #    "a local stack component and will not be available in "
-                #    "the Lightning step.\nPlease ensure that you always "
-                #    "use non-local stack components with the Lightning "
-                #    "orchestrator."
-                # )
-
-            return True, ""
-
-        return StackValidator(
-            custom_validation_function=_validate_remote_components,
-        )
+class LightningOrchestrator(BaseOrchestrator):
+    """Lightning orchestrator."""
 
     def _set_lightning_env_vars(
         self,
@@ -194,67 +156,39 @@ class LightningOrchestrator(WheeledOrchestrator):
         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 Lightning.
+    ) -> 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
-        task for each step (`_construct_lightning_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(
             LightningOrchestratorSettings, self.get_settings(deployment)
         )
         if deployment.schedule:
-            if (
-                deployment.schedule.catchup
-                or deployment.schedule.interval_second
-            ):
-                logger.warning(
-                    "Lightning orchestrator only uses schedules with the "
-                    "`cron_expression` property, with optional `start_time` and/or `end_time`. "
-                    "All other properties are ignored."
-                )
-            if deployment.schedule.cron_expression is None:
-                raise ValueError(
-                    "Property `cron_expression` must be set when passing "
-                    "schedule to a Lightning orchestrator."
-                )
-            if deployment.schedule.cron_expression:
-                raise ValueError(
-                    "Property `schedule_timezone` must be set when passing "
-                    "`cron_expression` to a Lightning orchestrator."
-                    "Lightning orchestrator requires a Java Timezone ID to run the pipeline on schedule."
-                    "Please refer to https://docs.oracle.com/middleware/1221/wcs/tag-ref/MISC/TimeZones.html for more information."
-                )
+            logger.warning(
+                "Lightning Orchestrator currently does not support the "
+                "use of schedules. The `schedule` will be ignored "
+                "and the pipeline will be run immediately."
+            )
 
         # Get deployment id
         deployment_id = deployment.id
@@ -262,13 +196,6 @@ class LightningOrchestrator(WheeledOrchestrator):
         pipeline_name = deployment.pipeline_configuration.name
         orchestrator_run_name = get_orchestrator_run_name(pipeline_name)
 
-        # Copy the repository to a temporary directory and add a setup.py file
-        # repository_temp_dir = (
-        #    self.copy_repository_to_temp_dir_and_add_setup_py()
-        # )
-
-        # Create a wheel for the package in the temporary directory
-        # wheel_path = self.create_wheel(temp_dir=repository_temp_dir)
         code_archive = code_utils.CodeArchive(
             root=source_utils.get_source_root()
         )
@@ -287,7 +214,6 @@ class LightningOrchestrator(WheeledOrchestrator):
         env_vars[ENV_ZENML_CUSTOM_SOURCE_ROOT] = (
             LIGHTNING_ZENML_DEFAULT_CUSTOM_REPOSITORY_PATH
         )
-        env_vars[ENV_ZENML_WHEEL_PACKAGE_NAME] = self.package_name
 
         # Create a line-by-line export of environment variables
         env_exports = "\n".join(
@@ -410,7 +336,7 @@ class LightningOrchestrator(WheeledOrchestrator):
                 studio.run(
                     f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {custom_command}"
                 )
-            # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
+
             logger.info("Running pipeline in async mode")
             studio.run(
                 f"nohup bash -c 'cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {entrypoint_string}' > log_{filename.rsplit('.', 2)[0]}.txt 2>&1 &"
@@ -430,6 +356,7 @@ class LightningOrchestrator(WheeledOrchestrator):
                 env_file_path,
             )
         os.unlink(env_file_path)
+        return None
 
     def _upload_and_run_pipeline(
         self,
@@ -507,8 +434,7 @@ class LightningOrchestrator(WheeledOrchestrator):
 
             studio.run("pip install uv")
             studio.run(f"uv pip install {requirements}")
-            studio.run("pip install zenml")
-            # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
+            studio.run(f"uv pip install zenml=={__version__}")
             for command in settings.custom_commands or []:
                 output = studio.run(
                     f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"
@@ -588,8 +514,7 @@ class LightningOrchestrator(WheeledOrchestrator):
 
         studio.run("pip install uv")
         studio.run(f"uv pip install {details['requirements']}")
-        studio.run("pip install zenml")
-        # studio.run(f"pip install {wheel_path.rsplit('/', 1)[-1]}")
+        studio.run(f"uv pip install zenml=={__version__}")
         for command in custom_commands or []:
             output = studio.run(
                 f"cd /teamspace/studios/this_studio/zenml_codes/{filename.rsplit('.', 2)[0]} && {command}"

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -15,7 +15,7 @@
 
 import os
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, cast
 from uuid import uuid4
 
 import sky
@@ -43,6 +43,7 @@ from zenml.integrations.skypilot.utils import (
 from zenml.logger import get_logger
 from zenml.orchestrators import (
     ContainerizedOrchestrator,
+    SubmissionResult,
 )
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
@@ -153,25 +154,33 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
             set: Whether to set the environment variables or not.
         """
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Runs each pipeline step in a separate Skypilot container.
+    ) -> 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:
             Exception: If the pipeline run fails.
             RuntimeError: If the code is running in a notebook.
+
+        Returns:
+            Optional submission result.
         """
         # First check whether the code is running in a notebook.
         if Environment.in_notebook():
@@ -366,7 +375,9 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     cluster_name,
                     **launch_kwargs,
                 )
-                sky_job_get(launch_job_id, settings.stream_logs, cluster_name)
+                return sky_job_get(
+                    launch_job_id, settings.stream_logs, cluster_name
+                )
 
             else:
                 # Prepare exec parameters with additional launch settings
@@ -399,7 +410,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     **exec_kwargs,
                 )
                 assert settings.cluster_name is not None
-                sky_job_get(
+                return sky_job_get(
                     exec_job_id, settings.stream_logs, settings.cluster_name
                 )
 

zenml/integrations/skypilot/utils.py

@@ -9,6 +9,7 @@ from zenml.integrations.skypilot.flavors.skypilot_orchestrator_base_vm_config im
     SkypilotBaseOrchestratorSettings,
 )
 from zenml.logger import get_logger
+from zenml.orchestrators import SubmissionResult
 
 logger = get_logger(__name__)
 
@@ -235,7 +236,9 @@ def prepare_launch_kwargs(
     return {k: v for k, v in launch_kwargs.items() if v is not None}
 
 
-def sky_job_get(request_id: str, stream_logs: bool, cluster_name: str) -> Any:
+def sky_job_get(
+    request_id: str, stream_logs: bool, cluster_name: str
+) -> Optional[SubmissionResult]:
     """Handle SkyPilot request results based on stream_logs setting.
 
     SkyPilot API exec and launch methods are asynchronous and return a request ID.
@@ -249,10 +252,7 @@ def sky_job_get(request_id: str, stream_logs: bool, cluster_name: str) -> Any:
         cluster_name: The name of the cluster to tail logs for.
 
     Returns:
-        The result of the SkyPilot operation.
-
-    Raises:
-        Exception: If the SkyPilot job fails.
+        Optional submission result.
     """
     if stream_logs:
         # Stream logs and wait for completion
@@ -261,13 +261,17 @@ def sky_job_get(request_id: str, stream_logs: bool, cluster_name: str) -> Any:
         # Just wait for completion without streaming logs
         job_id, _ = sky.get(request_id)
 
-    status = 0  # 0=Successful, 100=Failed
+    _wait_for_completion = None
     if stream_logs:
-        status = sky.tail_logs(
-            cluster_name=cluster_name, job_id=job_id, follow=True
-        )
-
-    if status != 0:
-        raise Exception(f"SkyPilot job {job_id} failed with status {status}")
 
-    return job_id
+        def _wait_for_completion() -> None:
+            status = 0  # 0=Successful, 100=Failed
+            status = sky.tail_logs(
+                cluster_name=cluster_name, job_id=job_id, follow=True
+            )
+            if status != 0:
+                raise Exception(
+                    f"SkyPilot job {job_id} failed with status {status}"
+                )
+
+    return SubmissionResult(wait_for_completion=_wait_for_completion)

zenml/integrations/tekton/orchestrators/tekton_orchestrator.py

@@ -46,7 +46,7 @@ from zenml.integrations.tekton.flavors.tekton_orchestrator_flavor import (
 )
 from zenml.io import fileio
 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
 from zenml.utils import io_utils, yaml_utils
@@ -455,28 +455,33 @@ class TektonOrchestrator(ContainerizedOrchestrator):
 
         return dsl.container_component(dynamic_func)
 
-    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 on Tekton.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
 
-        This function first compiles the ZenML pipeline into a Tekton yaml
-        and then applies this configuration to run the pipeline.
+        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 you try to run the pipelines in a notebook
                 environment.
+
+        Returns:
+            Optional submission result.
         """
         # First check whether the code running in a notebook
         if Environment.in_notebook():
@@ -640,7 +645,7 @@ class TektonOrchestrator(ContainerizedOrchestrator):
 
         # using the kfp client uploads the pipeline to Tekton 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,
@@ -651,7 +656,7 @@ class TektonOrchestrator(ContainerizedOrchestrator):
         deployment: "PipelineDeploymentResponse",
         pipeline_file_path: str,
         run_name: str,
-    ) -> None:
+    ) -> Optional[SubmissionResult]:
         """Tries to upload and run a KFP pipeline.
 
         Args:
@@ -661,6 +666,9 @@ class TektonOrchestrator(ContainerizedOrchestrator):
 
         Raises:
             RuntimeError: If Tekton API returns an error.
+
+        Returns:
+            Optional submission result.
         """
         pipeline_name = deployment.pipeline_configuration.name
         settings = cast(
@@ -756,8 +764,14 @@ class TektonOrchestrator(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:
@@ -779,6 +793,8 @@ class TektonOrchestrator(ContainerizedOrchestrator):
                 f"Failed to upload Tekton pipeline: {error}. {msg}",
             )
 
+        return None
+
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.
 

zenml/models/v2/core/step_run.py

@@ -129,6 +129,7 @@ class StepRunRequest(ProjectScopedRequest):
     parent_step_ids: List[UUID] = Field(
         title="The IDs of the parent steps of this step run.",
         default_factory=list,
+        deprecated=True,
     )
     inputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the input artifact versions of the step run.",

zenml/orchestrators/__init__.py

@@ -26,6 +26,7 @@ from zenml.orchestrators.base_orchestrator import (
     BaseOrchestrator,
     BaseOrchestratorConfig,
     BaseOrchestratorFlavor,
+    SubmissionResult,
 )
 from zenml.orchestrators.containerized_orchestrator import (
     ContainerizedOrchestrator,
@@ -52,4 +53,5 @@ __all__ = [
     "LocalOrchestratorFlavor",
     "LocalDockerOrchestrator",
     "LocalDockerOrchestratorFlavor",
+    "SubmissionResult",
 ]