WIP step env on containers

Author: schustmi

src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -195,7 +195,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Creates and writes an Airflow DAG zip file.
@@ -248,7 +248,7 @@ def prepare_or_run_pipeline(
                 docker_image=image,
                 command=command,
                 arguments=arguments,
-                environment=environment,
+                environment=environment[step_name],
                 operator_source=settings.operator,
                 operator_args=operator_args,
             )

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

@@ -266,7 +266,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Iterator[Dict[str, MetadataType]]:
         """Prepares or runs a pipeline on Sagemaker.
@@ -299,23 +299,24 @@ def prepare_or_run_pipeline(
 
         session = self._get_sagemaker_session()
 
-        # Sagemaker does not allow environment variables longer than 256
-        # characters to be passed to Processor steps. If an environment variable
-        # is longer than 256 characters, we split it into multiple environment
-        # variables (chunks) and re-construct it on the other side using the
-        # custom entrypoint configuration.
-        split_environment_variables(
-            size_limit=SAGEMAKER_PROCESSOR_STEP_ENV_VAR_SIZE_LIMIT,
-            env=environment,
-        )
-
-        environment[ENV_ZENML_SAGEMAKER_RUN_ID] = (
-            ExecutionVariables.PIPELINE_EXECUTION_ARN
-        )
-
         sagemaker_steps = []
         for step_name, step in deployment.step_configurations.items():
-            step_environment = environment.copy()
+            step_environment = environment[step_name]
+
+            # Sagemaker does not allow environment variables longer than 256
+            # characters to be passed to Processor steps. If an environment variable
+            # is longer than 256 characters, we split it into multiple environment
+            # variables (chunks) and re-construct it on the other side using the
+            # custom entrypoint configuration.
+            split_environment_variables(
+                size_limit=SAGEMAKER_PROCESSOR_STEP_ENV_VAR_SIZE_LIMIT,
+                env=step_environment,
+            )
+
+            step_environment[ENV_ZENML_SAGEMAKER_RUN_ID] = (
+                ExecutionVariables.PIPELINE_EXECUTION_ARN
+            )
+
             image = self.get_image(deployment=deployment, step_name=step_name)
             command = SagemakerEntrypointConfiguration.get_entrypoint_command()
             arguments = (

src/zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -202,7 +202,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Iterator[Dict[str, MetadataType]]:
         """Prepares or runs a pipeline on AzureML.
@@ -243,6 +243,7 @@ def prepare_or_run_pipeline(
         # Create components
         components = {}
         for step_name, step in deployment.step_configurations.items():
+            step_environment = environment[step_name]
             # Get the image for each step
             image = self.get_image(deployment=deployment, step_name=step_name)
 
@@ -252,7 +253,9 @@ def prepare_or_run_pipeline(
                 AzureMLEntrypointConfiguration.get_entrypoint_arguments(
                     step_name=step_name,
                     deployment_id=deployment.id,
-                    zenml_env_variables=b64_encode(json.dumps(environment)),
+                    zenml_env_variables=b64_encode(
+                        json.dumps(step_environment)
+                    ),
                 )
             )
 

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

@@ -558,6 +558,8 @@ def dynamic_pipeline() -> None:
                         step_settings.custom_job_parameters is not None
                     )
 
+                    step_environment = environment[component_name]
+
                     if use_custom_training_job:
                         if not step.config.resource_settings.empty:
                             logger.warning(
@@ -576,7 +578,7 @@ def dynamic_pipeline() -> None:
                         component = self._convert_to_custom_training_job(
                             component,
                             settings=step_settings,
-                            environment=environment,
+                            environment=step_environment,
                         )
                         task = (
                             component()
@@ -593,7 +595,7 @@ def dynamic_pipeline() -> None:
                             .set_caching_options(enable_caching=False)
                             .after(*upstream_step_components)
                         )
-                        for key, value in environment.items():
+                        for key, value in step_environment.items():
                             task = task.set_env_variable(name=key, value=value)
 
                         pod_settings = step_settings.pod_settings

src/zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -72,7 +72,7 @@
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
-from zenml.utils import io_utils, settings_utils, yaml_utils
+from zenml.utils import io_utils, settings_utils
 
 if TYPE_CHECKING:
     from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
@@ -470,7 +470,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Creates a kfp yaml file.
@@ -590,6 +590,7 @@ def dynamic_pipeline() -> None:
                     component_name,
                     component,
                 ) in step_name_to_dynamic_component.items():
+                    step_environment = environment[component_name]
                     # for each component, check to see what other steps are
                     # upstream of it
                     step = deployment.step_configurations[component_name]
@@ -609,6 +610,9 @@ def dynamic_pipeline() -> None:
                         )
                         .after(*upstream_step_components)
                     )
