mount cache dir to the NIM model server container (#2965)

* mount cache dir to nim container Signed-off-by: Samhita Alla <[email protected]> * add comment Signed-off-by: Samhita Alla <[email protected]> --------- Signed-off-by: Samhita Alla <[email protected]>
flyteorg · Dec 3, 2024 · 5ec1aeb · 5ec1aeb
1 parent ba8deff
commit 5ec1aeb
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/plugins/flytekit-inference/flytekitplugins/inference/nim/serve.py b/plugins/flytekit-inference/flytekitplugins/inference/nim/serve.py
@@ -114,6 +114,9 @@ def setup_nim_pod_template(self):
         model_server_container.volume_mounts = [V1VolumeMount(name="dshm", mount_path="/dev/shm")]
         model_server_container.security_context = V1SecurityContext(run_as_user=1000)
 
+        self.pod_template.pod_spec.volumes.append(V1Volume(name="cache", empty_dir={}))
+        model_server_container.volume_mounts.append(V1VolumeMount(name="cache", mount_path="/opt/nim/.cache"))
+
         # Download HF LoRA adapters
         if self._hf_repo_ids:
             if not self._lora_adapter_mem:
@@ -131,12 +134,13 @@ def setup_nim_pod_template(self):
                 None,
             )
             if local_peft_dir_env:
-                mount_path = local_peft_dir_env.value
+                peft_mount_path = local_peft_dir_env.value
             else:
+                # This is the directory where all LoRAs are stored for a particular model.
                 raise ValueError("NIM_PEFT_SOURCE environment variable must be set.")
 
             self.pod_template.pod_spec.volumes.append(V1Volume(name="lora", empty_dir={}))
-            model_server_container.volume_mounts.append(V1VolumeMount(name="lora", mount_path=mount_path))
+            model_server_container.volume_mounts.append(V1VolumeMount(name="lora", mount_path=peft_mount_path))
 
             self.pod_template.pod_spec.init_containers.insert(
                 0,
@@ -149,7 +153,7 @@ def setup_nim_pod_template(self):
                         f"""
             pip install -U "huggingface_hub[cli]"
 
-            export LOCAL_PEFT_DIRECTORY={mount_path}
+            export LOCAL_PEFT_DIRECTORY={peft_mount_path}
             mkdir -p $LOCAL_PEFT_DIRECTORY
 
             TOKEN_VAR_NAME={self._secrets.secrets_prefix}{hf_key}
@@ -175,7 +179,7 @@ def setup_nim_pod_template(self):
                     volume_mounts=[
                         V1VolumeMount(
                             name="lora",
-                            mount_path=mount_path,
+                            mount_path=peft_mount_path,
                         )
                     ],
                 ),