[SDK] Allow customizing base trainer and storage images in Train API

Allow customizing base storage_initializer and trainer images through Env vars. Signed-off-by: Varsha Prasad Narsing <[email protected]>
kubeflow · Sep 18, 2024 · 40c5c10 · 40c5c10
1 parent 126110f
commit 40c5c10
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
@@ -15,6 +15,7 @@
 import json
 import logging
 import multiprocessing
+import os
 import queue
 import time
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -242,7 +243,9 @@ def train(
         # create init container spec
         init_container_spec = utils.get_container_spec(
             name=constants.STORAGE_INITIALIZER,
-            base_image=constants.STORAGE_INITIALIZER_IMAGE,
+            base_image=os.getenv(
+                "STORAGE_INITIALIZER_IMAGE", constants.STORAGE_INITIALIZER_IMAGE_DEFAULT
+            ),
             args=[
                 "--model_provider",
                 mp,
@@ -259,7 +262,10 @@ def train(
         # create app container spec
         container_spec = utils.get_container_spec(
             name=constants.JOB_PARAMETERS[constants.PYTORCHJOB_KIND]["container"],
-            base_image=constants.TRAINER_TRANSFORMER_IMAGE,
+            base_image=os.getenv(
+                "TRAINER_TRANSFORMER_IMAGE_DEFAULT",
+                constants.TRAINER_TRANSFORMER_IMAGE_DEFAULT,
+            ),
             args=[
                 "--model_uri",
                 model_provider_parameters.model_uri,

diff --git a/sdk/python/kubeflow/training/constants/constants.py b/sdk/python/kubeflow/training/constants/constants.py
@@ -82,14 +82,14 @@
 
 
 # TODO (andreyvelich): We should add image tag for Storage Initializer and Trainer.
-STORAGE_INITIALIZER_IMAGE = "docker.io/kubeflow/storage-initializer"
+STORAGE_INITIALIZER_IMAGE_DEFAULT = "docker.io/kubeflow/storage-initializer"
 
 STORAGE_INITIALIZER_VOLUME_MOUNT = models.V1VolumeMount(
     name=STORAGE_INITIALIZER,
     mount_path=INIT_CONTAINER_MOUNT_PATH,
 )
 
-TRAINER_TRANSFORMER_IMAGE = "docker.io/kubeflow/trainer-huggingface"
+TRAINER_TRANSFORMER_IMAGE_DEFAULT = "docker.io/kubeflow/trainer-huggingface"
 
 # TFJob constants.
 TFJOB_KIND = "TFJob"