diff --git a/scripts/model_compiler.py b/scripts/model_compiler.py index 3383eb2ca..5b6b55e90 100644 --- a/scripts/model_compiler.py +++ b/scripts/model_compiler.py @@ -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 diff --git a/starpilot/assets/tests/test_model_pipeline.py b/starpilot/assets/tests/test_model_pipeline.py index 5e6eedd90..5b63e391d 100644 --- a/starpilot/assets/tests/test_model_pipeline.py +++ b/starpilot/assets/tests/test_model_pipeline.py @@ -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")