fcloud-sdk 0.1.0__py3-none-any.whl

fcloud/cli/migration.py

@@ -0,0 +1,276 @@
+"""Riding a spot migration: resume a foreground process on its new host.
+
+Reconnect UX, the restore budget, the resumed-process drain, and the
+log-follow reconnect/tail-delta helpers.
+"""
+
+import os
+import sys
+import time
+
+from fcloud.cli.attach import (
+    _ENDED_SESSION_STATES,
+    _RESUME_BUDGET_S,
+    _TERMINAL_SESSION_STATES,
+    _authoritative_session_status,
+)
+from fcloud.client import Client, FcloudError, PaymentOverdueError
+from fcloud.types import ProcessStatus
+
+
+def _migration_ui():
+    """User-facing reconnect notice callbacks, shared by every CLI path that
+    rides a connection drop (``logs --follow``, ``exec``) so the message is
+    identical everywhere. Returns ``(on_migrating, on_reconnected)`` to hand to
+    ``Client.reconnect_session``.
+
+    Whether the drop is an actual host migration is only known after the
+    re-attach (via the ``migrated`` flag), so the up-front notice stays
+    neutral: a data-plane flap on a healthy host must not be narrated as
+    infra churn.
+    """
+    def on_migrating() -> None:
+        sys.stderr.write(
+            "\n\033[33m⏳ connection interrupted — reconnecting, "
+            "no action needed.\033[0m\n")
+        sys.stderr.flush()
+
+    def on_reconnected(host_id: str, migrated: bool) -> None:
+        where = (f"resumed on a new host ({host_id})" if migrated
+                 else f"same host ({host_id})")
+        sys.stderr.write(
+            f"\033[32m✅ reconnected — {where}. Continuing.\033[0m\n")
+        sys.stderr.flush()
+
+    return on_migrating, on_reconnected
+
+
+def _ride_budget(wait_deadline: float | None) -> float:
+    """Remaining reconnect budget for one ride step.
+
+    The caller's --wait deadline bounds the default restore budget: a user who
+    asked for --wait 5m must not be blocked for the full 10-minute restore
+    window (that once held --wait 5m callers for 5-10 minutes). No deadline (wait_s=0,
+    or fcloud run's fixed in-run timeout) keeps the full budget.
+    """
+    if wait_deadline is None:
+        return _RESUME_BUDGET_S
+    return max(0.0, min(wait_deadline - time.monotonic(), _RESUME_BUDGET_S))
+
+
+def _still_migrating(result, sid: str, pid: str):
+    """Shape a ride that ran out of budget: the restore may still land, so
+    don't claim success or death — status "migrating" (exit 75, retryable)
+    lets JSON callers tell a ride-in-progress from a live process."""
+    result.status = "migrating"
+    result.returncode = None
+    result.message = (
+        f"Session {sid} is still being restored. Follow it with: "
+        f"fcloud logs {sid} {pid} --follow")
+    return result
+
+
+def _ride_migration(client: Client, result, on_output,
+                    wait_deadline: float | None = None):
+    """Ride a spot migration for a foreground run/exec so the user's command
+    resumes in place instead of falsely dying or hanging.
+
+    The bridge returns status="host_lost" when the host vanished with no
+    teardown notice — a checkpoint/restore migration is (likely) in flight.
+    Confirm with the dispatcher that the session is genuinely coming back
+    (NEVER resurrect one that was stopped/suspended/killed), wait for the
+    restore, re-attach to the SAME process, and stream it to completion. Every
+    other result passes through untouched. ``wait_deadline`` (monotonic) is
+    the caller's --wait budget; the ride never outlives it.
+    """
+    if result is None or result.status != "host_lost":
+        return result
+    sid, pid = result.session_id, result.process_id
+    if _authoritative_session_status(client, sid) in _TERMINAL_SESSION_STATES:
+        # Ended by stop/suspension/failure — resuming would rebuild a host for
+        # a deliberately-killed session. Report a non-success terminal so the
+        # caller never reads it as exit 0.
+        return _host_lost_terminal(
+            result, f"Session {sid} ended; not resuming. "
+            f"Captured output: fcloud logs {sid} {pid} --output all")
+    on_migrating, on_reconnected = _migration_ui()
+    try:
+        client.reconnect_session(
+            sid, on_migrating=on_migrating, on_reconnected=on_reconnected,
+            budget=_ride_budget(wait_deadline),
+        )
+    except PaymentOverdueError:
+        return _host_lost_terminal(
+            result, f"Session {sid} suspended (payment overdue); not resuming.")
+    except FcloudError:
+        # Terminal status surfaced during re-attach (closed/failed).
+        return _host_lost_terminal(
+            result, f"Session {sid} could not be restored; "
+            "the process did not survive.")
+    except Exception:
+        return _still_migrating(result, sid, pid)
+    return _drain_resumed_process(
+        client, sid, pid, on_output, result, wait_deadline)
+
+
+def _host_lost_terminal(result, message: str):
+    """Shape a lost-host result as a non-success terminal (exit 143)."""
+    result.status = "host_terminated"
+    result.returncode = None
+    result.message = message
+    return result
+
+
+def _rebuild_restart_mode() -> str:
+    """FCLOUD_MIGRATE_RESTART: "auto" (default) re-executes a workload whose
+    process was lost with its host on a plain rebuild; "never" fails fast
+    with a distinct host_lost result instead (for non-idempotent scripts)."""
+    return os.environ.get("FCLOUD_MIGRATE_RESTART", "auto").strip().lower()
+
+
+def _maybe_respawn_after_rebuild(client: Client, sid: str, pid: str):
+    """After a ride reconnects, the process may simply not exist anymore: a
+    plain REBUILD restores the session's files, not its processes (a true
+    checkpoint-restore resumes them; the client used to print "reconnected"
+    and poll the corpse pid for its whole budget).
+    The durable index says "lost_host" for exactly this case. Returns
+    ``(new_process, refused)``: a re-executed process (auto mode), or
+    ``(None, True)`` when the record is lost but a restart is disabled or
+    impossible. ``(None, False)`` means the record is fine — keep polling.
+    """
+    try:
+        log = client.session_process_logs(sid, pid)
+    except Exception:
+        return None, False
+    if log.status != "lost_host":
+        return None, False
+    if log.session_status and log.session_status in _ENDED_SESSION_STATES:
+        return None, True
+    if _rebuild_restart_mode() == "never" or not log.command:
+        return None, True
+    proc = client.session_spawn(sid, log.command)
+    print(
+        f"\N{ANTICLOCKWISE OPEN CIRCLE ARROW} process {pid} was lost with "
+        f"its host — restarted from the beginning as {proc.process_id} "
+        f"(disable with FCLOUD_MIGRATE_RESTART=never)",
+        file=sys.stderr,
+    )
+    return proc, False
+
+
+def _drain_resumed_process(client: Client, sid: str, pid: str, on_output,
+                           result, wait_deadline: float | None = None):
+    """Stream a restored process to exit on its new host, surviving further
+    migrations, and return its real exit result. Live output goes to the
+    console sink (human mode) or is captured onto the result (e.g. --json).
+    ``wait_deadline`` bounds the drain like session.run's own timeout would."""
+    from fcloud.types import Process
+    proc = Process(process_id=pid, session_id=sid, _client=client)
+    prev = ""
+    # Host B replays the checkpointed backlog before resuming; let it advance
+    # past the checkpoint, then baseline so we emit only genuinely-new output.
+    time.sleep(3)
+    while True:
+        try:
+            proc.poll(lines=1000)
+        except Exception:
+            try:
+                on_migrating, on_reconnected = _migration_ui()
+                client.reconnect_session(
+                    sid, on_migrating=on_migrating,
+                    on_reconnected=on_reconnected,
+                    budget=_ride_budget(wait_deadline))
+            except Exception:
+                return _still_migrating(result, sid, pid)
+            # The session is back — but after a plain rebuild the process
+            # is not. Consult durable truth (the rebuilt host marks lost
+            # records lost_host) and re-execute or fail fast instead of
+            # polling a corpse pid for the rest of the budget.
+            respawned, refused = _maybe_respawn_after_rebuild(client, sid, pid)
+            if refused:
+                result.status = "host_lost"
+                result.returncode = None
+                result.message = (
+                    f"Process {pid} was lost with its host (the session was "
+                    f"rebuilt, not restored) and was not restarted. Re-run "
+                    f"the command, or leave FCLOUD_MIGRATE_RESTART=auto. "
+                    f"Captured output: fcloud logs {sid} {pid} --output all")
+                return result
+            if respawned is not None:
+                pid = respawned.process_id
+                proc = Process(process_id=pid, session_id=sid, _client=client)
+                result.process_id = pid
+            prev = ""
+            time.sleep(3)
+            continue
+        if proc.output != prev:
+            delta = _tail_follow_delta(prev, proc.output)
+            if delta and on_output is not None:
+                on_output(delta)
+            prev = proc.output
+        if proc.status != ProcessStatus.RUNNING:
+            break
+        if wait_deadline is not None and time.monotonic() >= wait_deadline:
+            # The resumed process is genuinely alive but the caller's --wait
+            # is spent: report it exactly like session.run's own timeout
+            # (status "running", exit 0), not as a failed ride.
+            result.status = "running"
+            result.returncode = None
+            result.message = (
+                f"Command is still running. Follow logs: "
+                f"fcloud logs {sid} {pid} --follow")
+            return result
+        time.sleep(1)
+    if on_output is None:
+        result.stdout = proc.output
+    result.status = "exited"
+    result.returncode = proc.exit_code
+    result.message = ""
+    return result
+
+
+def _follow_reconnect(client: Client, sid: str, detach, *,
+                      budget: float = _RESUME_BUDGET_S):
+    """Re-attach to a session after its host was pre-empted/migrated, for the
+    log-follow loop. Thin wrapper over ``Client.reconnect_session`` (the shared
+    library mechanic): drops the dead connection, prints the migration UX, waits
+    for the restore, and returns a fresh detach fn. Raises if the session does
+    not come back within ``budget`` seconds. The default budget matches the
+    run/exec migration ride (_RESUME_BUDGET_S = the dispatcher's restore-phase
+    deadline) — --follow must not give up on restores the ride would survive.
+    """
+    try:
+        detach()
+    except Exception:
+        pass
+    on_migrating, on_reconnected = _migration_ui()
+    client.reconnect_session(
+        sid, on_migrating=on_migrating, on_reconnected=on_reconnected,
+        budget=budget,
+    )
+    return lambda: client.detach_session(sid)
+
+
+def _tail_follow_delta(prev: str, cur: str) -> str:
+    """Return only the genuinely-new tail of ``cur`` given what we already
+    emitted (``prev``).
+
+    The host serves ``poll`` as a *sliding tail window* (the last N lines,
+    proc.go ReadOutput), not cumulative output. Once a process prints more
+    than N lines the window scrolls, so ``cur`` no longer starts with
+    ``prev`` — and a naive prefix check (``cur.startswith(prev)``) then
+    reprints the entire window on every poll. For a long-running stream
+    (e.g. a trainer printing one line/sec) that turns ``--follow`` into a
+    flickering, endlessly-reprinting block. Align on the scrolled overlap
+    instead: emit only what follows the longest suffix-of-prev that is also
+    a prefix-of-cur. Falls back to the whole window when there is no overlap
+    (the first poll, or a post-migration rebaseline).
+    """
+    if not prev:
+        return cur
+    if cur.startswith(prev):
+        return cur[len(prev):]
+    for k in range(min(len(prev), len(cur)), 0, -1):
+        if prev[-k:] == cur[:k]:
+            return cur[k:]
+    return cur