+                    for key, value in step_environment.items():
+                        task = task.set_env_variable(name=key, value=value)
+
                     self._configure_container_resources(
                         task,
                         step.config.resource_settings,
@@ -617,39 +621,6 @@ def dynamic_pipeline() -> None:
 
             return dynamic_pipeline
 
-        def _update_yaml_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_yaml(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_yaml(pipeline_file_path, pipeline_definition)
-
-            print(
-                f"Updated YAML file with environment variables at {yaml_file_path}"
-            )
-
         # Get a filepath to use to save the finished yaml to
         fileio.makedirs(self.pipeline_directory)
         pipeline_file_path = os.path.join(
@@ -663,9 +634,6 @@ def _update_yaml_with_environment(
             pipeline_name=orchestrator_run_name,
         )
 
-        # Let's update the YAML file with the environment variables
-        _update_yaml_with_environment(pipeline_file_path, environment)
-
         logger.info(
             "Writing Kubeflow workflow definition to `%s`.", pipeline_file_path
         )

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -392,7 +392,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs the pipeline in Kubernetes.
@@ -501,6 +501,9 @@ def prepare_or_run_pipeline(
                 }
             )
 
+        # Use the env from any step for the orchestrator pod
+        orchestrator_pod_env = environment.popitem()[1]
+
         # Schedule as CRON job if CRON schedule is given.
         if deployment.schedule:
             if not deployment.schedule.cron_expression:
@@ -521,7 +524,7 @@ def prepare_or_run_pipeline(
                 service_account_name=service_account_name,
                 privileged=False,
                 pod_settings=orchestrator_pod_settings,
-                env=environment,
+                env=orchestrator_pod_env,
                 mount_local_stores=self.config.is_local,
             )
 
@@ -546,7 +549,7 @@ def prepare_or_run_pipeline(
                 privileged=False,
                 pod_settings=orchestrator_pod_settings,
                 service_account_name=service_account_name,
-                env=environment,
+                env=orchestrator_pod_env,
                 mount_local_stores=self.config.is_local,
             )
 

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -45,6 +45,7 @@
     get_config_environment_vars,
     get_orchestrator_run_name,
 )
+from zenml.utils import env_utils
 
 logger = get_logger(__name__)
 
@@ -94,8 +95,8 @@ def main() -> None:
     kube_client = orchestrator.get_kube_client(incluster=True)
     core_api = k8s_client.CoreV1Api(kube_client)
 
-    env = get_config_environment_vars()
-    env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
+    shared_env = get_config_environment_vars()
+    shared_env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
 
     def run_step_on_kubernetes(step_name: str) -> None:
         """Run a pipeline step in a separate Kubernetes pod.
@@ -112,6 +113,13 @@ def run_step_on_kubernetes(step_name: str) -> None:
             settings.model_dump() if settings else {}
         )
 
+        step_env = shared_env.copy()
+        step_env.update(
+            env_utils.get_step_environment(
+                step_config=step_config, stack=active_stack
+            )
+        )
+
         if settings.pod_name_prefix and not orchestrator_run_id.startswith(
             settings.pod_name_prefix
         ):

src/zenml/integrations/tekton/orchestrators/tekton_orchestrator.py

@@ -49,7 +49,7 @@
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
-from zenml.utils import io_utils, yaml_utils
+from zenml.utils import io_utils
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
@@ -459,7 +459,7 @@ def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
-        environment: Dict[str, str],
+        environment: Dict[str, Dict[str, str]],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs the pipeline on Tekton.
@@ -559,6 +559,7 @@ def dynamic_pipeline() -> None:
                     component_name,
                     component,
                 ) in step_name_to_dynamic_component.items():
+                    step_environment = environment[component_name]
                     # for each component, check to see what other steps are
                     # upstream of it
                     step = deployment.step_configurations[component_name]
@@ -578,6 +579,8 @@ def dynamic_pipeline() -> None:
                         )
                         .after(*upstream_step_components)
                     )
+                    for key, value in step_environment.items():
+                        task = task.set_env_variable(name=key, value=value)
                     self._configure_container_resources(
                         task,
                         step.config.resource_settings,
@@ -586,39 +589,6 @@ def dynamic_pipeline() -> None:
 
             return dynamic_pipeline
 
-        def _update_yaml_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_yaml(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_yaml(pipeline_file_path, pipeline_definition)
-
-            print(
-                f"Updated YAML file with environment variables at {yaml_file_path}"
-            )
-
         # Get a filepath to use to save the finished yaml to
         fileio.makedirs(self.pipeline_directory)
         pipeline_file_path = os.path.join(
@@ -631,9 +601,6 @@ def _update_yaml_with_environment(
             pipeline_name=orchestrator_run_name,
         )
 
-        # Let's update the YAML file with the environment variables
-        _update_yaml_with_environment(pipeline_file_path, environment)
-
         logger.info(
             "Writing Tekton workflow definition to `%s`.", pipeline_file_path
         )