zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -31,21 +31,35 @@
 
 import os
 import re
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Type, cast
+import types
+import urllib
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    cast,
+)
 from uuid import UUID
 
-import kfp
 from google.api_core import exceptions as google_exceptions
 from google.cloud import aiplatform
+from google.cloud.aiplatform_v1.types import PipelineState
 from kfp import dsl
-from kfp.v2 import dsl as dslv2
-from kfp.v2.compiler import Compiler as KFPV2Compiler
+from kfp.compiler import Compiler
 
+from zenml.config.resource_settings import ResourceSettings
 from zenml.constants import (
+    METADATA_ORCHESTRATOR_LOGS_URL,
+    METADATA_ORCHESTRATOR_RUN_ID,
     METADATA_ORCHESTRATOR_URL,
 )
 from zenml.entrypoints import StepEntrypointConfiguration
-from zenml.enums import StackComponentType
+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,
@@ -54,39 +68,26 @@ from zenml.integrations.gcp.flavors.vertex_orchestrator_flavor import (
     VertexOrchestratorConfig,
     VertexOrchestratorSettings,
 )
-from zenml.integrations.gcp.google_cloud_function import create_cloud_function
-from zenml.integrations.gcp.google_cloud_scheduler import create_scheduler_job
 from zenml.integrations.gcp.google_credentials_mixin import (
     GoogleCredentialsMixin,
 )
-from zenml.integrations.gcp.orchestrators import vertex_scheduler
-from zenml.integrations.gcp.orchestrators.vertex_scheduler.main import (
-    ENABLE_CACHING,
-    ENCRYPTION_SPEC_KEY_NAME,
-    JOB_ID,
-    LABELS,
-    LOCATION,
-    NETWORK,
-    PARAMETER_VALUES,
-    PIPELINE_ROOT,
-    PROJECT,
-    TEMPLATE_PATH,
-    WORKLOAD_SERVICE_ACCOUNT,
-)
-from zenml.integrations.kubeflow.utils import apply_pod_settings
 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.utils import get_orchestrator_run_name
 from zenml.stack.stack_validator import StackValidator
+from zenml.utils import yaml_utils
 from zenml.utils.io_utils import get_global_config_directory
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
-    from zenml.models import PipelineDeploymentResponse, ScheduleResponse
+    from zenml.models import (
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+        ScheduleResponse,
+    )
     from zenml.stack import Stack
-    from zenml.steps import ResourceSettings
 
 logger = get_logger(__name__)
 ENV_ZENML_VERTEX_RUN_ID = "ZENML_VERTEX_RUN_ID"
@@ -249,14 +250,12 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         if deployment.schedule:
             if (
                 deployment.schedule.catchup
-                or deployment.schedule.start_time
-                or deployment.schedule.end_time
                 or deployment.schedule.interval_second
             ):
                 logger.warning(
                     "Vertex orchestrator only uses schedules with the "
-                    "`cron_expression` property. All other properties "
-                    "are ignored."
+                    "`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(
@@ -264,61 +263,55 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     "schedule to a Vertex orchestrator."
                 )
 
