Skip to content

Commit 1d0a526

Browse files
committed
[python] return neuron cores instead of neuron devices
1 parent 0495859 commit 1d0a526

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

engines/python/setup/djl_python/neuron_utils/neuron_smart_default_utils.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,9 @@ def get_available_cores():
112112
command = "neuron-ls --json-output"
113113
try:
114114
output = sp.check_output(command, shell=True).decode("utf-8")
115-
return len(json.loads(output))
115+
output = json.loads(output)
116+
neuron_cores = sum(device.get("nc_count", 0) for device in output)
117+
return neuron_cores
116118
except Exception as e:
117119
logger.debug(f"Failed to get available cores: {e}")
118120
return 0

0 commit comments

Comments
 (0)