zenml-nightly 0.75.0.dev20250313__py3-none-any.whl → 0.75.0.dev20250315__py3-none-any.whl

zenml/config/global_config.py

@@ -48,7 +48,7 @@ from zenml.logger import get_logger
 from zenml.utils import io_utils, yaml_utils
 
 if TYPE_CHECKING:
-    from zenml.models import StackResponse, WorkspaceResponse
+    from zenml.models import ProjectResponse, StackResponse
     from zenml.zen_stores.base_zen_store import BaseZenStore
 
 logger = get_logger(__name__)
@@ -113,7 +113,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
             global config.
         store: Store configuration.
         active_stack_id: The ID of the active stack.
-        active_workspace_name: The name of the active workspace.
+        active_project_name: The name of the active project.
     """
 
     user_id: uuid.UUID = Field(default_factory=uuid.uuid4)
@@ -123,10 +123,10 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
     version: Optional[str] = None
     store: Optional[SerializeAsAny[StoreConfiguration]] = None
     active_stack_id: Optional[uuid.UUID] = None
-    active_workspace_name: Optional[str] = None
+    active_project_name: Optional[str] = None
 
     _zen_store: Optional["BaseZenStore"] = None
-    _active_workspace: Optional["WorkspaceResponse"] = None
+    _active_project: Optional["ProjectResponse"] = None
     _active_stack: Optional["StackResponse"] = None
 
     def __init__(self, **data: Any) -> None:
@@ -385,24 +385,24 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
     def _sanitize_config(self) -> None:
         """Sanitize and save the global configuration.
 
-        This method is called to ensure that the active stack and workspace
+        This method is called to ensure that the active stack and project
         are set to their default values, if possible.
         """
         # If running in a ZenML server environment, the active stack and
-        # workspace are not relevant
+        # project are not relevant
         if ENV_ZENML_SERVER in os.environ:
             return
-        active_workspace, active_stack = self.zen_store.validate_active_config(
-            self.active_workspace_name,
+        active_project, active_stack = self.zen_store.validate_active_config(
+            self.active_project_name,
             self.active_stack_id,
             config_name="global",
         )
-        if active_workspace:
-            self.active_workspace_name = active_workspace.name
-            self._active_workspace = active_workspace
+        if active_project:
+            self.active_project_name = active_project.name
+            self._active_project = active_project
         else:
-            self.active_workspace_name = None
-            self._active_workspace = None
+            self.active_project_name = None
+            self._active_project = None
 
         self.set_active_stack(active_stack)
 
@@ -719,22 +719,22 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
 
         return self._zen_store
 
-    def set_active_workspace(
-        self, workspace: "WorkspaceResponse"
-    ) -> "WorkspaceResponse":
-        """Set the workspace for the local client.
+    def set_active_project(
+        self, project: "ProjectResponse"
+    ) -> "ProjectResponse":
+        """Set the project for the local client.
 
         Args:
-            workspace: The workspace to set active.
+            project: The project to set active.
 
         Returns:
-            The workspace that was set active.
+            The project that was set active.
         """
-        self.active_workspace_name = workspace.name
-        self._active_workspace = workspace
-        # Sanitize the global configuration to reflect the new workspace
+        self.active_project_name = project.name
+        self._active_project = project
+        # Sanitize the global configuration to reflect the new project
         self._sanitize_config()
-        return workspace
+        return project
 
     def set_active_stack(self, stack: "StackResponse") -> None:
         """Set the active stack for the local client.
@@ -745,43 +745,41 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         self.active_stack_id = stack.id
         self._active_stack = stack
 
-    def get_active_workspace(self) -> "WorkspaceResponse":
-        """Get a model of the active workspace for the local client.
+    def get_active_project(self) -> "ProjectResponse":
+        """Get a model of the active project for the local client.
 
         Returns:
-            The model of the active workspace.
+            The model of the active project.
         """
-        workspace_name = self.get_active_workspace_name()
+        project_name = self.get_active_project_name()
 
-        if self._active_workspace is not None:
-            return self._active_workspace
+        if self._active_project is not None:
+            return self._active_project
 
