diff --git a/src/gpu_rent/ready.py b/src/gpu_rent/ready.py index bf61636..e5db4d4 100644 --- a/src/gpu_rent/ready.py +++ b/src/gpu_rent/ready.py @@ -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")] diff --git a/src/gpu_rent/ssh_ops.py b/src/gpu_rent/ssh_ops.py index 5a5bd22..0fd984c 100644 --- a/src/gpu_rent/ssh_ops.py +++ b/src/gpu_rent/ssh_ops.py @@ -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() diff --git a/tests/test_verify_stack.py b/tests/test_verify_stack.py index bd29c57..e54b69c 100644 --- a/tests/test_verify_stack.py +++ b/tests/test_verify_stack.py @@ -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)