From 0c98b21ba95b332d3536fa9a5db1f30612ab5734 Mon Sep 17 00:00:00 2001
From: Varsha Prasad Narsing <varshaprasad96@gmail.com>
Date: Mon, 23 Sep 2024 08:30:11 -0700
Subject: [PATCH] [SDK] move env var to constants.py

Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
---
 sdk/python/kubeflow/training/api/training_client.py | 10 ++--------
 sdk/python/kubeflow/training/constants/constants.py |  9 +++++++--
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
index 1626f18820..60c37cd261 100644
--- a/sdk/python/kubeflow/training/api/training_client.py
+++ b/sdk/python/kubeflow/training/api/training_client.py
@@ -15,7 +15,6 @@
 import json
 import logging
 import multiprocessing
-import os
 import queue
 import time
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -243,9 +242,7 @@ def train(
         # create init container spec
         init_container_spec = utils.get_container_spec(
             name=constants.STORAGE_INITIALIZER,
-            base_image=os.getenv(
-                "STORAGE_INITIALIZER_IMAGE", constants.STORAGE_INITIALIZER_IMAGE_DEFAULT
-            ),
+            base_image=constants.STORAGE_INITIALIZER_IMAGE,
             args=[
                 "--model_provider",
                 mp,
@@ -262,10 +259,7 @@ def train(
         # create app container spec
         container_spec = utils.get_container_spec(
             name=constants.JOB_PARAMETERS[constants.PYTORCHJOB_KIND]["container"],
-            base_image=os.getenv(
-                "TRAINER_TRANSFORMER_IMAGE_DEFAULT",
-                constants.TRAINER_TRANSFORMER_IMAGE_DEFAULT,
-            ),
+            base_image=constants.TRAINER_TRANSFORMER_IMAGE,
             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
index e745468736..dbbd885baa 100644
--- a/sdk/python/kubeflow/training/constants/constants.py
+++ b/sdk/python/kubeflow/training/constants/constants.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import Union
 
 from kubeflow.storage_initializer.constants import INIT_CONTAINER_MOUNT_PATH
@@ -82,14 +83,18 @@
 
 
 # TODO (andreyvelich): We should add image tag for Storage Initializer and Trainer.
-STORAGE_INITIALIZER_IMAGE_DEFAULT = "docker.io/kubeflow/storage-initializer"
+STORAGE_INITIALIZER_IMAGE = os.getenv(
+    "STORAGE_INITIALIZER_IMAGE", "docker.io/kubeflow/storage-initializer"
+)
 
 STORAGE_INITIALIZER_VOLUME_MOUNT = models.V1VolumeMount(
     name=STORAGE_INITIALIZER,
     mount_path=INIT_CONTAINER_MOUNT_PATH,
 )
 
-TRAINER_TRANSFORMER_IMAGE_DEFAULT = "docker.io/kubeflow/trainer-huggingface"
+TRAINER_TRANSFORMER_IMAGE = os.getenv(
+    "TRAINER_TRANSFORMER_IMAGE", "docker.io/kubeflow/trainer-huggingface"
+)
 
 # TFJob constants.
 TFJOB_KIND = "TFJob"