[SDK] move env var to constants.py (#2268)

Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>

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,

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"