Add support for dynamic pipelines to the Vertex orchestrator

Author: schustmi

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 Kubernetes.
+
+        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,