zenml-nightly 0.75.0.dev20250314__tar.gz → 0.75.0.dev20250315__tar.gz

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.75.0.dev20250314
+Version: 0.75.0.dev20250315
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.75.0.dev20250314"
+version = "0.75.0.dev20250315"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]
@@ -449,6 +449,7 @@ module = [
     "sagemaker.*",
     "azureml.*",
     "google.*",
+    "google_cloud_pipeline_components.*",
     "neuralprophet.*",
     "lightgbm.*",
     "scipy.*",

zenml_nightly-0.75.0.dev20250315/src/zenml/VERSION

@@ -0,0 +1 @@
+0.75.0.dev20250315

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/cli/login.py

@@ -246,7 +246,7 @@ def connect_to_pro_server(
     """
     from zenml.login.credentials_store import get_credentials_store
     from zenml.login.pro.client import ZenMLProClient
-    from zenml.login.pro.tenant.models import TenantStatus
+    from zenml.login.pro.workspace.models import WorkspaceStatus
 
     pro_api_url = pro_api_url or ZENML_PRO_API_URL
     pro_api_url = pro_api_url.rstrip("/")
@@ -295,7 +295,7 @@ def connect_to_pro_server(
         # We also need to remove all existing API tokens associated with the
         # target ZenML Pro API, otherwise they will continue to be used after
         # the re-login flow.
-        credentials_store.clear_all_pro_tokens(pro_api_url)
+        credentials_store.clear_all_pro_tokens()
         try:
             token = web_login(
                 pro_api_url=pro_api_url,
@@ -310,13 +310,14 @@ def connect_to_pro_server(
             "your session expires."
         )
 
-        tenant_id: Optional[str] = None
+        workspace_id: Optional[str] = None
         if token.device_metadata:
-            tenant_id = token.device_metadata.get("tenant_id")
+            # TODO: is this still correct?
+            workspace_id = token.device_metadata.get("tenant_id")
 
-        if tenant_id is None and pro_server is None:
+        if workspace_id is None and pro_server is None:
             # This is not really supposed to happen, because the implementation
-            # of the web login workflow should always return a tenant ID, but
+            # of the web login workflow should always return a workspace ID, but
             # we're handling it just in case.
             cli_utils.declare(
                 "A valid server was not selected during the login process. "
@@ -328,14 +329,14 @@ def connect_to_pro_server(
 
         # The server selected during the web login process overrides any
         # server argument passed to the command.
-        server_id = UUID(tenant_id)
+        server_id = UUID(workspace_id)
 
     client = ZenMLProClient(pro_api_url)
 
     if server_id:
-        server = client.tenant.get(server_id)
+        server = client.workspace.get(server_id)
     elif server_url:
-        servers = client.tenant.list(url=server_url, member_only=True)
+        servers = client.workspace.list(url=server_url, member_only=True)
         if not servers:
             raise AuthorizationException(
                 f"The '{server_url}' URL belongs to a ZenML Pro server, "
@@ -345,7 +346,9 @@ def connect_to_pro_server(
 
         server = servers[0]
     elif server_name:
-        servers = client.tenant.list(tenant_name=server_name, member_only=True)
+        servers = client.workspace.list(
+            workspace_name=server_name, member_only=True
+        )
         if not servers:
             raise AuthorizationException(
                 f"No ZenML Pro server with the name '{server_name}' exists "
@@ -361,15 +364,15 @@ def connect_to_pro_server(
 
     server_id = server.id
 
-    if server.status == TenantStatus.PENDING:
+    if server.status == WorkspaceStatus.PENDING:
         with console.status(
             f"Waiting for your `{server.name}` ZenML Pro server to be set up..."
         ):
             timeout = 180  # 3 minutes
             while True:
                 time.sleep(5)
-                server = client.tenant.get(server_id)
-                if server.status != TenantStatus.PENDING:
+                server = client.workspace.get(server_id)
+                if server.status != WorkspaceStatus.PENDING:
                     break
                 timeout -= 5
                 if timeout <= 0:
@@ -380,7 +383,7 @@ def connect_to_pro_server(
                         f"ZenML Pro dashboard at {server.dashboard_url}."
                     )
 
-    if server.status == TenantStatus.FAILED:
+    if server.status == WorkspaceStatus.FAILED:
         cli_utils.error(
             f"Your `{server.name}` ZenML Pro server is currently in a "
             "failed state. Please manage the server state by visiting the "
@@ -388,7 +391,7 @@ def connect_to_pro_server(
             "your server administrator."
         )
 
-    elif server.status == TenantStatus.DEACTIVATED:
+    elif server.status == WorkspaceStatus.DEACTIVATED:
         cli_utils.error(
             f"Your `{server.name}` ZenML Pro server is currently "
             "deactivated. Please manage the server state by visiting the "
@@ -396,7 +399,7 @@ def connect_to_pro_server(
             "your server administrator."
         )
 
-    elif server.status == TenantStatus.AVAILABLE:
+    elif server.status == WorkspaceStatus.AVAILABLE:
         if not server.url:
             cli_utils.error(
                 f"The ZenML Pro server '{server.name}' is not currently "
@@ -418,7 +421,7 @@ def connect_to_pro_server(
     connect_to_server(server.url, api_key=api_key, pro_server=True)
 
     # Update the stored server info with more accurate data taken from the
-    # ZenML Pro tenant object.
+    # ZenML Pro workspace object.
     credentials_store.update_server_info(server.url, server)
 
     cli_utils.declare(f"Connected to ZenML Pro server: {server.name}.")
@@ -836,7 +839,7 @@ def login(
                 pro_api_url=pro_api_url,
             )
 
-    elif current_non_local_server:
+    elif current_non_local_server and not refresh:
         # The server argument is not provided, so we default to
         # re-authenticating to the current non-local server that the client is
         # connected to.

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/cli/server.py

@@ -220,7 +220,7 @@ def status() -> None:
                 )
                 if pro_credentials:
                     pro_client = ZenMLProClient(pro_credentials.url)
-                    pro_servers = pro_client.tenant.list(
+                    pro_servers = pro_client.workspace.list(
                         url=store_cfg.url, member_only=True
                     )
                     if pro_servers:
@@ -575,7 +575,7 @@ def server_list(
     from zenml.login.credentials_store import get_credentials_store
     from zenml.login.pro.client import ZenMLProClient
     from zenml.login.pro.constants import ZENML_PRO_API_URL
-    from zenml.login.pro.tenant.models import TenantRead, TenantStatus
+    from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
 
     pro_api_url = pro_api_url or ZENML_PRO_API_URL
     pro_api_url = pro_api_url.rstrip("/")
@@ -601,10 +601,10 @@ def server_list(
         # that the user has never connected to (and are therefore not stored in
         # the credentials store).
 
-        accessible_pro_servers: List[TenantRead] = []
+        accessible_pro_servers: List[WorkspaceRead] = []
         try:
             client = ZenMLProClient(pro_api_url)
-            accessible_pro_servers = client.tenant.list(member_only=not all)
+            accessible_pro_servers = client.workspace.list(member_only=not all)
         except AuthorizationException as e:
             cli_utils.warning(f"ZenML Pro authorization error: {e}")
 
@@ -638,7 +638,7 @@ def server_list(
             accessible_pro_servers = [
                 s
                 for s in accessible_pro_servers
-                if s.status == TenantStatus.AVAILABLE
+                if s.status == WorkspaceStatus.AVAILABLE
             ]
 
         if not accessible_pro_servers:

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/client.py

@@ -68,6 +68,7 @@ from zenml.enums import (
     SorterOps,
     StackComponentType,
     StoreType,
+    TaggableResourceTypes,
 )
 from zenml.exceptions import (
     AuthorizationException,
@@ -7770,6 +7771,7 @@ class Client(metaclass=ClientMetaClass):
         name: Optional[str] = None,
         color: Optional[Union[str, ColorVariants]] = None,
         exclusive: Optional[bool] = None,
+        resource_type: Optional[Union[str, TaggableResourceTypes]] = None,
         hydrate: bool = False,
     ) -> Page[TagResponse]:
         """Get tags by filter.
@@ -7777,7 +7779,7 @@ class Client(metaclass=ClientMetaClass):
         Args:
             sort_by: The column to sort by.
             page: The page of items.
-            size: The maximum size of all pages.
+            size: The maximum size of all pages
             logical_operator: Which logical operator to use [and, or].
             id: Use the id of stacks to filter by.
             user: Use the user to filter by.
@@ -7786,6 +7788,7 @@ class Client(metaclass=ClientMetaClass):
             name: The name of the tag.
             color: The color of the tag.
             exclusive: Flag indicating whether the tag is exclusive.
+            resource_type: Filter tags associated with a specific resource type.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -7805,6 +7808,7 @@ class Client(metaclass=ClientMetaClass):
                 name=name,
                 color=color,
                 exclusive=exclusive,
+                resource_type=resource_type,
             ),
             hydrate=hydrate,
         )

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/config/server_config.py

