Test for neuron-ls issue

huggingface · Sep 7, 2023 · d3f1929 · d3f1929
1 parent 81ad5bb
commit d3f1929
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/optimum/neuron/utils/cache_utils.py b/optimum/neuron/utils/cache_utils.py
@@ -258,11 +258,10 @@ def set_neuron_cache_path(neuron_cache_path: Union[str, Path], ignore_no_cache:
 
 def get_num_neuron_cores() -> int:
     path = os.environ["PATH"]
-    env = dict(os.environ)
     if NEURON_BINARIES_PATH not in path:
         path = f"{NEURON_BINARIES_PATH}:{path}"
-        env["PATH"] = path
-    proc = subprocess.Popen(["neuron-ls", "-j"], stdout=subprocess.PIPE, env=env)
+        os.environ["PATH"] = path
+    proc = subprocess.Popen(["neuron-ls", "-j"], stdout=subprocess.PIPE)
     stdout, _ = proc.communicate()
     stdout = stdout.decode("utf-8")
     json_stdout = json.loads(stdout)

diff --git a/tests/distributed/test_model_parallelization.py b/tests/distributed/test_model_parallelization.py
@@ -44,6 +44,7 @@
 )
 
 from optimum.neuron.utils.cache_utils import get_num_neuron_cores
+from optimum.neuron.utils.import_utils import is_neuronx_available
 
 from ..test_utils import is_trainium_test
 
@@ -53,7 +54,10 @@
 
 
 TEMPLATE_FILE_NAME = "model_parallel_test_template.txt"
-NUM_NEURON_CORES_AVAILABLE = get_num_neuron_cores()
+if is_neuronx_available():
+    NUM_NEURON_CORES_AVAILABLE = get_num_neuron_cores()
+else:
+    NUM_NEURON_CORES_AVAILABLE = 0
 
 
 def _generate_supported_model_class_names(