[python] return neuron cores instead of neuron devices

deepjavalibrary · Jan 29, 2025 · 3f1b901 · 3f1b901
1 parent 0495859
commit 3f1b901
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/neuron_utils/neuron_smart_default_utils.py b/engines/python/setup/djl_python/neuron_utils/neuron_smart_default_utils.py
@@ -112,7 +112,9 @@ def get_available_cores():
         command = "neuron-ls --json-output"
         try:
             output = sp.check_output(command, shell=True).decode("utf-8")
-            return len(json.loads(output))
+            output = json.loads(output)
+            neuron_cores = sum(device.get("nc_count", 0) for device in output)
+            return neuron_cores
         except Exception as e:
             logger.debug(f"Failed to get available cores: {e}")
         return 0

diff --git a/engines/python/setup/djl_python/tests/neuron_test_scripts/test_neuron_smart_default_utils.py b/engines/python/setup/djl_python/tests/neuron_test_scripts/test_neuron_smart_default_utils.py
@@ -66,7 +66,9 @@ class TestNeuronSmartDefaultUtils(unittest.TestCase):
 
     def test_get_available_cores(self):
         with patch('subprocess.check_output') as mock_check_output:
-            mock_check_output.return_value = b'[{"core": 0}, {"core": 1}, {"core": 2}, {"core": 3}]'
+            mock_check_output.return_value = (
+                b'[{"core": 0, "nc_count": 2}, {"core": 1, "nc_count": 2}, '
+                b'{"core": 2, "nc_count": 2}, {"core": 3, "nc_count": 2}]')
             assert NeuronSmartDefaultUtils.get_available_cores() == 4
 
     def test_get_available_cores_exception(self):