@@ -592,23 +592,23 @@ class ServerConfiguration(BaseModel):
             server_config.external_user_info_url = (
                 f"{server_pro_config.api_url}/users/authorize_server"
             )
-            server_config.external_server_id = server_pro_config.tenant_id
+            server_config.external_server_id = server_pro_config.workspace_id
             server_config.rbac_implementation_source = (
                 "zenml.zen_server.rbac.zenml_cloud_rbac.ZenMLCloudRBAC"
             )
             server_config.feature_gate_implementation_source = "zenml.zen_server.feature_gate.zenml_cloud_feature_gate.ZenMLCloudFeatureGateInterface"
             server_config.reportable_resources = DEFAULT_REPORTABLE_RESOURCES
-            server_config.dashboard_url = f"{server_pro_config.dashboard_url}/organizations/{server_pro_config.organization_id}/tenants/{server_pro_config.tenant_id}"
+            server_config.dashboard_url = f"{server_pro_config.dashboard_url}/workspaces/{server_pro_config.workspace_id}"
             server_config.metadata.update(
                 dict(
                     account_id=str(server_pro_config.organization_id),
                     organization_id=str(server_pro_config.organization_id),
-                    tenant_id=str(server_pro_config.tenant_id),
+                    workspace_id=str(server_pro_config.workspace_id),
                 )
             )
