From 69daf81a43d092e115bc5d932a87bb24e9be92bb Mon Sep 17 00:00:00 2001
From: Leonid Pershin <lpershin@syntellect.ru>
Date: Fri, 21 Aug 2026 09:02:07 +0300
Subject: [PATCH] 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.
---
 src/gpu_rent/ready.py      |  4 ++--
 src/gpu_rent/ssh_ops.py    | 18 ++++++++++++++++--
 tests/test_verify_stack.py |  8 +++++---
 3 files changed, 23 insertions(+), 7 deletions(-)

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)