-        workspace = self.zen_store.get_workspace(
-            workspace_name_or_id=workspace_name,
+        project = self.zen_store.get_project(
+            project_name_or_id=project_name,
         )
-        return self.set_active_workspace(workspace)
+        return self.set_active_project(project)
 
-    def get_active_workspace_name(self) -> str:
-        """Get the name of the active workspace.
-
-        If the active workspace doesn't exist yet, the ZenStore is reinitialized.
+    def get_active_project_name(self) -> str:
+        """Get the name of the active project.
 
         Returns:
-            The name of the active workspace.
+            The name of the active project.
 
         Raises:
-            RuntimeError: If the active workspace is not set.
+            RuntimeError: If the active project is not set.
         """
-        if self.active_workspace_name is None:
+        if self.active_project_name is None:
             _ = self.zen_store
-            if self.active_workspace_name is None:
+            if self.active_project_name is None:
                 raise RuntimeError(
-                    "No workspace is currently set as active. Please set the "
-                    "active workspace using the `zenml workspace set` CLI "
+                    "No project is currently set as active. Please set the "
+                    "active project using the `zenml project set` CLI "
                     "command."
                 )
 
-        return self.active_workspace_name
+        return self.active_project_name
 
     def get_active_stack_id(self) -> UUID:
         """Get the ID of the active stack.

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/constants.py

@@ -143,13 +143,13 @@ ENV_ZENML_REPOSITORY_PATH = "ZENML_REPOSITORY_PATH"
 ENV_ZENML_PREVENT_PIPELINE_EXECUTION = "ZENML_PREVENT_PIPELINE_EXECUTION"
 ENV_ZENML_ENABLE_RICH_TRACEBACK = "ZENML_ENABLE_RICH_TRACEBACK"
 ENV_ZENML_ACTIVE_STACK_ID = "ZENML_ACTIVE_STACK_ID"
-ENV_ZENML_ACTIVE_WORKSPACE_ID = "ZENML_ACTIVE_WORKSPACE_ID"
+ENV_ZENML_ACTIVE_PROJECT_ID = "ZENML_ACTIVE_PROJECT_ID"
 ENV_ZENML_SUPPRESS_LOGS = "ZENML_SUPPRESS_LOGS"
 ENV_ZENML_ENABLE_REPO_INIT_WARNINGS = "ZENML_ENABLE_REPO_INIT_WARNINGS"
 ENV_ZENML_SECRET_VALIDATION_LEVEL = "ZENML_SECRET_VALIDATION_LEVEL"
 ENV_ZENML_DEFAULT_USER_NAME = "ZENML_DEFAULT_USER_NAME"
 ENV_ZENML_DEFAULT_USER_PASSWORD = "ZENML_DEFAULT_USER_PASSWORD"
-ENV_ZENML_DEFAULT_WORKSPACE_NAME = "ZENML_DEFAULT_WORKSPACE_NAME"
+ENV_ZENML_DEFAULT_PROJECT_NAME = "ZENML_DEFAULT_PROJECT_NAME"
 ENV_ZENML_STORE_PREFIX = "ZENML_STORE_"
 ENV_ZENML_SECRETS_STORE_PREFIX = "ZENML_SECRETS_STORE_"
 ENV_ZENML_BACKUP_SECRETS_STORE_PREFIX = "ZENML_BACKUP_SECRETS_STORE_"
@@ -239,7 +239,7 @@ SQL_STORE_BACKUP_DIRECTORY_NAME = "database_backup"
 
 DEFAULT_USERNAME = "default"
 DEFAULT_PASSWORD = ""
-DEFAULT_WORKSPACE_NAME = "default"
+DEFAULT_PROJECT_NAME = "default"
 DEFAULT_STACK_AND_COMPONENT_NAME = "default"
 
 # Rich config
@@ -409,6 +409,7 @@ VERSION_1 = "/v1"
 VISUALIZE = "/visualize"
 WEBHOOKS = "/webhooks"
 WORKSPACES = "/workspaces"
+PROJECTS = "/projects"
 
 # model metadata yaml file name
 MODEL_METADATA_YAML_FILE_NAME = "model_metadata.yaml"
@@ -450,6 +451,7 @@ STACK_RECIPES_GITHUB_REPO = "https://github.com/zenml-io/mlops-stacks.git"
 
 # Parameters for internal ZenML Models
 TEXT_FIELD_MAX_LENGTH = 65535
+STR_ID_FIELD_MAX_LENGTH = 50
 STR_FIELD_MAX_LENGTH = 255
 MEDIUMTEXT_MAX_LENGTH = 2**24 - 1
 

zenml/event_hub/event_hub.py

@@ -126,7 +126,7 @@ class InternalEventHub(BaseEventHub):
         triggers: List[TriggerResponse] = depaginate(
             self.zen_store.list_triggers,
             trigger_filter_model=TriggerFilter(
-                workspace=event_source.workspace.id,
+                project=event_source.project.id,
                 event_source_id=event_source.id,
                 is_active=True,
             ),

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/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/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/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
         )