-    def _configure_container_resources(
+    def _create_dynamic_component(
         self,
-        container_op: dsl.ContainerOp,
-        resource_settings: "ResourceSettings",
-        node_selector_constraint: Optional[Tuple[str, str]] = None,
-    ) -> None:
-        """Adds resource requirements to the container.
+        image: str,
+        command: List[str],
+        arguments: List[str],
+        component_name: str,
+    ) -> dsl.PipelineTask:
+        """Creates a dynamic container component for a Vertex pipeline.
 
         Args:
-            container_op: The kubeflow container operation to configure.
-            resource_settings: The resource settings to use for this
-                container.
-            node_selector_constraint: Node selector constraint to apply to
-                the container.
+            image: The image to use for the component.
+            command: The command to use for the component.
+            arguments: The arguments to use for the component.
+            component_name: The name of the component.
+
+        Returns:
+            The dynamic container component.
         """
-        # Set optional CPU, RAM and GPU constraints for the pipeline
 
-        cpu_limit = resource_settings.cpu_count or self.config.cpu_limit
+        def dynamic_container_component() -> dsl.ContainerSpec:
+            """Dynamic container component.
 
-        if cpu_limit is not None:
-            container_op = container_op.set_cpu_limit(str(cpu_limit))
-
-        memory_limit = (
-            resource_settings.memory[:-1]
-            if resource_settings.memory
-            else self.config.memory_limit
-        )
-        if memory_limit is not None:
-            container_op = container_op.set_memory_limit(memory_limit)
+            Returns:
+                The dynamic container component.
+            """
+            return dsl.ContainerSpec(
+                image=image,
+                command=command,
+                args=arguments,
+            )
 
-        gpu_limit = (
-            resource_settings.gpu_count
-            if resource_settings.gpu_count is not None
-            else self.config.gpu_limit
+        # Change the name of the function
+        new_container_spec_func = types.FunctionType(
+            dynamic_container_component.__code__,
+            dynamic_container_component.__globals__,
+            name=component_name,
+            argdefs=dynamic_container_component.__defaults__,
+            closure=dynamic_container_component.__closure__,
         )
-        if gpu_limit is not None and gpu_limit > 0:
-            container_op = container_op.set_gpu_limit(gpu_limit)
+        pipeline_task = dsl.container_component(new_container_spec_func)
 
-        if node_selector_constraint:
-            constraint_label, value = node_selector_constraint
-            if not (
-                constraint_label
-                == GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
-                and gpu_limit == 0
-            ):
-                container_op.add_node_selector_constraint(
-                    constraint_label, value
-                )
+        return pipeline_task
 
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> Any:
+    ) -> Iterator[Dict[str, MetadataType]]:
         """Creates a KFP JSON pipeline.
 
         # noqa: DAR402
@@ -353,12 +346,15 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 environment.
 
         Raises:
-            ValueError: If the attribute `pipeline_root` is not set and it
+            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.
         """
         orchestrator_run_name = get_orchestrator_run_name(
             pipeline_name=deployment.pipeline_configuration.name
@@ -381,77 +377,166 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         else:
             self._pipeline_root = self.config.pipeline_root
 
-        def _construct_kfp_pipeline() -> None:
-            """Create a `ContainerOp` for each step.
-
-            This should contain the name of the Docker image and configures the
-            entrypoint of the Docker image to run the step.
+        def _create_dynamic_pipeline() -> Any:
+            """Create a dynamic pipeline including each step.
 
-            Additionally, this gives each `ContainerOp` information about its
-            direct downstream steps.
-
-            If this callable is passed to the `compile()` method of
-            `KFPV2Compiler` all `dsl.ContainerOp` instances will be
-            automatically added to a singular `dsl.Pipeline` instance.
+            Returns:
+                pipeline_func
             """
-            command = StepEntrypointConfiguration.get_entrypoint_command()
-            step_name_to_container_op: Dict[str, dsl.ContainerOp] = {}
+            step_name_to_dynamic_component: Dict[str, Any] = {}
 
             for step_name, step in deployment.step_configurations.items():
                 image = self.get_image(
-                    deployment=deployment, step_name=step_name
+                    deployment=deployment,
+                    step_name=step_name,
                 )
+                command = StepEntrypointConfiguration.get_entrypoint_command()
                 arguments = (
                     StepEntrypointConfiguration.get_entrypoint_arguments(
-                        step_name=step_name, deployment_id=deployment.id
+                        step_name=step_name,
+                        deployment_id=deployment.id,
                     )
                 )
-
-                # Create the `ContainerOp` for the step. Using the
-                # `dsl.ContainerOp`
-                # class directly is deprecated when using the Kubeflow SDK v2.
-                container_op = kfp.components.load_component_from_text(
-                    f"""
-                    name: {step_name}
-                    implementation:
-                        container:
-                            image: {image}
-                            command: {command + arguments}"""
-                )()
-
-                container_op.set_env_variable(
-                    name=ENV_ZENML_VERTEX_RUN_ID,
-                    value=dslv2.PIPELINE_JOB_NAME_PLACEHOLDER,
+                dynamic_component = self._create_dynamic_component(
+                    image, command, arguments, step_name
                 )
-
-                for key, value in environment.items():
-                    container_op.set_env_variable(name=key, value=value)
-
-                # Set upstream tasks as a dependency of the current step
-                for upstream_step_name in step.spec.upstream_steps:
-                    upstream_container_op = step_name_to_container_op[
-                        upstream_step_name
+                step_settings = cast(
+                    VertexOrchestratorSettings, self.get_settings(step)
+                )
+                pod_settings = step_settings.pod_settings
+                if pod_settings:
+                    if pod_settings.host_ipc:
+                        logger.warning(
+                            "Host IPC is set to `True` but not supported in "
+                            "this orchestrator. Ignoring..."
+                        )
+                    if pod_settings.affinity:
+                        logger.warning(
+                            "Affinity is set but not supported in Vertex with "
+                            "Kubeflow Pipelines 2.x. Ignoring..."
+                        )
+                    if pod_settings.tolerations:
+                        logger.warning(
+                            "Tolerations are set but not supported in "
+                            "Vertex with Kubeflow Pipelines 2.x. Ignoring..."
+                        )
+                    if pod_settings.volumes:
+                        logger.warning(
+                            "Volumes are set but not supported in Vertex with "
+                            "Kubeflow Pipelines 2.x. Ignoring..."
+                        )
+                    if pod_settings.volume_mounts:
+                        logger.warning(
+                            "Volume mounts are set but not supported in "
+                            "Vertex with Kubeflow Pipelines 2.x. Ignoring..."
+                        )
+                    for key in pod_settings.node_selectors:
+                        if (
+                            key
+                            != GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                        ):
+                            logger.warning(
+                                "Vertex only allows the %s node selector, "
+                                "ignoring the node selector %s.",
+                                GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                                key,
+                            )
+
+                step_name_to_dynamic_component[step_name] = dynamic_component
+
+            @dsl.pipeline(  # type: ignore[misc]
+                display_name=orchestrator_run_name,
+            )
+            def dynamic_pipeline() -> None:
+                """Dynamic pipeline."""
+                # iterate through the components one by one
+                # (from step_name_to_dynamic_component)
+                for (
+                    component_name,
+                    component,
+                ) in step_name_to_dynamic_component.items():
+                    # for each component, check to see what other steps are
+                    # upstream of it
+                    step = deployment.step_configurations[component_name]
+                    upstream_step_components = [
+                        step_name_to_dynamic_component[upstream_step_name]
+                        for upstream_step_name in step.spec.upstream_steps
                     ]
-                    container_op.after(upstream_container_op)
+                    task = (
+                        component()
+                        .set_display_name(
+                            name=component_name,
+                        )
+                        .set_caching_options(enable_caching=False)
+                        .set_env_variable(
+                            name=ENV_ZENML_VERTEX_RUN_ID,
+                            value=dsl.PIPELINE_JOB_NAME_PLACEHOLDER,
+                        )
+                        .after(*upstream_step_components)
+                    )
 
-                settings = cast(
-                    VertexOrchestratorSettings,
-                    self.get_settings(step),
-                )
-                if settings.pod_settings:
-                    apply_pod_settings(
-                        container_op=container_op,
-                        settings=settings.pod_settings,
+                    step_settings = cast(
+                        VertexOrchestratorSettings, self.get_settings(step)
+                    )
+                    pod_settings = step_settings.pod_settings
+
+                    node_selector_constraint: Optional[Tuple[str, str]] = None
+                    if pod_settings and (
+                        GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                        in pod_settings.node_selectors.keys()
+                    ):
+                        node_selector_constraint = (
+                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                            pod_settings.node_selectors[
+                                GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL
+                            ],
+                        )
+                    elif step_settings.node_selector_constraint:
+                        node_selector_constraint = (
+                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
+                            step_settings.node_selector_constraint[1],
+                        )
+
+                    self._configure_container_resources(
+                        dynamic_component=task,
+                        resource_settings=step.config.resource_settings,
+                        node_selector_constraint=node_selector_constraint,
                     )
 
-                self._configure_container_resources(
-                    container_op=container_op,
-                    resource_settings=step.config.resource_settings,
-                    node_selector_constraint=settings.node_selector_constraint,
-                )
-                container_op.set_caching_options(enable_caching=False)
+            return dynamic_pipeline
 
-                step_name_to_container_op[step_name] = container_op
+        def _update_json_with_environment(
+            yaml_file_path: str, environment: Dict[str, str]
+        ) -> None:
+            """Updates the env section of the steps in the YAML file with the given environment variables.
+
+            Args:
+                yaml_file_path: The path to the YAML file to update.
+                environment: A dictionary of environment variables to add.
+            """
+            pipeline_definition = yaml_utils.read_json(pipeline_file_path)
+
+            # Iterate through each component and add the environment variables
+            for executor in pipeline_definition["deploymentSpec"]["executors"]:
+                if (
+                    "container"
+                    in pipeline_definition["deploymentSpec"]["executors"][
+                        executor
+                    ]
+                ):
+                    container = pipeline_definition["deploymentSpec"][
+                        "executors"
+                    ][executor]["container"]
+                    if "env" not in container:
+                        container["env"] = []
+                    for key, value in environment.items():
+                        container["env"].append({"name": key, "value": value})
+
+            yaml_utils.write_json(pipeline_file_path, pipeline_definition)
+
+            print(
+                f"Updated YAML file with environment variables at {yaml_file_path}"
+            )
 
         # Save the generated pipeline to a file.
         fileio.makedirs(self.pipeline_directory)
@@ -463,13 +548,17 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         # Compile the pipeline using the Kubeflow SDK V2 compiler that allows
         # to generate a JSON representation of the pipeline that can be later
         # upload to Vertex AI Pipelines service.
-        KFPV2Compiler().compile(
-            pipeline_func=_construct_kfp_pipeline,
+        Compiler().compile(
+            pipeline_func=_create_dynamic_pipeline(),
             package_path=pipeline_file_path,
             pipeline_name=_clean_pipeline_name(
                 deployment.pipeline_configuration.name
             ),
         )
+
+        # Let's update the YAML file with the environment variables
+        _update_json_with_environment(pipeline_file_path, environment)
+
         logger.info(
             "Writing Vertex workflow definition to `%s`.", pipeline_file_path
         )
@@ -478,144 +567,16 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             VertexOrchestratorSettings, self.get_settings(deployment)
         )
 
-        if deployment.schedule:
-            logger.info(
-                "Scheduling job using Google Cloud Scheduler and Google "
-                "Cloud Functions..."
-            )
-            self._upload_and_schedule_pipeline(
-                pipeline_name=deployment.pipeline_configuration.name,
-                run_name=orchestrator_run_name,
-                stack=stack,
-                schedule=deployment.schedule,
-                pipeline_file_path=pipeline_file_path,
-                settings=settings,
-            )
-
-        else:
-            logger.info("No schedule detected. Creating one-off vertex job...")
-            # Using the Google Cloud AIPlatform client, upload and execute the
-            # pipeline
-            # on the Vertex AI Pipelines service.
-            self._upload_and_run_pipeline(
-                pipeline_name=deployment.pipeline_configuration.name,
-                pipeline_file_path=pipeline_file_path,
-                run_name=orchestrator_run_name,
-                settings=settings,
-            )
-
-    def _upload_and_schedule_pipeline(
-        self,
-        pipeline_name: str,
-        run_name: str,
-        stack: "Stack",
-        schedule: "ScheduleResponse",
-        pipeline_file_path: str,
-        settings: VertexOrchestratorSettings,
-    ) -> None:
-        """Uploads and schedules pipeline on GCP.
-
-        Args:
-            pipeline_name: Name of the pipeline.
-            run_name: Orchestrator run name.
-            stack: The stack the pipeline will run on.
-            schedule: The schedule the pipeline will run on.
-            pipeline_file_path: Path of the JSON file containing the compiled
-                Kubeflow pipeline (compiled with Kubeflow SDK v2).
-            settings: Pipeline level settings for this orchestrator.
-
-        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`.
-        """
-        # First, do some validation
-        artifact_store = stack.artifact_store
-        if artifact_store.flavor != GCP_ARTIFACT_STORE_FLAVOR:
-            raise ValueError(
-                "Currently, the Vertex AI orchestrator only supports "
-                "scheduled runs in combination with an artifact store of "
-                f"flavor: {GCP_ARTIFACT_STORE_FLAVOR}. The current stacks "
-                f"artifact store is of flavor: {artifact_store.flavor}. "
-                "Please update your stack accordingly."
-            )
-
-        # Get the credentials that would be used to create resources.
-        credentials, project_id = self._get_authentication()
-
-        scheduler_service_account_email: Optional[str] = None
-        if self.config.scheduler_service_account:
-            scheduler_service_account_email = (
-                self.config.scheduler_service_account
-            )
-        elif hasattr(credentials, "signer_email"):
-            scheduler_service_account_email = credentials.signer_email
-        else:
-            scheduler_service_account_email = (
-                self.config.function_service_account
-                or self.config.workload_service_account
-            )
-
-        if not scheduler_service_account_email:
-            raise ValueError(
-                "A GCP service account is required to schedule a pipeline run. "
-                "The credentials used to authenticate with GCP do not have a "
-                "service account associated with them and a service account "
-                "was not configured in the `scheduler_service_account` field "
-                "of the orchestrator config. Please update your orchestrator "
-                "configuration or credentials accordingly."
-            )
-
-        # Copy over the scheduled pipeline to the artifact store
-        artifact_store_base_uri = f"{artifact_store.path.rstrip('/')}/vertex_scheduled_pipelines/{pipeline_name}/{run_name}"
-        artifact_store_pipeline_uri = (
-            f"{artifact_store_base_uri}/vertex_pipeline.json"
-        )
-        fileio.copy(pipeline_file_path, artifact_store_pipeline_uri)
-        logger.info(
-            "The scheduled pipeline representation has been "
-            "automatically copied to this path of the `GCPArtifactStore`: "
-            f"{artifact_store_pipeline_uri}",
-        )
-
-        # Create cloud function
-        function_uri = create_cloud_function(
-            directory_path=vertex_scheduler.__path__[0],  # fixed path
-            upload_path=f"{artifact_store_base_uri}/code.zip",
-            project=project_id,
-            location=self.config.location,
-            function_name=run_name,
-            credentials=credentials,
-            function_service_account_email=self.config.function_service_account,
-        )
-
-        # Create the scheduler job
-        body = {
-            TEMPLATE_PATH: artifact_store_pipeline_uri,
-            JOB_ID: _clean_pipeline_name(pipeline_name),
-            PIPELINE_ROOT: self._pipeline_root,
-            PARAMETER_VALUES: None,
-            ENABLE_CACHING: False,
-            ENCRYPTION_SPEC_KEY_NAME: self.config.encryption_spec_key_name,
-            LABELS: settings.labels,
-            PROJECT: project_id,
-            LOCATION: self.config.location,
-            WORKLOAD_SERVICE_ACCOUNT: self.config.workload_service_account,
-            NETWORK: self.config.network,
-        }
-
-        create_scheduler_job(
-            project=project_id,
-            region=self.config.location,
-            http_uri=function_uri,
-            body=body,
-            schedule=str(schedule.cron_expression),
-            credentials=credentials,
-            service_account_email=scheduler_service_account_email,
-        )
+        # Using the Google Cloud AIPlatform client, upload and execute the
+        # pipeline on the Vertex AI Pipelines service.
+        if metadata := 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,
@@ -623,7 +584,8 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         pipeline_file_path: str,
         run_name: str,
         settings: VertexOrchestratorSettings,
-    ) -> None:
+        schedule: Optional["ScheduleResponse"] = None,
+    ) -> Iterator[Dict[str, MetadataType]]:
         """Uploads and run the pipeline on the Vertex AI Pipelines service.
 
         Args:
@@ -632,6 +594,14 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 Kubeflow pipeline (compiled with Kubeflow SDK v2).
             run_name: Orchestrator run name.
             settings: Pipeline level settings for this orchestrator.
+            schedule: The schedule the pipeline will run on.
+
+        Raises:
+            RuntimeError: If the Vertex Orchestrator fails to provision or any
+                other Runtime errors.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
         """
         # We have to replace the hyphens in the run name with underscores
         # and lower case the string, because the Vertex AI Pipelines service
@@ -639,8 +609,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         job_id = _clean_pipeline_name(run_name)
 
         # Get the credentials that would be used to create the Vertex AI
-        # Pipelines
-        # job.
+        # Pipelines job.
         credentials, project_id = self._get_authentication()
 
         # Instantiate the Vertex AI Pipelines job
@@ -658,52 +627,74 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             location=self.config.location,
         )
 
-        logger.info(
-            "Submitting pipeline job with job_id `%s` to Vertex AI Pipelines "
-            "service.",
-            job_id,
-        )
+        if self.config.workload_service_account:
+            logger.info(
+                "The Vertex AI Pipelines job workload will be executed "
+                "using the `%s` "
+                "service account.",
+                self.config.workload_service_account,
+            )
+        if self.config.network:
+            logger.info(
+                "The Vertex AI Pipelines job will be peered with the `%s` "
+                "network.",
+                self.config.network,
+            )
 
-        # Submit the job to Vertex AI Pipelines service.
         try:
-            if self.config.workload_service_account:
+            if schedule:
                 logger.info(
-                    "The Vertex AI Pipelines job workload will be executed "
-                    "using the `%s` "
-                    "service account.",
-                    self.config.workload_service_account,
+                    "Scheduling job using native Vertex AI Pipelines "
+                    "scheduling..."
+                )
+                run.create_schedule(
+                    display_name=schedule.name,
+                    cron=schedule.cron_expression,
+                    start_time=schedule.utc_start_time,
+                    end_time=schedule.utc_end_time,
+                    service_account=self.config.workload_service_account,
+                    network=self.config.network,
                 )
 
-            if self.config.network:
+            else:
                 logger.info(
-                    "The Vertex AI Pipelines job will be peered with the `%s` "
-                    "network.",
-                    self.config.network,
+                    "No schedule detected. Creating one-off Vertex job..."
+                )
+                logger.info(
+                    "Submitting pipeline job with job_id `%s` to Vertex AI "
+                    "Pipelines service.",
+                    job_id,
                 )
 
-            run.submit(
-                service_account=self.config.workload_service_account,
-                network=self.config.network,
-            )
-            logger.info(
-                "View the Vertex AI Pipelines job at %s", run._dashboard_uri()
-            )
-
-            if settings.synchronous:
+                # Submit the job to Vertex AI Pipelines service.
+                run.submit(
+                    service_account=self.config.workload_service_account,
+                    network=self.config.network,
+                )
                 logger.info(
-                    "Waiting for the Vertex AI Pipelines job to finish..."
+                    "View the Vertex AI Pipelines job at %s",
+                    run._dashboard_uri(),
                 )
-                run.wait()
+
+                # Yield metadata based on the generated job object
+                yield from self.compute_metadata(run)
+
+                if settings.synchronous:
+                    logger.info(
+                        "Waiting for the Vertex AI Pipelines job to finish..."
+                    )
+                    run.wait()
 
         except google_exceptions.ClientError as e:
-            logger.warning(
-                "Failed to create the Vertex AI Pipelines job: %s", e
+            logger.error("Failed to create the Vertex AI Pipelines job: %s", e)
+            raise RuntimeError(
+                f"Failed to create the Vertex AI Pipelines job: {e}"
             )
-
         except RuntimeError as e:
             logger.error(
                 "The Vertex AI Pipelines job execution has failed: %s", e
             )
+            raise
 
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.
@@ -744,3 +735,235 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
         return {
             METADATA_ORCHESTRATOR_URL: Uri(run_url),
         }
+
+    def _configure_container_resources(
+        self,
+        dynamic_component: dsl.PipelineTask,
+        resource_settings: "ResourceSettings",
+        node_selector_constraint: Optional[Tuple[str, str]] = None,
+    ) -> dsl.PipelineTask:
+        """Adds resource requirements to the container.
+
+        Args:
+            dynamic_component: The dynamic component to add the resource
+                settings to.
+            resource_settings: The resource settings to use for this
+                container.
+            node_selector_constraint: Node selector constraint to apply to
+                the container.
+
+        Returns:
+            The dynamic component with the resource settings applied.
+        """
+        # Set optional CPU, RAM and GPU constraints for the pipeline
+        cpu_limit = None
+        if resource_settings:
+            cpu_limit = resource_settings.cpu_count or self.config.cpu_limit
+
+        if cpu_limit is not None:
+            dynamic_component = dynamic_component.set_cpu_limit(str(cpu_limit))
+
+        memory_limit = (
+            resource_settings.memory[:-1]
+            if resource_settings.memory
+            else self.config.memory_limit
+        )
+        if memory_limit is not None:
+            dynamic_component = dynamic_component.set_memory_limit(
+                memory_limit
+            )
+
+        gpu_limit = (
+            resource_settings.gpu_count
+            if resource_settings.gpu_count is not None
+            else self.config.gpu_limit
+        )
+
+        if node_selector_constraint:
+            _, value = node_selector_constraint
+            if gpu_limit is not None and gpu_limit > 0:
+                dynamic_component = (
+                    dynamic_component.set_accelerator_type(value)
+                    .set_accelerator_limit(gpu_limit)
+                    .set_gpu_limit(gpu_limit)
+                )
+            else:
+                logger.warning(
+                    "Accelerator type %s specified, but the GPU limit is not "
+                    "set or set to 0. The accelerator type will be ignored. "
+                    "To fix this warning, either remove the specified "
+                    "accelerator type or set the `gpu_count` using the "
+                    "ResourceSettings (https://docs.zenml.io/how-to/advanced-topics/training-with-gpus)."
+                )
+
+        return dynamic_component
+
+    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: The run that was executed by this orchestrator.
+
+        Returns:
+            the actual status of the pipeline job.
+
+        Raises:
+            AssertionError: If the run was not executed by to this orchestrator.
+            ValueError: If it fetches an unknown state or if we can not fetch
+                the orchestrator run ID.
+        """
+        # Make sure that the stack exists and is accessible
+        if run.stack is None:
+            raise ValueError(
+                "The stack that the run was executed on is not available "
+                "anymore."
+            )
+
+        # Make sure that the run belongs to this orchestrator
+        assert (
+            self.id
+            == run.stack.components[StackComponentType.ORCHESTRATOR][0].id
+        )
+
+        # Initialize the Vertex client
+        credentials, project_id = self._get_authentication()
+        aiplatform.init(
+            project=project_id,
+            location=self.config.location,
+            credentials=credentials,
+        )
+
+        # Fetch the status of the PipelineJob
+        if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
+        elif run.orchestrator_run_id is not None:
+            run_id = run.orchestrator_run_id
+        else:
+            raise ValueError(
+                "Can not find the orchestrator run ID, thus can not fetch "
+                "the status."
+            )
+        status = aiplatform.PipelineJob.get(run_id).state
+
+        # Map the potential outputs to ZenML ExecutionStatus. Potential values:
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/describe_pipeline_execution.html#
+        if status in [PipelineState.PIPELINE_STATE_UNSPECIFIED]:
+            return run.status
+        elif status in [
+            PipelineState.PIPELINE_STATE_QUEUED,
+            PipelineState.PIPELINE_STATE_PENDING,
+        ]:
+            return ExecutionStatus.INITIALIZING
+        elif status in [
+            PipelineState.PIPELINE_STATE_RUNNING,
+            PipelineState.PIPELINE_STATE_PAUSED,
+        ]:
+            return ExecutionStatus.RUNNING
+        elif status in [PipelineState.PIPELINE_STATE_SUCCEEDED]:
+            return ExecutionStatus.COMPLETED
+
+        elif status in [
+            PipelineState.PIPELINE_STATE_FAILED,
+            PipelineState.PIPELINE_STATE_CANCELLING,
+            PipelineState.PIPELINE_STATE_CANCELLED,
+        ]:
+            return ExecutionStatus.FAILED
+        else:
+            raise ValueError("Unknown status for the pipeline job.")
+
+    def compute_metadata(
+        self, job: aiplatform.PipelineJob
+    ) -> Iterator[Dict[str, MetadataType]]:
+        """Generate run metadata based on the corresponding Vertex PipelineJob.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Yields:
+            A dictionary of metadata related to the pipeline run.
+        """
+        metadata: Dict[str, MetadataType] = {}
+
+        # Orchestrator Run ID
+        if run_id := self._compute_orchestrator_run_id(job):
+            metadata[METADATA_ORCHESTRATOR_RUN_ID] = run_id
+
+        # URL to the Vertex's pipeline view
+        if orchestrator_url := self._compute_orchestrator_url(job):
+            metadata[METADATA_ORCHESTRATOR_URL] = Uri(orchestrator_url)
+
+        # URL to the corresponding Logs Explorer page
+        if logs_url := self._compute_orchestrator_logs_url(job):
+            metadata[METADATA_ORCHESTRATOR_LOGS_URL] = Uri(logs_url)
+
+        yield metadata
+
+    @staticmethod
+    def _compute_orchestrator_url(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Generate the Orchestrator Dashboard URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+             the URL to the dashboard view in Vertex.
+        """
+        try:
+            return str(job._dashboard_uri())
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline url: {e}"
+            )
+            return None
+
+    @staticmethod
+    def _compute_orchestrator_logs_url(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Generate the Logs Explorer URL upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+            the URL querying the pipeline logs in Logs Explorer on GCP.
+        """
+        try:
+            base_url = "https://console.cloud.google.com/logs/query"
+            query = f"""
+             resource.type="aiplatform.googleapis.com/PipelineJob"
+             resource.labels.pipeline_job_id="{job.job_id}"
+             """
+            encoded_query = urllib.parse.quote(query)
+            return f"{base_url}?project={job.project}&query={encoded_query}"
+
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the logs url: {e}"
+            )
+            return None
+
+    @staticmethod
+    def _compute_orchestrator_run_id(
+        job: aiplatform.PipelineJob,
+    ) -> Optional[str]:
+        """Fetch the Orchestrator Run ID upon pipeline execution.
+
+        Args:
+            job: The corresponding PipelineJob object.
+
+        Returns:
+            the Execution ID of the run in Vertex.
+        """
+        try:
+            if job.job_id:
+                return str(job.job_id)
+
+            return None
+        except Exception as e:
+            logger.warning(
+                f"There was an issue while extracting the pipeline run ID: {e}"
+            )
+            return None