Make external GPU CPU pinning conditional

This commit is contained in:
firestar5683
2026-09-08 18:48:40 -05:00
parent 2360ff9b0f
commit ca3d8a3816
2 changed files with 17 additions and 2 deletions
+7 -2
View File
@@ -84,9 +84,14 @@ def wait_for_external_gpu() -> None:
def external_gpu_compile_command(command: list[str]) -> list[str]:
"""Pin USB-GPU compilation to AGNOS' isolated CPU without changing host builds."""
"""Pin USB-GPU compilation when AGNOS exposes the isolated CPU."""
if sys.platform == "linux" and platform.machine() == "aarch64":
return ["taskset", "-c", "7", *command]
try:
available_cpus = os.sched_getaffinity(0)
if 7 in available_cpus:
return ["taskset", "-c", "7", *command]
except (AttributeError, OSError):
pass
return command
@@ -405,10 +405,20 @@ def test_external_gpu_compile_uses_agnos_isolated_cpu(monkeypatch):
command = ["python3", "compile_modeld.py"]
monkeypatch.setattr(model_compiler.sys, "platform", "linux")
monkeypatch.setattr(model_compiler.platform, "machine", lambda: "aarch64")
monkeypatch.setattr(model_compiler.os, "sched_getaffinity", lambda _: {7}, raising=False)
assert model_compiler.external_gpu_compile_command(command) == ["taskset", "-c", "7", *command]
def test_external_gpu_compile_skips_unavailable_agnos_cpu(monkeypatch):
command = ["python3", "compile_modeld.py"]
monkeypatch.setattr(model_compiler.sys, "platform", "linux")
monkeypatch.setattr(model_compiler.platform, "machine", lambda: "aarch64")
monkeypatch.setattr(model_compiler.os, "sched_getaffinity", lambda _: {0, 1, 2, 3}, raising=False)
assert model_compiler.external_gpu_compile_command(command) is command
def test_external_gpu_compile_does_not_pin_other_platforms(monkeypatch):
command = ["python3", "compile_modeld.py"]
monkeypatch.setattr(model_compiler.sys, "platform", "darwin")