adding env vars

Signed-off-by: Tarek Abouzeid <[email protected]> using same implementation in katib Signed-off-by: Tarek Abouzeid <[email protected]> updating docs Signed-off-by: Tarek Abouzeid <[email protected]>
kubeflow · Oct 16, 2024 · 652a0d5 · 652a0d5
1 parent 078a16c
commit 652a0d5
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
@@ -102,6 +102,12 @@ def train(
         model_provider_parameters=None,
         dataset_provider_parameters=None,
         trainer_parameters=None,
+        init_env_vars: Optional[
+            Union[Dict[str, str], List[Union[models.V1EnvVar, models.V1EnvVar]]]
+        ] = None,
+        worker_env_vars: Optional[
+            Union[Dict[str, str], List[Union[models.V1EnvVar, models.V1EnvVar]]]
+        ] = None,
         storage_config: Dict[str, Optional[Union[str, List[str]]]] = {
             "size": constants.PVC_DEFAULT_SIZE,
             "storage_class": None,
@@ -164,6 +170,20 @@ def train(
                 and HuggingFace training arguments like optimizer or number of training epochs.
                 This argument must be the type of
                 `kubeflow.storage_initializer.HuggingFaceTrainerParams`
+            init_env_vars: Environment variable(s) to be attached to init container.
+                You can specify a dictionary as a mapping object representing the environment
+                variables. Otherwise, you can specify a list, in which the element can either
+                be a kubernetes.client.models.V1EnvVar (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvVar.md)
+                or a kubernetes.client.models.V1EnvFromSource (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvFromSource.md)
+            worker_env_vars: Environment variable(s) to be attached to training container.
+                You can specify a dictionary as a mapping object representing the environment
+                variables. Otherwise, you can specify a list, in which the element can either
+                be a kubernetes.client.models.V1EnvVar (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvVar.md)
+                or a kubernetes.client.models.V1EnvFromSource (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvFromSource.md)
             storage_config: Configuration for Storage Initializer PVC to download pre-trained model
                 and dataset. You can configure PVC size and storage class name in this argument.
         """
@@ -254,6 +274,7 @@ def train(
                 json.dumps(dataset_provider_parameters.__dict__),
             ],
             volume_mounts=[constants.STORAGE_INITIALIZER_VOLUME_MOUNT],
+            env=init_env_vars,
         )
 
         # create app container spec
@@ -280,6 +301,7 @@ def train(
             ],
             volume_mounts=[constants.STORAGE_INITIALIZER_VOLUME_MOUNT],
             resources=resources_per_worker,
+            env=worker_env_vars,
         )
 
         storage_initializer_volume = models.V1Volume(
@@ -329,6 +351,9 @@ def create_job(
         num_ps_replicas: Optional[int] = None,
         packages_to_install: Optional[List[str]] = None,
         pip_index_url: str = constants.DEFAULT_PIP_INDEX_URL,
+        env_vars: Optional[
+            Union[Dict[str, str], List[Union[models.V1EnvVar, models.V1EnvVar]]]
+        ] = None,
     ):
         """Create the Training Job.
         Job can be created using one of the following options:
@@ -386,6 +411,13 @@ def create_job(
                 to the base image packages if `train_func` parameter is set.
                 These packages are installed before executing the objective function.
             pip_index_url: The PyPI url from which to install Python packages.
+            env_vars: Environment variable(s) to be attached to training container.
+                You can specify a dictionary as a mapping object representing the environment
+                variables. Otherwise, you can specify a list, in which the element can either
+                be a kubernetes.client.models.V1EnvVar (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvVar.md)
+                or a kubernetes.client.models.V1EnvFromSource (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvFromSource.md)
 
         Raises:
             ValueError: Invalid input parameters.
@@ -463,6 +495,7 @@ def create_job(
                 command=command,
                 args=args,
                 resources=resources_per_worker,
+                env=env_vars,
             )
 
             # Get Pod template spec using the above container.

diff --git a/sdk/python/kubeflow/training/utils/utils.py b/sdk/python/kubeflow/training/utils/utils.py
@@ -195,6 +195,9 @@ def get_container_spec(
     args: Optional[List[str]] = None,
     resources: Union[dict, models.V1ResourceRequirements, None] = None,
     volume_mounts: Optional[List[models.V1VolumeMount]] = None,
+    env: Optional[
+        Union[Dict[str, str], List[Union[models.V1EnvVar, models.V1EnvVar]]]
+    ] = None,
 ) -> models.V1Container:
     """
     Get container spec for the given parameters.
@@ -210,6 +213,7 @@ def get_container_spec(
         command=command,
         args=args,
         volume_mounts=volume_mounts,
+        env=env,
     )
 
     # Convert dict to the Kubernetes container resources if that is required.