-            if server_pro_config.tenant_name:
+            if server_pro_config.workspace_name:
                 server_config.metadata.update(
-                    dict(tenant_name=server_pro_config.tenant_name)
+                    dict(workspace_name=server_pro_config.workspace_name)
                 )
 
             extra_cors_allow_origins = [
@@ -660,8 +660,8 @@ class ServerProConfiguration(BaseModel):
         oauth2_audience: The OAuth2 audience.
         organization_id: The ZenML Pro organization ID.
         organization_name: The ZenML Pro organization name.
-        tenant_id: The ZenML Pro tenant ID.
-        tenant_name: The ZenML Pro tenant name.
+        workspace_id: The ZenML Pro workspace ID.
+        workspace_name: The ZenML Pro workspace name.
     """
 
     api_url: str
@@ -670,8 +670,8 @@ class ServerProConfiguration(BaseModel):
     oauth2_audience: str
     organization_id: UUID
     organization_name: Optional[str] = None
-    tenant_id: UUID
-    tenant_name: Optional[str] = None
+    workspace_id: UUID
+    workspace_name: Optional[str] = None
 
     @field_validator("api_url", "dashboard_url")
     @classmethod

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/integrations/gcp/__init__.py

@@ -52,6 +52,7 @@ class GcpIntegration(Integration):
         "google-cloud-storage>=2.9.0",
         "google-cloud-aiplatform>=1.34.0",  # includes shapely pin fix
         "google-cloud-build>=3.11.0",
+        "google-cloud-pipeline-components>=2.19.0",
         "kubernetes",
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes","kfp"]

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py

@@ -23,6 +23,9 @@ from zenml.integrations.gcp import (
 from zenml.integrations.gcp.google_credentials_mixin import (
     GoogleCredentialsConfigMixin,
 )
+from zenml.integrations.gcp.vertex_custom_job_parameters import (
+    VertexCustomJobParameters,
+)
 from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.models import ServiceConnectorRequirements
 from zenml.orchestrators import BaseOrchestratorConfig, BaseOrchestratorFlavor
@@ -61,6 +64,8 @@ class VertexOrchestratorSettings(BaseSettings):
     node_selector_constraint: Optional[Tuple[str, str]] = None
     pod_settings: Optional[KubernetesPodSettings] = None
 
+    custom_job_parameters: Optional[VertexCustomJobParameters] = None
+
     _node_selector_deprecation = (
         deprecation_utils.deprecate_pydantic_attributes(
             "node_selector_constraint"

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/integrations/gcp/flavors/vertex_step_operator_flavor.py

@@ -23,6 +23,9 @@ from zenml.integrations.gcp import (
 from zenml.integrations.gcp.google_credentials_mixin import (
     GoogleCredentialsConfigMixin,
 )
+from zenml.integrations.gcp.vertex_custom_job_parameters import (
+    VertexCustomJobParameters,
+)
 from zenml.models import ServiceConnectorRequirements
 from zenml.step_operators.base_step_operator import (
     BaseStepOperatorConfig,
@@ -33,34 +36,8 @@ if TYPE_CHECKING:
     from zenml.integrations.gcp.step_operators import VertexStepOperator
 
 
-class VertexStepOperatorSettings(BaseSettings):
-    """Settings for the Vertex step operator.
-
-    Attributes:
-        accelerator_type: Defines which accelerator (GPU, TPU) is used for the
-            job. Check out out this table to see which accelerator
-            type and count are compatible with your chosen machine type:
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
-        accelerator_count: Defines number of accelerators to be used for the
-            job. Check out out this table to see which accelerator
-            type and count are compatible with your chosen machine type:
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#gpu-compatibility-table.
-        machine_type: Machine type specified here
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types.
-        boot_disk_size_gb: Size of the boot disk in GB. (Default: 100)
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
-        boot_disk_type: Type of the boot disk. (Default: pd-ssd)
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
-        persistent_resource_id: The ID of the persistent resource to use for the job.
-            https://cloud.google.com/vertex-ai/docs/training/persistent-resource-overview
-    """
-
-    accelerator_type: Optional[str] = None
-    accelerator_count: int = 0
-    machine_type: str = "n1-standard-4"
-    boot_disk_size_gb: int = 100
-    boot_disk_type: str = "pd-ssd"
-    persistent_resource_id: Optional[str] = None
+class VertexStepOperatorSettings(VertexCustomJobParameters, BaseSettings):
+    """Settings for the Vertex step operator."""
 
 
 class VertexStepOperatorConfig(

{zenml_nightly-0.75.0.dev20250314 → zenml_nightly-0.75.0.dev20250315}/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -49,8 +49,12 @@ from uuid import UUID
 from google.api_core import exceptions as google_exceptions
 from google.cloud import aiplatform
 from google.cloud.aiplatform_v1.types import PipelineState
+from google_cloud_pipeline_components.v1.custom_job.utils import (
+    create_custom_training_job_from_component,
+)
 from kfp import dsl
 from kfp.compiler import Compiler
+from kfp.dsl.base_component import BaseComponent
 
 from zenml.config.resource_settings import ResourceSettings
 from zenml.constants import (
@@ -71,13 +75,15 @@ from zenml.integrations.gcp.flavors.vertex_orchestrator_flavor import (
 from zenml.integrations.gcp.google_credentials_mixin import (
     GoogleCredentialsMixin,
 )
+from zenml.integrations.gcp.vertex_custom_job_parameters import (
+    VertexCustomJobParameters,
+)
 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:
@@ -263,14 +269,14 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     "schedule to a Vertex orchestrator."
                 )
 
-    def _create_dynamic_component(
+    def _create_container_component(
         self,
         image: str,
         command: List[str],
         arguments: List[str],
         component_name: str,
-    ) -> dsl.PipelineTask:
-        """Creates a dynamic container component for a Vertex pipeline.
+    ) -> BaseComponent:
+        """Creates a container component for a Vertex pipeline.
 
         Args:
             image: The image to use for the component.
@@ -279,7 +285,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             component_name: The name of the component.
 
         Returns:
-            The dynamic container component.
+            The container component.
         """
 
         def dynamic_container_component() -> dsl.ContainerSpec:
@@ -294,7 +300,6 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 args=arguments,
             )
 
-        # Change the name of the function
         new_container_spec_func = types.FunctionType(
             dynamic_container_component.__code__,
             dynamic_container_component.__globals__,
@@ -303,9 +308,50 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             closure=dynamic_container_component.__closure__,
         )
         pipeline_task = dsl.container_component(new_container_spec_func)
-
         return pipeline_task
 
+    def _convert_to_custom_training_job(
+        self,
+        component: BaseComponent,
+        settings: VertexOrchestratorSettings,
+        environment: Dict[str, str],
+    ) -> BaseComponent:
+        """Convert a component to a custom training job component.
+
+        Args:
+            component: The component to convert.
+            settings: The settings for the custom training job.
+            environment: The environment variables to set in the custom
+                training job.
+
+        Returns:
+            The custom training job component.
+        """
+        custom_job_parameters = (
+            settings.custom_job_parameters or VertexCustomJobParameters()
+        )
+        if (
+            custom_job_parameters.persistent_resource_id
+            and not custom_job_parameters.service_account
+        ):
+            # Persistent resources require an explicit service account, but
+            # none was provided in the custom job parameters. We try to fall
+            # back to the workload service account.
+            custom_job_parameters.service_account = (
+                self.config.workload_service_account
+            )
+
+        custom_job_component = create_custom_training_job_from_component(
+            component_spec=component,
+            env=[
+                {"name": key, "value": value}
+                for key, value in environment.items()
+            ],
+            **custom_job_parameters.model_dump(),
+        )
+
+        return custom_job_component
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -383,7 +429,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             Returns:
                 pipeline_func
             """
-            step_name_to_dynamic_component: Dict[str, Any] = {}
+            step_name_to_dynamic_component: Dict[str, BaseComponent] = {}
 
             for step_name, step in deployment.step_configurations.items():
                 image = self.get_image(
@@ -397,7 +443,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                         deployment_id=deployment.id,
                     )
                 )
-                dynamic_component = self._create_dynamic_component(
+                component = self._create_container_component(
                     image, command, arguments, step_name
                 )
                 step_settings = cast(
@@ -442,7 +488,11 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                                 key,
                             )
 
-                step_name_to_dynamic_component[step_name] = dynamic_component
+                step_name_to_dynamic_component[step_name] = component
+
+            environment[ENV_ZENML_VERTEX_RUN_ID] = (
+                dsl.PIPELINE_JOB_NAME_PLACEHOLDER
+            )
 
             @dsl.pipeline(  # type: ignore[misc]
                 display_name=orchestrator_run_name,
@@ -462,81 +512,81 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                         step_name_to_dynamic_component[upstream_step_name]
                         for upstream_step_name in step.spec.upstream_steps
                     ]
-                    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)
-                    )
 
                     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
-                            ],
+
+                    use_custom_training_job = (
+                        step_settings.custom_job_parameters is not None
+                    )
+
+                    if use_custom_training_job:
+                        if not step.config.resource_settings.empty:
+                            logger.warning(
+                                "Ignoring resource settings because "
+                                "the step is running as a custom training job. "
+                                "Use `custom_job_parameters.machine_type` "
+                                "to configure the machine type instead."
+                            )
+                        if step_settings.node_selector_constraint:
+                            logger.warning(
+                                "Ignoring node selector constraint because "
+                                "the step is running as a custom training job. "
+                                "Use `custom_job_parameters.accelerator_type` "
+                                "to configure the accelerator type instead."
+                            )
+                        component = self._convert_to_custom_training_job(
+                            component,
+                            settings=step_settings,
+                            environment=environment,
+                        )
+                        task = (
+                            component()
+                            .set_display_name(name=component_name)
+                            .set_caching_options(enable_caching=False)
+                            .after(*upstream_step_components)
                         )
-                    elif step_settings.node_selector_constraint:
-                        node_selector_constraint = (
-                            GKE_ACCELERATOR_NODE_SELECTOR_CONSTRAINT_LABEL,
-                            step_settings.node_selector_constraint[1],
+                    else:
+                        task = (
+                            component()
+                            .set_display_name(
+                                name=component_name,
+                            )
+                            .set_caching_options(enable_caching=False)
+                            .after(*upstream_step_components)
                         )
+                        for key, value in environment.items():
+                            task = task.set_env_variable(name=key, value=value)
 
-                    self._configure_container_resources(
-                        dynamic_component=task,
-                        resource_settings=step.config.resource_settings,
-                        node_selector_constraint=node_selector_constraint,
-                    )
+                        pod_settings = step_settings.pod_settings
 
-            return dynamic_pipeline
+                        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],
+                            )
 
-        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.
+                        self._configure_container_resources(
+                            dynamic_component=task,
+                            resource_settings=step.config.resource_settings,
+                            node_selector_constraint=node_selector_constraint,
+                        )
 
-            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}"
-            )
+            return dynamic_pipeline
 
         # Save the generated pipeline to a file.
         fileio.makedirs(self.pipeline_directory)
@@ -556,9 +606,6 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             ),
         )
 
-        # 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
         )