Add changes for RHEL AI image 1.3

opendatahub-io · Dec 10, 2024 · 6035ea0 · 6035ea0
1 parent ce27915
commit 6035ea0
Showing 1 changed file with 7 additions and 15 deletions.
diff --git a/training/components.py b/training/components.py
@@ -252,39 +252,31 @@ def list_phase1_final_model():
             persistent_volume_claim=models.V1PersistentVolumeClaimVolumeSource(
                 claim_name=output_pvc_name
             ),
-        )
+        ),
     ]
 
     # Set volume mounts
     volume_mounts_common = [
         models.V1VolumeMount(
             mount_path="/input_data", name="input-data", read_only=True
         ),
-        models.V1VolumeMount(
-            mount_path="/input_model", name="model", read_only=True
-        )
+        models.V1VolumeMount(mount_path="/input_model", name="model", read_only=True),
     ]
     volume_mounts_master = volume_mounts_common.append(
-        models.V1VolumeMount(
-            mount_path="/output", name="output"
-        )
+        models.V1VolumeMount(mount_path="/output", name="output")
     )
     volume_mounts_worker = volume_mounts_common.append(
-        models.V1VolumeMount(
-            mount_path="/output", name="output", read_only=True
-        )
+        models.V1VolumeMount(mount_path="/output", name="output", read_only=True)
     )
 
     # Set env variables
     env_vars = [
         models.V1EnvVar(name="NNODES", value=f"{nnodes}"),
-        models.V1EnvVar(
-            name="NPROC_PER_NODE", value=f"{nproc_per_node}"
-        ),
+        models.V1EnvVar(name="NPROC_PER_NODE", value=f"{nproc_per_node}"),
         models.V1EnvVar(name="XDG_CACHE_HOME", value="/tmp"),
         models.V1EnvVar(name="TRITON_CACHE_DIR", value="/tmp"),
         models.V1EnvVar(name="HF_HOME", value="/tmp"),
-        models.V1EnvVar(name="TRANSFORMERS_CACHE", value="/tmp")
+        models.V1EnvVar(name="TRANSFORMERS_CACHE", value="/tmp"),
     ]
 
     # Get master and worker container specs
@@ -409,4 +401,4 @@ def wait_for_job_get_logs(
     )
     if delete_after_done:
         logging.info("Deleting job after completion.")
-        training_client.delete_job(name, namespace)
+        training_client.delete_job(name, namespace)