Add support for dynamic pipelines to the Vertex orchestrator

Author: schustmi

docs/book/how-to/steps-pipelines/dynamic_pipelines.md

@@ -5,7 +5,7 @@ description: Write dynamic pipelines
 # Dynamic Pipelines (Experimental)
 
 {% hint style="warning" %}
-**Experimental Feature**: Dynamic pipelines are currently an experimental feature. There are known issues and limitations, and the interface is subject to change. This feature is only supported by the `local` and `kubernetes` orchestrators. If you encounter any issues or have feedback, please let us know at [https://github.com/zenml-io/zenml/issues](https://github.com/zenml-io/zenml/issues).
+**Experimental Feature**: Dynamic pipelines are currently an experimental feature. There are known issues and limitations, and the interface is subject to change. This feature is only supported by the `local`, `kubernetes`, `sagemaker` and `vertex` orchestrators. If you encounter any issues or have feedback, please let us know at [https://github.com/zenml-io/zenml/issues](https://github.com/zenml-io/zenml/issues).
 {% endhint %}
 
 {% hint style="info" %}
@@ -265,26 +265,11 @@ When running multiple steps concurrently using `step.submit()`, a failure in one
 Dynamic pipelines are currently only supported by:
 - `local` orchestrator
 - `kubernetes` orchestrator
+- `sagemaker` orchestrator
+- `vertex` orchestrator
 
 Other orchestrators will raise an error if you try to run a dynamic pipeline with them.
 
-### Remote Execution Requirement
-
-When running dynamic pipelines remotely (e.g., with the `kubernetes` orchestrator), you **must** include `depends_on` for at least one step in your pipeline definition. This is currently required due to a bug in remote execution.
-
-{% hint style="warning" %}
-**Required for Remote Execution**: Without `depends_on`, remote execution will fail. This requirement does not apply when running locally with the `local` orchestrator.
-{% endhint %}
-
-For example:
-
-```python
-@pipeline(dynamic=True, depends_on=[some_step])
-def dynamic_pipeline():
-    some_step()
-    # ... rest of your pipeline
-```
-
 ### Artifact Loading
 
 When you call `.load()` on an artifact in a dynamic pipeline, it synchronously loads the data. For large artifacts or when you want to maintain parallelism, consider passing the step outputs (future or artifact) directly to downstream steps instead of loading them.

src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -1050,10 +1050,10 @@ def _wait_for_completion() -> None:
             metadata=metadata,
         )
 
-    def launch_dynamic_step(
+    def run_isolated_step(
         self, step_run_info: "StepRunInfo", environment: Dict[str, str]
     ) -> None:
-        """Launch a dynamic step.
+        """Runs an isolated step on Sagemaker.
 
         Args:
             step_run_info: The step run information.

src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -63,17 +63,20 @@
 from kfp.compiler import Compiler
 from kfp.dsl.base_component import BaseComponent
 
+from zenml import __version__
 from zenml.config.resource_settings import ResourceSettings
 from zenml.constants import (
     METADATA_ORCHESTRATOR_LOGS_URL,
     METADATA_ORCHESTRATOR_RUN_ID,
     METADATA_ORCHESTRATOR_URL,
+    ORCHESTRATOR_DOCKER_IMAGE_KEY,
 )
 from zenml.entrypoints import StepEntrypointConfiguration
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.integrations.gcp import GCP_ARTIFACT_STORE_FLAVOR
 from zenml.integrations.gcp.constants import (
     GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+    VERTEX_ENDPOINT_SUFFIX,
 )
 from zenml.integrations.gcp.flavors.vertex_orchestrator_flavor import (
     VertexOrchestratorConfig,
@@ -82,6 +85,10 @@
 from zenml.integrations.gcp.google_credentials_mixin import (
     GoogleCredentialsMixin,
 )
+from zenml.integrations.gcp.utils import (
+    build_job_request,
+    monitor_job,
+)
 from zenml.integrations.gcp.vertex_custom_job_parameters import (
     VertexCustomJobParameters,
 )
@@ -90,11 +97,18 @@
 from zenml.metadata.metadata_types import MetadataType, Uri
 from zenml.orchestrators import ContainerizedOrchestrator, SubmissionResult
 from zenml.orchestrators.utils import get_orchestrator_run_name
+from zenml.pipelines.dynamic.entrypoint_configuration import (
+    DynamicPipelineEntrypointConfiguration,
+)
 from zenml.stack.stack_validator import StackValidator
+from zenml.step_operators.step_operator_entrypoint_configuration import (
+    StepOperatorEntrypointConfiguration,
+)
 from zenml.utils.io_utils import get_global_config_directory
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
+    from zenml.config.step_run_info import StepRunInfo
     from zenml.models import (
         PipelineRunResponse,
         PipelineSnapshotResponse,
@@ -621,6 +635,151 @@ def dynamic_pipeline() -> None:
             schedule=snapshot.schedule,
         )
 
+    def submit_dynamic_pipeline(
+        self,
+        snapshot: "PipelineSnapshotResponse",
+        stack: "Stack",
+        environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
+    ) -> Optional[SubmissionResult]:
+        """Submits a dynamic pipeline to the orchestrator.
+
+        Args:
+            snapshot: The pipeline snapshot to submit.
+            stack: The stack the pipeline will run on.
+            environment: Environment variables to set in the orchestration
+                environment.
+            placeholder_run: An optional placeholder run.
+
+        Raises:
+            RuntimeError: If the snapshot contains a schedule.
+
+        Returns:
+            Optional submission result.
+        """
+        if snapshot.schedule:
+            raise RuntimeError(
+                "Scheduling dynamic pipelines is not supported for the "
+                "Vertex orchestrator yet."
+            )
+
+        settings = cast(
+            VertexOrchestratorSettings, self.get_settings(snapshot)
+        )
+
+        command = (
+            DynamicPipelineEntrypointConfiguration.get_entrypoint_command()
+        )
+        args = DynamicPipelineEntrypointConfiguration.get_entrypoint_arguments(
+            snapshot_id=snapshot.id,
+            run_id=placeholder_run.id if placeholder_run else None,
+        )
+
+        image = self.get_image(snapshot=snapshot)
+        labels = settings.labels.copy()
+        labels["source"] = f"zenml-{__version__.replace('.', '_')}"
+
+        job_request = build_job_request(
+            display_name=get_orchestrator_run_name(
+                pipeline_name=snapshot.pipeline_configuration.name
+            ),
+            image=image,
+            entrypoint_command=command + args,
+            custom_job_settings=settings.custom_job_parameters
+            or VertexCustomJobParameters(),
+            resource_settings=snapshot.pipeline_configuration.resource_settings,
+            environment=environment,
+            labels=labels,
+            encryption_spec_key_name=self.config.encryption_spec_key_name,
+            service_account=self.config.workload_service_account,
+            network=self.config.network,
+        )
+
+        credentials, project_id = self._get_authentication()
+        client_options = {
+            "api_endpoint": self.config.location + VERTEX_ENDPOINT_SUFFIX
+        }
+        client = aiplatform.gapic.JobServiceClient(
+            credentials=credentials, client_options=client_options
+        )
+        parent = f"projects/{project_id}/locations/{self.config.location}"
+        logger.info(
+            "Submitting custom job='%s', path='%s' to Vertex AI Training.",
+            job_request["display_name"],
+            parent,
+        )
+        job = client.create_custom_job(parent=parent, custom_job=job_request)
+
+        wait_for_completion = None
+        if settings.synchronous:
+            wait_for_completion = lambda: monitor_job(
+                job_id=job.name,
+                credentials_source=self,
+                client_options=client_options,
+            )
+
+        return SubmissionResult(
+            wait_for_completion=wait_for_completion,
+        )
+
+    def run_isolated_step(
+        self, step_run_info: "StepRunInfo", environment: Dict[str, str]
+    ) -> None:
+        """Runs an isolated step on Vertex.
+
+        Args:
+            step_run_info: The step run information.
+            environment: The environment variables to set.
+        """
+        settings = cast(
+            VertexOrchestratorSettings, self.get_settings(step_run_info)
+        )
+
+        image = step_run_info.get_image(key=ORCHESTRATOR_DOCKER_IMAGE_KEY)
+        command = StepOperatorEntrypointConfiguration.get_entrypoint_command()
+        args = StepOperatorEntrypointConfiguration.get_entrypoint_arguments(
+            step_name=step_run_info.pipeline_step_name,
+            snapshot_id=(step_run_info.snapshot.id),
+            step_run_id=str(step_run_info.step_run_id),
+        )
+
+        labels = settings.labels.copy()
+        labels["source"] = f"zenml-{__version__.replace('.', '_')}"
+
+        job_request = build_job_request(
+            display_name=f"{step_run_info.run_name}-{step_run_info.pipeline_step_name}",
+            image=image,
+            entrypoint_command=command + args,
+            custom_job_settings=settings.custom_job_parameters
+            or VertexCustomJobParameters(),
+            resource_settings=step_run_info.config.resource_settings,
+            environment=environment,
+            labels=labels,
+            encryption_spec_key_name=self.config.encryption_spec_key_name,
+            service_account=self.config.workload_service_account,
+            network=self.config.network,
+        )
+
+        credentials, project_id = self._get_authentication()
+        client_options = {
+            "api_endpoint": self.config.location + VERTEX_ENDPOINT_SUFFIX
+        }
+        client = aiplatform.gapic.JobServiceClient(
+            credentials=credentials, client_options=client_options
+        )
+        parent = f"projects/{project_id}/locations/{self.config.location}"
+        logger.info(
+            "Submitting custom job='%s', path='%s' to Vertex AI Training.",
+            job_request["display_name"],
+            parent,
+        )
+        job = client.create_custom_job(parent=parent, custom_job=job_request)
+        monitor_job(
+            job_id=job.name,
+            credentials_source=self,
+            client_options=client_options,
+        )
+
     def _upload_and_run_pipeline(
         self,
         pipeline_name: str,
@@ -786,19 +945,19 @@ def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.
 
         Raises:
-            RuntimeError: If the environment variable specifying the run id
-                is not set.
+            RuntimeError: If the orchestrator run id cannot be read from the
+                environment.
 
         Returns:
             The orchestrator run id.
         """
-        try:
-            return os.environ[ENV_ZENML_VERTEX_RUN_ID]
-        except KeyError:
-            raise RuntimeError(
-                "Unable to read run id from environment variable "
-                f"{ENV_ZENML_VERTEX_RUN_ID}."
-            )
+        for env in [ENV_ZENML_VERTEX_RUN_ID, "CLOUD_ML_JOB_ID"]:
+            if env in os.environ:
+                return os.environ[env]
+
+        raise RuntimeError(
+            "Unable to get orchestrator run id from environment."
+        )
 
     def get_pipeline_run_metadata(
         self, run_id: UUID