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:
@@ -323,7 +323,6 @@ def verify_gpu_env(
|
|||||||
encoding="utf-8"
|
encoding="utf-8"
|
||||||
)
|
)
|
||||||
swarm_flag = "1" if want_swarm else "0"
|
swarm_flag = "1" if want_swarm else "0"
|
||||||
prefix = f"import os\nos.environ['GPU_RENT_CHECK_SWARM']={swarm_flag!r}\n"
|
|
||||||
log(
|
log(
|
||||||
"проверка GPU-стека: nvidia-smi, CUDA"
|
"проверка GPU-стека: nvidia-smi, CUDA"
|
||||||
+ (", torch в Comfy venv" if want_swarm else " (llm-only — без torch)")
|
+ (", torch в Comfy venv" if want_swarm else " (llm-only — без torch)")
|
||||||
@@ -337,10 +336,11 @@ def verify_gpu_env(
|
|||||||
out = run_python(
|
out = run_python(
|
||||||
cfg,
|
cfg,
|
||||||
host,
|
host,
|
||||||
prefix + script,
|
script,
|
||||||
remote_path="/tmp/gpu-rent-stack_env_probe.py",
|
remote_path="/tmp/gpu-rent-stack_env_probe.py",
|
||||||
timeout=180,
|
timeout=180,
|
||||||
log=None,
|
log=None,
|
||||||
|
env={"GPU_RENT_CHECK_SWARM": swarm_flag},
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
|
last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
|
||||||
|
|||||||
+16
-2
@@ -352,8 +352,13 @@ def run_python(
|
|||||||
remote_path: str,
|
remote_path: str,
|
||||||
timeout: int = 1800,
|
timeout: int = 1800,
|
||||||
log: Callable[[str], None] | None = None,
|
log: Callable[[str], None] | None = None,
|
||||||
|
env: dict[str, str] | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Upload a Python script and run it as SSH user (not root)."""
|
"""Upload a Python script and run it as SSH user (not root).
|
||||||
|
|
||||||
|
Do not prepend Python code to ``script`` (breaks ``from __future__``).
|
||||||
|
Pass runtime flags via ``env`` instead.
|
||||||
|
"""
|
||||||
client = _connect(cfg, host)
|
client = _connect(cfg, host)
|
||||||
chunks: list[str] = []
|
chunks: list[str] = []
|
||||||
try:
|
try:
|
||||||
@@ -362,7 +367,16 @@ def run_python(
|
|||||||
fh.write(script)
|
fh.write(script)
|
||||||
sftp.chmod(remote_path, 0o755)
|
sftp.chmod(remote_path, 0o755)
|
||||||
sftp.close()
|
sftp.close()
|
||||||
command = f"python3 {shlex.quote(remote_path)}"
|
exports = ""
|
||||||
|
if env:
|
||||||
|
parts = [
|
||||||
|
f"{key}={shlex.quote(str(val))}"
|
||||||
|
for key, val in env.items()
|
||||||
|
if key
|
||||||
|
]
|
||||||
|
if parts:
|
||||||
|
exports = " ".join(parts) + " "
|
||||||
|
command = f"{exports}python3 {shlex.quote(remote_path)}"
|
||||||
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
_stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
|
||||||
_stream_pty_output(stdout, log=log, chunks=chunks)
|
_stream_pty_output(stdout, log=log, chunks=chunks)
|
||||||
code = stdout.channel.recv_exit_status()
|
code = stdout.channel.recv_exit_status()
|
||||||
|
|||||||
@@ -81,7 +81,9 @@ def test_verify_gpu_env_ok(monkeypatch):
|
|||||||
}
|
}
|
||||||
|
|
||||||
def fake_run_python(cfg, host, script, **kw):
|
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"
|
return json.dumps(payload) + "\n"
|
||||||
|
|
||||||
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
|
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):
|
def fake_run_python(cfg, host, script, **kw):
|
||||||
assert "GPU_RENT_CHECK_SWARM" in script
|
assert "from __future__" in script
|
||||||
assert "os.environ['GPU_RENT_CHECK_SWARM']='0'" in script
|
assert (kw.get("env") or {}).get("GPU_RENT_CHECK_SWARM") == "0"
|
||||||
return json.dumps(payload)
|
return json.dumps(payload)
|
||||||
|
|
||||||
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
|
monkeypatch.setattr(ssh_ops, "run_python", fake_run_python)
|
||||||
|
|||||||
Reference in New Issue
Block a user