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
+2 -2
View File
@@ -323,7 +323,6 @@ def verify_gpu_env(
encoding="utf-8"
)
swarm_flag = "1" if want_swarm else "0"
prefix = f"import os\nos.environ['GPU_RENT_CHECK_SWARM']={swarm_flag!r}\n"
log(
"проверка GPU-стека: nvidia-smi, CUDA"
+ (", torch в Comfy venv" if want_swarm else " (llm-only — без torch)")
@@ -337,10 +336,11 @@ def verify_gpu_env(
out = run_python(
cfg,
host,
prefix + script,
script,
remote_path="/tmp/gpu-rent-stack_env_probe.py",
timeout=180,
log=None,
env={"GPU_RENT_CHECK_SWARM": swarm_flag},
)
except Exception as exc:
last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
+16 -2
View File
@@ -352,8 +352,13 @@ def run_python(
remote_path: str,
timeout: int = 1800,
log: Callable[[str], None] | None = None,
env: dict[str, str] | None = None,
) -> 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)
chunks: list[str] = []
try:
@@ -362,7 +367,16 @@ def run_python(
fh.write(script)
sftp.chmod(remote_path, 0o755)
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)
_stream_pty_output(stdout, log=log, chunks=chunks)
code = stdout.channel.recv_exit_status()