chore: import upstream snapshot with attribution

python/ray/tests/accelerators/test_neuron.py

@@ -0,0 +1,104 @@
+import subprocess
+import sys
+from unittest.mock import patch
+
+import pytest
+
+import ray
+from ray._private.accelerators import NeuronAcceleratorManager
+
+
+def test_user_configured_more_than_visible(monkeypatch, call_ray_stop_only):
+    # Test more neuron_cores are configured than visible.
+    monkeypatch.setenv("NEURON_RT_VISIBLE_CORES", "0,1,2")
+    with pytest.raises(ValueError):
+        ray.init(resources={"neuron_cores": 4})
+
+
+@patch(
+    "ray._private.accelerators.NeuronAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=4,
+)
+def test_auto_detected_more_than_visible(
+    mock_get_num_accelerators, monkeypatch, shutdown_only
+):
+    # Test more neuron_cores are detected than visible.
+    monkeypatch.setenv("NEURON_RT_VISIBLE_CORES", "0,1,2")
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["neuron_cores"] == 3
+
+
+@patch(
+    "ray._private.accelerators.NeuronAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
+    return_value=2,
+)
+def test_auto_detect_resources(mock_get_num_accelerators, shutdown_only):
+    # Test that ray node resources are filled with auto detected count.
+    ray.init()
+    _ = mock_get_num_accelerators.called
+    assert ray.available_resources()["neuron_cores"] == 2
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout=(
+            b'[{"neuron_device":0,"bdf":"00:1e.0",'
+            b'"connected_to":null,"nc_count":2,'
+            b'"memory_size":34359738368,"neuron_processes":[]}]'
+        ),
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_single_device(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 2
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout=(
+            b'[{"neuron_device":0,"bdf":"00:1e.0",'
+            b'"connected_to":null,"nc_count":2,'
+            b'"memory_size":34359738368,"neuron_processes":[]},'
+            b'{"neuron_device":1,"bdf":"00:1f.0","connected_to":null,'
+            b'"nc_count":2,"memory_size":34359738368,"neuron_processes":[]}]'
+        ),
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_multiple_devices(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 4
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(
+        args=[], returncode=1, stdout=b"AccessDenied"
+    ),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_failure_with_error(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+@patch(
+    "subprocess.run",
+    return_value=subprocess.CompletedProcess(args=[], returncode=0, stdout=b"[{}]"),
+)
+@patch("os.path.isdir", return_value=True)
+@patch("sys.platform", "linux")
+def test_get_neuron_core_count_failure_with_empty_results(mock_isdir, mock_subprocess):
+    assert NeuronAcceleratorManager.get_current_node_num_accelerators() == 0
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))