Refactor GPU environment verification and SSH execution for improved flexibility

- Removed the direct injection of environment variables into the script, preventing issues with `from __future__` imports.
- Updated the `run_python` function to accept an `env` parameter for passing runtime environment variables, enhancing script execution control.
- Adjusted tests to verify the new behavior, ensuring that environment variables are correctly set without modifying the script content.
This commit is contained in:
Leonid Pershin
2026-08-21 09:02:07 +03:00
parent f17b1c9731
commit 69daf81a43
3 changed files with 23 additions and 7 deletions
+5 -3
View File
@@ -81,7 +81,9 @@ def test_verify_gpu_env_ok(monkeypatch):
}
def fake_run_python(cfg, host, script, **kw):
assert "GPU_RENT_CHECK_SWARM" in script
assert "from __future__" in script
assert not script.lstrip().startswith("import os")
assert (kw.get("env") or {}).get("GPU_RENT_CHECK_SWARM") == "1"
return json.dumps(payload) + "\n"
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
@@ -171,8 +173,8 @@ def test_verify_gpu_env_llm_only_skips_torch_requirement(monkeypatch):
}
def fake_run_python(cfg, host, script, **kw):
assert "GPU_RENT_CHECK_SWARM" in script
assert "os.environ['GPU_RENT_CHECK_SWARM']='0'" in script
assert "from __future__" in script
assert (kw.get("env") or {}).get("GPU_RENT_CHECK_SWARM") == "0"
return json.dumps(payload)
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)