daemon: handle put/get file transfers

daemon/client.py

@@ -21,13 +21,14 @@ import sys
 import time
 import signal
 import socket
+import shutil
 import getpass
 import threading
 import subprocess
 import uuid
 import queue as queuelib
 
-from gradio_client import Client
+from gradio_client import Client, handle_file
 
 
 def _primary_ip() -> str:
@@ -61,6 +62,7 @@ CLIENT_ID = os.environ.get("CLIENT_ID") or _default_client_id()
 POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "5"))
 DEFAULT_TIMEOUT = int(os.environ.get("CMD_TIMEOUT", "120"))
 MAX_OUTPUT_CHARS = 200_000
+MAX_FILE_BYTES = 100 * 1024 * 1024  # 100 MB cap on transfers
 
 _log_lock = threading.Lock()
 
@@ -207,7 +209,65 @@ def _meta() -> str:
     return f"{os.uname().sysname} {os.uname().release} pid={os.getpid()}"
 
 
+def _expand(p: str) -> str:
+    return os.path.expanduser(os.path.expandvars(p or ""))
+
+
+def _result_path(ret) -> str:
+    """gradio_client returns a File output as a path str, or a FileData dict."""
+    if isinstance(ret, dict):
+        return ret.get("path") or ret.get("name") or ""
+    if isinstance(ret, (list, tuple)) and ret:
+        return _result_path(ret[0])
+    return ret
+
+
+def do_put(client: Client, cmd: dict):
+    """Server -> client: download the payload from the relay and write it to dest."""
+    cid = cmd["id"]
+    dest = _expand(cmd.get("dest") or cmd.get("filename") or "downloaded.bin")
+    log(f"[file] put → {dest}")
+    try:
+        ret = client.predict(cmd["file_id"], SECRET, api_name="/fetch_file")
+        src = _result_path(ret)
+        if os.path.isdir(dest) or dest.endswith(os.sep):
+            dest = os.path.join(dest, cmd.get("filename") or os.path.basename(src))
+        parent = os.path.dirname(os.path.abspath(dest))
+        os.makedirs(parent, exist_ok=True)
+        shutil.copy(src, dest)
+        size = os.path.getsize(dest)
+        out, code = (f"wrote {size} bytes to {dest}", 0)
+    except Exception as e:
+        out, code = (f"[relay: put failed: {e!r}]", 1)
+    client.predict(cid, CLIENT_ID, SECRET, code, out, "file", api_name="/result")
+    log(f"[file] put → exit {code}")
+
+
+def do_get(client: Client, cmd: dict):
+    """Client -> server: read a local file and upload it to the relay."""
+    cid = cmd["id"]
+    remote = _expand(cmd.get("remote_path"))
+    log(f"[file] get {remote}")
+    try:
+        if not os.path.isfile(remote):
+            raise FileNotFoundError(f"not a file: {remote}")
+        size = os.path.getsize(remote)
+        if size > MAX_FILE_BYTES:
+            raise ValueError(f"{size} bytes exceeds the {MAX_FILE_BYTES}-byte limit")
+        client.predict(cid, CLIENT_ID, SECRET, handle_file(remote), "",
+                       api_name="/upload_result")
+        log(f"[file] get → sent {size} bytes")
+    except Exception as e:
+        client.predict(cid, CLIENT_ID, SECRET, None, f"{e}", api_name="/upload_result")
+        log(f"[file] get → error: {e}")
+
+
 def handle(client: Client, cmd: dict):
+    typ = cmd.get("type", "exec")
+    if typ == "put":
+        return do_put(client, cmd)
+    if typ == "get":
+        return do_get(client, cmd)
     cid = cmd["id"]
     session = cmd.get("session") or "main"
     command = cmd["command"]