fcloud/cli/mount.py

@@ -0,0 +1,150 @@
+"""`fcloud mount`: browse a session's files locally via rclone (read-only)."""
+
+import sys
+
+from fcloud.cli.registry import cmd
+
+
+@cmd("mount", "fcloud mount <session-id> <mountpoint> [-- rclone-args...]",
+     "Mount a session's files locally via rclone (read-only)")
+def cmd_mount(args: list[str]) -> None:
+    if len(args) < 2:
+        print("Usage: fcloud mount <session-id> <mountpoint> [-- rclone-args...]")
+        print("\nBrowse a session's files in Finder/Explorer (read-only). Needs")
+        print("only rclone (brew install rclone) — it uses the OS NFS client, so")
+        print("NO macFUSE / kernel extension / recovery mode is required.")
+        print("For CLI/programmatic access use `fcloud ls` and `fcloud download`.")
+        sys.exit(1)
+    import shutil
+    import subprocess
+
+    if not shutil.which("rclone"):
+        print("Error: rclone not found. Install: https://rclone.org/install/",
+              file=sys.stderr)
+        sys.exit(2)
+
+    import os
+
+    import requests
+
+    from fcloud.config import resolve_api_key, resolve_url
+
+    sid, mountpoint = args[0], os.path.abspath(os.path.expanduser(args[1]))
+    extra = args[3:] if len(args) > 2 and args[2] == "--" else args[2:]
+    # Safety: fcloud mount only ever uses a path that does NOT exist yet. It
+    # creates the mountpoint and removes it on unmount, so it can only touch a
+    # directory it made itself — it will never shadow or delete an existing
+    # folder. So `fcloud mount ~/important` is refused outright, before any
+    # filesystem change. (os.path.lexists also catches a symlink target.)
+    if os.path.lexists(mountpoint):
+        print(f"Error: {mountpoint} already exists — refusing. fcloud mount "
+              f"creates its own mountpoint and deletes it on unmount, so it "
+              f"only ever touches a folder it made. Pick a path that does not "
+              f"exist yet, e.g. ~/fcloud-{sid}.", file=sys.stderr)
+        sys.exit(2)
+    try:
+        os.makedirs(mountpoint)  # no exist_ok: it must not pre-exist
+    except OSError as e:
+        print(f"Error: cannot create mountpoint {mountpoint}: {e}", file=sys.stderr)
+        sys.exit(1)
+    url = resolve_url(None)
+    key = resolve_api_key(None)
+    if not key:
+        print("Error: no API key. Run: fcloud set_token <key>", file=sys.stderr)
+        sys.exit(2)
+
+    # Mint a short-lived capability token, then build a HEADER-LESS mount
+    # URL: the token lives in the path. rclone needs no auth header, so it
+    # follows each file's redirect to S3 cleanly (an Authorization header
+    # would be re-sent across the redirect and S3 rejects dual auth). The
+    # dispatcher only serves listings + redirects — file bytes stream
+    # straight from S3 to this machine.
+    resp = requests.post(
+        f"{url}/v1/fs/{sid}/mount-url",
+        json={}, headers={"Authorization": f"Bearer {key}"}, timeout=30,
+    )
+    if resp.status_code == 404:
+        print("Not found: unknown session or not yours.", file=sys.stderr)
+        sys.exit(1)
+    if resp.status_code >= 400:
+        print(f"Error minting mount token: {resp.status_code}: {resp.text[:200]}",
+              file=sys.stderr)
+        sys.exit(1)
+    mount_url = url.rstrip("/") + resp.json()["path"]
+
+    import re
+
+    # Use rclone's NFS mount (the OS NFS client), not FUSE: it needs NO
+    # macFUSE / kernel extension / recovery-mode approval — just rclone.
+    #
+    # Define a cleanly-NAMED remote via env vars rather than the on-the-fly
+    # ":http:" remote. The anonymous remote is auto-named ":http{xxxx}",
+    # and that name leaks into the Finder volume label and the VFS cache
+    # path (~/Library/Caches/rclone/vfs/:http{xxxx}/...), where the ':' and
+    # '{}' break cache-file creation. A clean per-session name fixes the
+    # label, makes the cache work, and avoids cache collisions across mounts.
+    remote = "fcloudmnt" + re.sub(r"[^a-z0-9]", "", sid.lower())
+    env = dict(os.environ)
+    env[f"RCLONE_CONFIG_{remote.upper()}_TYPE"] = "http"
+    env[f"RCLONE_CONFIG_{remote.upper()}_URL"] = mount_url
+
+    # Name the volume after the mountpoint the user chose, so Finder's label
+    # matches the folder (~/fcloud -> "fcloud") instead of an internal name.
+    volname = os.path.basename(os.path.normpath(mountpoint)) or f"fcloud-{sid}"
+    # vfs-cache-mode=off: read-only browsing doesn't need a local file
+    # cache, and rclone's full-cache path errors trying to create nested
+    # cache dirs. With off, reads stream straight from S3 (Range-capable),
+    # which is exactly what we want — no cache, no errors.
+    rclone_cmd = [
+        "rclone", "nfsmount", f"{remote}:", mountpoint,
+        "--volname", volname,
+        "--read-only",
+        "--vfs-cache-mode", "off",
+        "--dir-cache-time", "10s",
+        *extra,
+    ]
+    print(f"Mounting session {sid} at {mountpoint} (read-only). "
+          f"Open it in Finder/Explorer; press Ctrl-C to unmount.")
+    print("(macOS may block Terminal from reading the mount — that's a "
+          "privacy/TCC quirk; browse in Finder, or use `fcloud ls` from the CLI.)")
+    import signal
+    import time
+
+    # Route SIGINT (Ctrl-C) and SIGTERM (kill / backgrounded runs) to the
+    # same unmount+cleanup path. Re-installing SIGINT matters because a
+    # shell that backgrounds the process sets SIGINT to "ignore".
+    def _request_unmount(signum, frame):
+        raise KeyboardInterrupt
+    try:
+        signal.signal(signal.SIGINT, _request_unmount)
+        signal.signal(signal.SIGTERM, _request_unmount)
+    except (ValueError, OSError):
+        pass  # not the main thread / unsupported — fall back to defaults
+
+    proc = subprocess.Popen(rclone_cmd, env=env)
+    rc = 0
+    try:
+        rc = proc.wait()
+    except KeyboardInterrupt:
+        # Forward Ctrl-C and wait for rclone to unmount + exit before we
+        # touch the mountpoint, so cleanup sees an empty (unmounted) dir.
+        try:
+            proc.send_signal(signal.SIGINT)
+        except Exception:
+            pass
+        try:
+            rc = proc.wait(timeout=20)
+        except Exception:
+            rc = 0
+    finally:
+        # We created this mountpoint (it did not exist — we refused above if
+        # it did), so remove it now that rclone has unmounted. Retry briefly
+        # for macOS's NFS-detach lag. os.rmdir only removes an EMPTY dir, a
+        # final belt-and-suspenders so nothing unexpected is ever deleted.
+        for _ in range(20):
+            try:
+                os.rmdir(mountpoint)
+                break
+            except OSError:
+                time.sleep(0.3)
+    sys.exit(rc)

fcloud/cli/output.py

@@ -0,0 +1,17 @@
+"""Process-wide CLI output mode.
+
+`_JSON_OUTPUT` is set once by `fcloud.cli.main` from the global `--json` flag
+and read by every command at call time (`output._JSON_OUTPUT`), so a single
+binding governs the whole invocation. Tests patch it here.
+"""
+
+import json as _json
+
+_JSON_OUTPUT = False
+
+
+def _out(data: dict) -> None:
+    """Print data as JSON if --json, otherwise ignored (callers
+    handle human output themselves)."""
+    if _JSON_OUTPUT:
+        print(_json.dumps(data))