[python] move the default compile path to /tmp/.djl.ai (#2697)

deepjavalibrary · Jan 31, 2025 · 89f0fe4 · 89f0fe4
1 parent d39b347
commit 89f0fe4
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/engines/python/setup/djl_python/transformers_neuronx.py b/engines/python/setup/djl_python/transformers_neuronx.py
@@ -142,14 +142,22 @@ def set_model_loader_class(self) -> None:
         if self.config.model_loader == "nxdi":
             os.environ[
                 'VLLM_NEURON_FRAMEWORK'] = "neuronx-distributed-inference"
+            djl_neuron_compiled_artifacts_path = os.path.join(
+                os.getenv("DJL_CACHE_DIR", "/tmp/.djl.ai"),
+                "neuron-compiled-artifacts")
+            nxdi_compiled_model_path = os.path.join(
+                self.config.model_id_or_path, NXDI_COMPILED_MODEL_FILE_NAME)
             if self.config.save_mp_checkpoint_path:
+                # If the compilation path is given by the user
                 os.environ[
                     "NEURON_COMPILED_ARTIFACTS"] = self.config.save_mp_checkpoint_path
-            nxdi_compiled_model_path = os.path.join(
-                self.config.model_id_or_path, NXDI_COMPILED_MODEL_FILE_NAME)
-            if os.path.isfile(nxdi_compiled_model_path):
+            elif os.path.isfile(nxdi_compiled_model_path):
+                # if the compilation path already exists
                 os.environ[
                     "NEURON_COMPILED_ARTIFACTS"] = self.config.model_id_or_path
+            else:
+                os.environ[
+                    "NEURON_COMPILED_ARTIFACTS"] = djl_neuron_compiled_artifacts_path
             return
 
         if self.config.model_loader == "vllm":