From 652a0d50d3444c2b1575b9abbaeca3f8297be5f2 Mon Sep 17 00:00:00 2001 From: Tarek Abouzeid Date: Tue, 15 Oct 2024 18:50:53 +0200 Subject: [PATCH] adding env vars Signed-off-by: Tarek Abouzeid using same implementation in katib Signed-off-by: Tarek Abouzeid updating docs Signed-off-by: Tarek Abouzeid --- .../kubeflow/training/api/training_client.py | 33 +++++++++++++++++++ sdk/python/kubeflow/training/utils/utils.py | 4 +++ 2 files changed, 37 insertions(+) diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py index 808f07065b..206ef1901e 100644 --- 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 index 4f2d9c97d7..c084f58cb4 100644 --- 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.