0.1.8/hf-dev

#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""Turn a Hugging Face Job into a low-friction remote development target."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import selectors
import shlex
import shutil
import socket
import subprocess
import sys
import tarfile
import time
import uuid
from pathlib import Path
from typing import Any

import tomllib


VERSION = "0.1.8"
MIN_HF_VERSION = (1, 28, 0)
TERMINAL_STAGES = {"CANCELED", "COMPLETED", "DELETED", "ERROR"}
DEFAULTS: dict[str, Any] = {
    "image": "huggingface/transformers-all-latest-gpu",
    "flavor": "l40sx1",
    "timeout": "30m",
    "bootstrap": (
        "{python} -m pip install -q debugpy && (uv pip install -e . || {python} -m pip install -e .)"
    ),
    "python": "auto",
    "forward_hf_token": True,
    "cache_bucket": None,
}
MUTAGEN_IGNORES = (
    ".venv",
    "venv",
    "__pycache__",
    "*.py[cod]",
    ".pytest_cache",
    ".mypy_cache",
    ".ruff_cache",
    ".tox",
    ".nox",
    "*.egg-info",
    ".eggs",
    "node_modules",
    "build",
    "dist",
)

REMOTE_DEBUG_LAUNCHER = r"""
import fcntl
import json
import os
from pathlib import Path
import runpy
import signal
import sys
import time

pid_path = Path(sys.argv[1])
owner_token = sys.argv[2]
remote_root = str(Path(sys.argv[3]).resolve())
port = int(sys.argv[4])
kind, target = sys.argv[5:7]
target_args = sys.argv[7:]
lock_path = Path(str(pid_path) + ".lock")


def read_owner():
    try:
        data = json.loads(pid_path.read_text())
        return int(data["pid"]), str(data["token"])
    except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError):
        return None


def process_exists(pid):
    try:
        if Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z":
            return False
    except (OSError, IndexError):
        pass
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def owns_process(pid, token):
    try:
        arguments = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    except OSError:
        return False
    return token.encode() in arguments


def is_legacy_debugger(pid):
    if pid == os.getpid():
        return False
    try:
        cwd = str(Path(f"/proc/{pid}/cwd").resolve())
        command = Path(f"/proc/{pid}/cmdline").read_bytes()
    except OSError:
        return False
    return cwd == remote_root and b"HF_DEV_DEBUG_READY" in command and b"debugpy.wait_for_client" in command


def terminate(pid):
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    deadline = time.monotonic() + 5
    while process_exists(pid) and time.monotonic() < deadline:
        time.sleep(0.1)
    if process_exists(pid):
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    previous = read_owner()
    if previous and previous[0] != os.getpid() and owns_process(*previous):
        terminate(previous[0])
    for process_dir in Path("/proc").iterdir():
        if process_dir.name.isdigit() and is_legacy_debugger(int(process_dir.name)):
            terminate(int(process_dir.name))
    temporary = pid_path.with_name(f"{pid_path.name}.{os.getpid()}.tmp")
    temporary.write_text(json.dumps({"pid": os.getpid(), "token": owner_token}))
    os.replace(temporary, pid_path)
    fcntl.flock(lock, fcntl.LOCK_UN)

try:
    import debugpy

    debugpy.listen(("127.0.0.1", port))
    print("HF_DEV_DEBUG_READY", flush=True)
    debugpy.wait_for_client()
    sys.argv = [target, *target_args]
    if kind == "module":
        runpy.run_module(target, run_name="__main__", alter_sys=True)
    else:
        runpy.run_path(target, run_name="__main__")
finally:
    with lock_path.open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if read_owner() == (os.getpid(), owner_token):
            pid_path.unlink(missing_ok=True)
""".strip()

REMOTE_DEBUG_STOPPER = r"""
import fcntl
import json
import os
from pathlib import Path
import signal
import sys
import time

pid_path = Path(sys.argv[1])
remote_root = str(Path(sys.argv[2]).resolve())
expected_token = sys.argv[3] if len(sys.argv) > 3 else None
lock_path = Path(str(pid_path) + ".lock")


def read_owner():
    try:
        data = json.loads(pid_path.read_text())
        return int(data["pid"]), str(data["token"])
    except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError):
        return None


def process_exists(pid):
    try:
        if Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z":
            return False
    except (OSError, IndexError):
        pass
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def owns_process(pid, token):
    try:
        arguments = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    except OSError:
        return False
    return token.encode() in arguments


def is_legacy_debugger(pid):
    if pid == os.getpid():
        return False
    try:
        cwd = str(Path(f"/proc/{pid}/cwd").resolve())
        command = Path(f"/proc/{pid}/cmdline").read_bytes()
    except OSError:
        return False
    return cwd == remote_root and b"HF_DEV_DEBUG_READY" in command and b"debugpy.wait_for_client" in command


def terminate(pid):
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    deadline = time.monotonic() + 5
    while process_exists(pid) and time.monotonic() < deadline:
        time.sleep(0.1)
    if process_exists(pid):
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    owner = read_owner()
    if owner and (expected_token is None or owner[1] == expected_token):
        if owns_process(*owner):
            terminate(owner[0])
        if read_owner() == owner:
            pid_path.unlink(missing_ok=True)
    if expected_token is None:
        for process_dir in Path("/proc").iterdir():
            if process_dir.name.isdigit() and is_legacy_debugger(int(process_dir.name)):
                terminate(int(process_dir.name))
""".strip()

REMOTE_NOTEBOOK_LAUNCHER = r"""
import fcntl
import json
import os
from pathlib import Path
import secrets
import signal
import socket
import subprocess
import sys
import time
import urllib.request

pid_path = Path(sys.argv[1])
owner_token = sys.argv[2]
python = sys.argv[3]
remote_root = str(Path(sys.argv[4]).resolve())
lock_path = Path(str(pid_path) + ".lock")


def read_owner():
    try:
        data = json.loads(pid_path.read_text())
        return int(data["pid"]), str(data["token"]), int(data.get("child_pid", 0))
    except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError):
        return None


def process_exists(pid):
    if not pid:
        return False
    try:
        if Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z":
            return False
    except (OSError, IndexError):
        pass
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def owns_process(pid, token):
    try:
        arguments = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    except OSError:
        return False
    return token.encode() in arguments


def owns_jupyter(pid):
    try:
        cwd = str(Path(f"/proc/{pid}/cwd").resolve())
        command = Path(f"/proc/{pid}/cmdline").read_bytes()
    except OSError:
        return False
    return cwd == remote_root and b"jupyter_server" in command


def terminate(pid):
    if not process_exists(pid):
        return
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    deadline = time.monotonic() + 5
    while process_exists(pid) and time.monotonic() < deadline:
        time.sleep(0.1)
    if process_exists(pid):
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


def write_owner(child_pid=0):
    temporary = pid_path.with_name(f"{pid_path.name}.{os.getpid()}.tmp")
    temporary.write_text(json.dumps({"pid": os.getpid(), "token": owner_token, "child_pid": child_pid}))
    os.replace(temporary, pid_path)


pid_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    previous = read_owner()
    write_owner()
    fcntl.flock(lock, fcntl.LOCK_UN)

if previous:
    previous_pid, previous_token, previous_child = previous
    if previous_pid != os.getpid() and owns_process(previous_pid, previous_token):
        terminate(previous_pid)
    if owns_jupyter(previous_child):
        terminate(previous_child)

with socket.socket() as port_socket:
    port_socket.bind(("127.0.0.1", 0))
    port = port_socket.getsockname()[1]

auth_token = secrets.token_urlsafe(32)
environment = dict(os.environ)
environment["JUPYTER_TOKEN"] = auth_token
command = [
    python,
    "-m",
    "jupyter_server",
    "--no-browser",
    "--ServerApp.ip=127.0.0.1",
    f"--ServerApp.port={port}",
    "--ServerApp.port_retries=0",
    f"--ServerApp.root_dir={remote_root}",
    "--ServerApp.allow_remote_access=False",
    "--ServerApp.allow_root=True",
    "--ServerApp.log_level=WARNING",
]
child = subprocess.Popen(
    command,
    cwd=remote_root,
    env=environment,
    stdin=subprocess.DEVNULL,
    stdout=sys.stderr,
    stderr=sys.stderr,
)

with lock_path.open("a+") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    write_owner(child.pid)
    fcntl.flock(lock, fcntl.LOCK_UN)


def stop(_signum=None, _frame=None):
    raise KeyboardInterrupt


signal.signal(signal.SIGTERM, stop)
signal.signal(signal.SIGINT, stop)

try:
    deadline = time.monotonic() + 120
    status_url = f"http://127.0.0.1:{port}/api/status?token={auth_token}"
    while time.monotonic() < deadline:
        if child.poll() is not None:
            raise RuntimeError(f"Jupyter server exited during startup with code {child.returncode}")
        try:
            with urllib.request.urlopen(status_url, timeout=1):
                break
        except Exception:
            time.sleep(0.25)
    else:
        raise RuntimeError("Timed out waiting for Jupyter server startup")
    print(
        "HF_DEV_REMOTE_JUPYTER_READY "
        + json.dumps({"port": port, "token": auth_token, "kernel_name": "hf-dev"}),
        flush=True,
    )
    child.wait()
except KeyboardInterrupt:
    pass
finally:
    terminate(child.pid)
    with lock_path.open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        current = read_owner()
        if current and current[:2] == (os.getpid(), owner_token):
            pid_path.unlink(missing_ok=True)
""".strip()

REMOTE_NOTEBOOK_STOPPER = r"""
import fcntl
import json
import os
from pathlib import Path
import signal
import sys
import time

pid_path = Path(sys.argv[1])
remote_root = str(Path(sys.argv[2]).resolve())
expected_token = sys.argv[3] if len(sys.argv) > 3 else None
lock_path = Path(str(pid_path) + ".lock")


def read_owner():
    try:
        data = json.loads(pid_path.read_text())
        return int(data["pid"]), str(data["token"]), int(data.get("child_pid", 0))
    except (FileNotFoundError, KeyError, TypeError, ValueError, json.JSONDecodeError):
        return None


def process_exists(pid):
    if not pid:
        return False
    try:
        if Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z":
            return False
    except (OSError, IndexError):
        pass
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def owns_process(pid, token):
    try:
        arguments = Path(f"/proc/{pid}/cmdline").read_bytes().split(b"\0")
    except OSError:
        return False
    return token.encode() in arguments


def owns_jupyter(pid):
    try:
        cwd = str(Path(f"/proc/{pid}/cwd").resolve())
        command = Path(f"/proc/{pid}/cmdline").read_bytes()
    except OSError:
        return False
    return cwd == remote_root and b"jupyter_server" in command


def terminate(pid):
    if not process_exists(pid):
        return
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    deadline = time.monotonic() + 5
    while process_exists(pid) and time.monotonic() < deadline:
        time.sleep(0.1)
    if process_exists(pid):
        try:
            os.kill(pid, signal.SIGKILL)
        except ProcessLookupError:
            pass


lock_path.parent.mkdir(parents=True, exist_ok=True)
with lock_path.open("a+") as lock:
    fcntl.flock(lock, fcntl.LOCK_EX)
    owner = read_owner()
    if owner and (expected_token is not None and owner[1] != expected_token):
        owner = None
    fcntl.flock(lock, fcntl.LOCK_UN)

if owner:
    pid, token, child_pid = owner
    if owns_process(pid, token):
        terminate(pid)
    if owns_jupyter(child_pid):
        terminate(child_pid)
    with lock_path.open("a+") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        if read_owner() == owner:
            pid_path.unlink(missing_ok=True)
""".strip()


class HfDevError(RuntimeError):
    pass


def command_path(env_name: str, default: str) -> str:
    override = os.environ.get(env_name)
    if override:
        return override
    resolved = shutil.which(default)
    if resolved is None:
        raise HfDevError(f"Required command is not on PATH: {default}")
    return resolved


def run(
    command: list[str],
    *,
    cwd: Path | None = None,
    check: bool = True,
    capture: bool = False,
) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        cwd=cwd,
        check=check,
        text=True,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.PIPE if capture else None,
    )


def repo_root(start: str | Path | None = None) -> Path:
    current = Path(start or os.getcwd()).expanduser().resolve()
    result = run(
        ["git", "-C", str(current), "rev-parse", "--show-toplevel"],
        check=False,
        capture=True,
    )
    if result.returncode == 0:
        return Path(result.stdout.strip()).resolve()
    return current


def repo_key(root: Path) -> str:
    return hashlib.sha256(str(root).encode()).hexdigest()[:10]


def safe_name(value: str) -> str:
    sanitized = re.sub(r"[^a-zA-Z0-9_-]", "-", value).strip("-")
    return sanitized[:40] or "workspace"


def duration_seconds(value: int | float | str) -> int:
    if isinstance(value, (int, float)):
        seconds = float(value)
    else:
        match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*([smhd]?)\s*", str(value))
        if match is None:
            raise HfDevError(
                f"Invalid timeout {value!r}; use seconds or a suffix such as 15m, 2h, or 1d"
            )
        amount, unit = match.groups()
        seconds = float(amount) * {"": 1, "s": 1, "m": 60, "h": 3600, "d": 86400}[unit]
    if seconds <= 0:
        raise HfDevError("The Job timeout must be greater than zero")
    return max(1, int(seconds))


def state_path(root: Path) -> Path:
    state_dir = Path(
        os.environ.get("HF_DEV_STATE_DIR", Path.home() / ".local" / "state" / "hf-dev")
    ).expanduser()
    return state_dir / f"{repo_key(root)}.json"


def load_state(root: Path, *, required: bool = True) -> dict[str, Any] | None:
    path = state_path(root)
    if not path.exists():
        if required:
            raise HfDevError(
                f"No HF development Job is recorded for {root}. Run `hf-dev up` first."
            )
        return None
    try:
        return json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as error:
        raise HfDevError(f"Cannot read state file {path}: {error}") from error


def save_state(root: Path, state: dict[str, Any]) -> None:
    path = state_path(root)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(".tmp")
    temporary.write_text(json.dumps(state, indent=2, sort_keys=True) + "\n")
    temporary.chmod(0o600)
    temporary.replace(path)


def remove_state(root: Path) -> None:
    path = state_path(root)
    if path.exists():
        path.unlink()


def load_toml(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        data = tomllib.loads(path.read_text())
    except (OSError, tomllib.TOMLDecodeError) as error:
        raise HfDevError(f"Cannot read configuration {path}: {error}") from error
    section = data.get("hf-dev", data)
    if not isinstance(section, dict):
        raise HfDevError(f"Expected a table in {path}")
    return section


def configuration(root: Path, args: argparse.Namespace) -> dict[str, Any]:
    config = dict(DEFAULTS)
    user_config = Path(
        os.environ.get(
            "HF_DEV_CONFIG", Path.home() / ".config" / "hf-dev" / "config.toml"
        )
    ).expanduser()
    config.update(load_toml(user_config))
    config.update(load_toml(root / ".hf-dev.toml"))

    for key in (
        "image",
        "flavor",
        "timeout",
        "namespace",
        "resource_group_id",
        "identity_file",
        "bootstrap",
        "python",
        "cache_bucket",
    ):
        value = getattr(args, key, None)
        if value is not None:
            config[key] = value
    if getattr(args, "no_forward_hf_token", False):
        config["forward_hf_token"] = False
    if getattr(args, "no_cache", False):
        config["cache_bucket"] = None
    return config


def hf_command() -> str:
    return command_path("HF_DEV_HF_BIN", "hf")


def mutagen_command() -> str:
    return command_path("HF_DEV_MUTAGEN_BIN", "mutagen")


def hf_version() -> tuple[int, ...]:
    result = run([hf_command(), "version"], capture=True)
    match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout)
    if match is None:
        raise HfDevError(f"Could not parse `hf version`: {result.stdout.strip()}")
    return tuple(int(part) for part in match.groups())


def require_current_hf() -> None:
    version = hf_version()
    if version < MIN_HF_VERSION:
        expected = ".".join(str(part) for part in MIN_HF_VERSION)
        actual = ".".join(str(part) for part in version)
        raise HfDevError(
            f"hf >= {expected} is required for Jobs SSH support; found {actual}"
        )


def current_namespace() -> str:
    result = run([hf_command(), "auth", "whoami", "--format", "json"], capture=True)
    data = json.loads(result.stdout)
    namespace = data.get("user") or data.get("name")
    if not namespace:
        raise HfDevError("Could not determine the current Hugging Face namespace")
    return str(namespace)


def prepare_cache(config: dict[str, Any], *, create: bool) -> None:
    configured = config.get("cache_bucket")
    if not configured:
        return
    bucket = str(configured)
    if bucket == "auto":
        bucket = f"{current_namespace()}/hf-dev-cache"
    elif "/" not in bucket:
        bucket = f"{current_namespace()}/{bucket}"
    if create:
        result = run(
            [
                hf_command(),
                "buckets",
                "create",
                bucket,
                "--private",
                "--exist-ok",
                "--format",
                "json",
            ],
            check=False,
            capture=True,
        )
        if result.returncode != 0:
            details = result.stderr.strip() or result.stdout.strip()
            raise HfDevError(
                f"Could not prepare persistent cache bucket {bucket}: {details}"
            )
    config["cache_bucket"] = bucket


def job_ref(state: dict[str, Any]) -> str:
    namespace = state.get("namespace")
    return f"{namespace}/{state['job_id']}" if namespace else state["job_id"]


def inspect_job(state: dict[str, Any]) -> dict[str, Any]:
    command = [hf_command(), "jobs", "inspect", job_ref(state), "--format", "json"]
    result = run(command, check=False, capture=True)
    if result.returncode != 0:
        raise HfDevError(
            result.stderr.strip() or result.stdout.strip() or "Could not inspect Job"
        )
    payload = json.loads(result.stdout)
    if isinstance(payload, list):
        if not payload:
            raise HfDevError(f"Job not found: {job_ref(state)}")
        return payload[0]
    return payload


def job_is_active(state: dict[str, Any]) -> bool:
    stage = str(inspect_job(state).get("status", {}).get("stage", "")).upper()
    return bool(stage) and stage not in TERMINAL_STAGES


def identity_args(config: dict[str, Any]) -> list[str]:
    configured = config.get("identity_file")
    if configured:
        path = Path(str(configured)).expanduser().resolve()
        if not path.is_file():
            raise HfDevError(f"SSH identity file does not exist: {path}")
        return ["-i", str(path), "-o", "IdentitiesOnly=yes"]
    return []


def ssh_base(state: dict[str, Any], config: dict[str, Any]) -> list[str]:
    return [
        command_path("HF_DEV_SSH_BIN", "ssh"),
        "-o",
        "StrictHostKeyChecking=accept-new",
        "-o",
        "ServerAliveInterval=30",
        "-o",
        "ServerAliveCountMax=3",
        *identity_args(config),
        f"{state['job_id']}@ssh.hf.jobs",
    ]


def ssh_exec(
    state: dict[str, Any],
    config: dict[str, Any],
    remote_command: str,
    *,
    check: bool = True,
    capture: bool = False,
) -> subprocess.CompletedProcess[str]:
    return run(ssh_base(state, config) + [remote_command], check=check, capture=capture)


def remote_debug_pid_path(root: Path) -> str:
    return f"/tmp/hf-dev-debug-{repo_key(root)}.json"


def remote_notebook_pid_path(root: Path) -> str:
    return f"/tmp/hf-dev-jupyter-{repo_key(root)}.json"


def stop_remote_debug(
    root: Path,
    state: dict[str, Any],
    config: dict[str, Any],
    *,
    owner_token: str | None = None,
) -> None:
    python = str(state.get("python") or config.get("python", "python3"))
    command = [
        python,
        "-c",
        REMOTE_DEBUG_STOPPER,
        remote_debug_pid_path(root),
        str(state["remote_root"]),
    ]
    if owner_token is not None:
        command.append(owner_token)
    ssh_exec(state, config, shlex.join(command), check=False, capture=True)


def stop_remote_notebook(
    root: Path,
    state: dict[str, Any],
    config: dict[str, Any],
    *,
    owner_token: str | None = None,
) -> None:
    python = str(state.get("python") or config.get("python", "python3"))
    command = [
        python,
        "-c",
        REMOTE_NOTEBOOK_STOPPER,
        remote_notebook_pid_path(root),
        str(state["remote_root"]),
    ]
    if owner_token is not None:
        command.append(owner_token)
    ssh_exec(state, config, shlex.join(command), check=False, capture=True)


def mutagen_session_name(root: Path) -> str:
    return f"hf-dev-{repo_key(root)}"


def stop_mutagen_session(session: str) -> None:
    run([mutagen_command(), "sync", "terminate", session], check=False, capture=True)


def mutagen_version() -> str:
    result = run([mutagen_command(), "version"], capture=True)
    match = re.search(r"\b\d+\.\d+\.\d+(?:[-+][a-zA-Z0-9.-]+)?\b", result.stdout)
    if match is None:
        raise HfDevError(f"Could not parse `mutagen version`: {result.stdout.strip()}")
    return match.group(0)


def install_mutagen_agent(state: dict[str, Any], config: dict[str, Any]) -> None:
    platform_result = ssh_exec(state, config, "uname -s && uname -m", capture=True)
    platform_parts = platform_result.stdout.strip().lower().splitlines()
    if len(platform_parts) != 2:
        raise HfDevError(
            f"Could not identify the remote platform: {platform_result.stdout.strip()}"
        )
    system, machine = platform_parts
    machine = {
        "aarch64": "arm64",
        "armv7l": "arm",
        "i386": "386",
        "i686": "386",
        "x86_64": "amd64",
    }.get(machine, machine)
    agent_name = f"{system}_{machine}"

    executable = Path(mutagen_command()).resolve()
    configured_bundle = os.environ.get("HF_DEV_MUTAGEN_AGENTS")
    candidates = (
        [Path(configured_bundle).expanduser()]
        if configured_bundle
        else [
            executable.parent / "mutagen-agents.tar.gz",
            executable.parent.parent / "libexec" / "mutagen-agents.tar.gz",
        ]
    )
    bundle = next((candidate for candidate in candidates if candidate.is_file()), None)
    if bundle is None:
        searched = ", ".join(str(candidate) for candidate in candidates)
        raise HfDevError(
            f"Mutagen agent bundle not found (searched: {searched}). Install Mutagen using its official archive "
            "or Homebrew formula, or set HF_DEV_MUTAGEN_AGENTS."
        )
    try:
        with tarfile.open(bundle, mode="r:gz") as archive:
            member = archive.getmember(agent_name)
            extracted = archive.extractfile(member)
            if extracted is None:
                raise HfDevError(
                    f"Mutagen agent bundle has no executable payload for {agent_name}"
                )
            agent = extracted.read()
    except (KeyError, OSError, tarfile.TarError) as error:
        raise HfDevError(
            f"Could not load Mutagen agent {agent_name} from {bundle}: {error}"
        ) from error

    version = mutagen_version()
    target = f".mutagen/agents/{version}/mutagen-agent"
    install = f"mkdir -p .mutagen/agents/{version} && dd of={target} status=none && chmod 700 {target}"
    print(
        f"Installing Mutagen {version} agent for {agent_name} through SSH...",
        flush=True,
    )
    process = subprocess.Popen(
        ssh_base(state, config) + [install],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = process.communicate(agent)
    if process.returncode != 0:
        details = (
            stderr.decode(errors="replace").strip()
            or stdout.decode(errors="replace").strip()
        )
        raise HfDevError(
            details
            or f"Mutagen agent installation failed with code {process.returncode}"
        )


def start_mutagen_session(
    root: Path, state: dict[str, Any], config: dict[str, Any]
) -> None:
    session = mutagen_session_name(root)
    stop_mutagen_session(session)
    ssh_exec(state, config, f"mkdir -p {shlex.quote(state['remote_root'])}")
    # HF Jobs' SSH proxy doesn't support the scp subprocess used by Mutagen's
    # automatic agent installation, so seed the exact versioned agent over SSH.
    install_mutagen_agent(state, config)
    remote = f"{state['job_id']}@ssh.hf.jobs:{state['remote_root']}"
    command = [
        mutagen_command(),
        "sync",
        "create",
        "--name",
        session,
        "--label",
        "hf-dev=true",
        "--label",
        f"hf-dev-job={state['job_id']}",
        "--mode",
        "one-way-replica",
        "--ignore-vcs",
        "--symlink-mode",
        "posix-raw",
        "--watch-mode-beta",
        "no-watch",
        "--default-file-mode-beta",
        "0644",
        "--default-directory-mode-beta",
        "0755",
        "--compression",
        "zstandard",
    ]
    for pattern in MUTAGEN_IGNORES:
        command.extend(["--ignore", pattern])
    command.extend([str(root), remote])
    print(f"Starting Mutagen session {session}...", flush=True)
    result = run(command, check=False, capture=True)
    if result.stdout.strip():
        print(result.stdout.rstrip())
    if result.returncode != 0:
        raise HfDevError(
            result.stderr.strip()
            or f"Mutagen session creation failed with code {result.returncode}"
        )
    state["mutagen_session"] = session
    save_state(root, state)
    flush_mutagen(state)


def flush_mutagen(state: dict[str, Any], attempts: int = 5) -> bool:
    session = state.get("mutagen_session")
    if not session:
        return False
    error = ""
    for attempt in range(attempts):
        result = run(
            [mutagen_command(), "sync", "flush", str(session)],
            check=False,
            capture=True,
        )
        if result.returncode == 0:
            return True
        error = (
            result.stderr.strip()
            or f"Mutagen flush failed with code {result.returncode}"
        )
        if attempt + 1 < attempts:
            time.sleep(2)
    raise HfDevError(error)


def ensure_mutagen_sync(
    root: Path, state: dict[str, Any], config: dict[str, Any]
) -> None:
    try:
        if flush_mutagen(state):
            return
    except HfDevError as error:
        print(
            f"Mutagen synchronization unavailable ({error}); recreating the session...",
            flush=True,
        )
    start_mutagen_session(root, state, config)


def wait_for_ssh(
    state: dict[str, Any], config: dict[str, Any], timeout: int = 900
) -> None:
    print(f"Waiting for {job_ref(state)} to accept SSH connections...", flush=True)
    deadline = time.monotonic() + timeout
    last_error = ""
    while time.monotonic() < deadline:
        result = ssh_exec(state, config, "true", check=False, capture=True)
        if result.returncode == 0:
            return
        last_error = result.stderr.strip()
        try:
            if not job_is_active(state):
                details = inspect_job(state)
                raise HfDevError(
                    f"Job stopped before SSH became ready: {details.get('status')}"
                )
        except HfDevError as error:
            last_error = str(error)
        time.sleep(5)
    raise HfDevError(f"Timed out waiting for SSH. Last error: {last_error}")


def discover_remote_python(
    root: Path, state: dict[str, Any], config: dict[str, Any]
) -> str:
    configured = str(config.get("python", "auto")).strip()
    candidates = []
    if configured and configured != "auto":
        candidates.append(configured)
    else:
        if state.get("python"):
            candidates.append(str(state["python"]))
        candidates.extend(
            [
                "python3",
                "python",
                "/opt/conda/bin/python",
                "/opt/venv/bin/python",
                "/usr/local/bin/python3",
                "/usr/bin/python3",
            ]
        )
    candidates = list(dict.fromkeys(candidates))

    probes = []
    for candidate in candidates:
        quoted = shlex.quote(candidate)
        if "/" in candidate:
            probes.append(f"([ -x {quoted} ] && printf '%s\\n' {quoted})")
        else:
            probes.append(f"command -v {quoted}")
    result = ssh_exec(state, config, " || ".join(probes), check=False, capture=True)
    resolved = result.stdout.strip().splitlines()
    if result.returncode != 0 or not resolved:
        details = result.stderr.strip()
        searched = ", ".join(candidates)
        raise HfDevError(
            f"No remote Python interpreter found (tried: {searched}). {details}".rstrip()
        )

    python = resolved[0]
    config["python"] = python
    state["python"] = python
    save_state(root, state)
    print(f"Using remote Python interpreter: {python}", flush=True)
    return python


def tracked_files(root: Path) -> list[Path]:
    result = subprocess.run(
        [
            "git",
            "-C",
            str(root),
            "ls-files",
            "--cached",
            "--others",
            "--exclude-standard",
            "-z",
        ],
        check=False,
        stdout=subprocess.PIPE,
    )
    if result.returncode == 0:
        paths = [
            Path(item.decode(errors="surrogateescape"))
            for item in result.stdout.split(b"\0")
            if item
        ]
        return [
            path
            for path in paths
            if (root / path).is_file() or (root / path).is_symlink()
        ]

    excluded = {
        ".git",
        ".hg",
        ".svn",
        ".venv",
        "__pycache__",
        ".pytest_cache",
        ".mypy_cache",
    }
    return [
        path.relative_to(root)
        for path in root.rglob("*")
        if not any(part in excluded for part in path.relative_to(root).parts)
        and (path.is_file() or path.is_symlink())
    ]


def sync_full(root: Path, state: dict[str, Any], config: dict[str, Any]) -> None:
    remote_root = str(state["remote_root"])
    expected_prefix = f"/workspace/hf-dev/{repo_key(root)}-"
    if not remote_root.startswith(expected_prefix):
        raise HfDevError(f"Refusing to replace unexpected remote path: {remote_root}")

    files = tracked_files(root)
    print(
        f"Syncing {len(files)} tracked/unignored files to {job_ref(state)}...",
        flush=True,
    )
    prepare = (
        f"rm -rf -- {shlex.quote(remote_root)} && mkdir -p {shlex.quote(remote_root)} "
        f"&& tar -xf - -C {shlex.quote(remote_root)}"
    )
    process = subprocess.Popen(
        ssh_base(state, config) + [prepare], stdin=subprocess.PIPE
    )
    assert process.stdin is not None
    try:
        with tarfile.open(fileobj=process.stdin, mode="w|") as archive:
            for relative in files:
                archive.add(root / relative, arcname=str(relative), recursive=False)
    except BrokenPipeError as error:
        raise HfDevError("SSH connection closed while syncing the workspace") from error
    finally:
        process.stdin.close()
    return_code = process.wait()
    if return_code != 0:
        raise HfDevError(f"Workspace sync failed with exit code {return_code}")


def relative_file(root: Path, value: str | Path) -> tuple[Path, Path]:
    absolute = Path(value).expanduser()
    if not absolute.is_absolute():
        absolute = root / absolute
    absolute = absolute.resolve(strict=False)
    try:
        relative = absolute.relative_to(root)
    except ValueError as error:
        raise HfDevError(
            f"Path is outside the active repository: {absolute}"
        ) from error
    return absolute, relative


def sync_file(
    root: Path, state: dict[str, Any], config: dict[str, Any], value: str
) -> None:
    absolute, relative = relative_file(root, value)
    remote = f"{state['remote_root']}/{relative.as_posix()}"
    if not absolute.exists():
        ssh_exec(state, config, f"rm -f -- {shlex.quote(remote)}")
        return
    if not absolute.is_file() and not absolute.is_symlink():
        return
    ssh_exec(state, config, f"mkdir -p {shlex.quote(str(Path(remote).parent))}")
    scp = [
        command_path("HF_DEV_SCP_BIN", "scp"),
        "-q",
        "-o",
        "StrictHostKeyChecking=accept-new",
        *identity_args(config),
        str(absolute),
        f"{state['job_id']}@ssh.hf.jobs:{shlex.quote(remote)}",
    ]
    run(scp)


def launch_command(root: Path, config: dict[str, Any]) -> list[str]:
    name = f"hf-dev-{safe_name(root.name)}-{repo_key(root)[:6]}"
    command = [
        hf_command(),
        "jobs",
        "run",
        "--ssh",
        "--detach",
        "--name",
        name,
        "--label",
        "hf-dev=true",
        "--label",
        f"hf-dev-repo={safe_name(root.name)}",
        "--flavor",
        str(config["flavor"]),
        "--timeout",
        str(config["timeout"]),
    ]
    if config.get("namespace"):
        command.extend(["--namespace", str(config["namespace"])])
    if config.get("resource_group_id"):
        command.extend(["--resource-group-id", str(config["resource_group_id"])])
    if config.get("forward_hf_token", True):
        command.extend(["--secrets", "HF_TOKEN"])
    if config.get("cache_bucket"):
        cache_home = "/cache/huggingface"
        command.extend(
            [
                "--volume",
                f"hf://buckets/{config['cache_bucket']}:{cache_home}",
                "--env",
                f"HF_HOME={cache_home}",
                "--env",
                "HF_HUB_DISABLE_SYMLINKS=1",
            ]
        )
    command.extend([str(config["image"]), "sleep", "infinity"])
    return command


def create_job(root: Path, config: dict[str, Any]) -> dict[str, Any]:
    result = run(launch_command(root, config), check=False, capture=True)
    output = "\n".join(part for part in (result.stdout, result.stderr) if part)
    if output.strip():
        print(output.rstrip())
    if result.returncode != 0:
        raise HfDevError(f"Job launch failed with exit code {result.returncode}")
    matches = re.findall(r"\b[0-9a-f]{24}\b", output)
    if not matches:
        raise HfDevError(
            "The Job was launched, but its ID could not be parsed from the hf CLI output"
        )
    namespace = str(config.get("namespace") or current_namespace())
    remote_root = f"/workspace/hf-dev/{repo_key(root)}-{safe_name(root.name)}"
    state = {
        "job_id": matches[-1],
        "namespace": namespace,
        "repo_root": str(root),
        "remote_root": remote_root,
        "image": config["image"],
        "flavor": config["flavor"],
        "timeout": config["timeout"],
        "cache_bucket": config.get("cache_bucket"),
        "cache_home": "/cache/huggingface" if config.get("cache_bucket") else None,
        "created_at": time.time(),
    }
    save_state(root, state)
    start_cancel_watchdog(root, state)
    return state


def start_cancel_watchdog(root: Path, state: dict[str, Any]) -> None:
    deadline = time.time() + duration_seconds(state["timeout"])
    command = [
        sys.executable,
        str(Path(__file__).resolve()),
        "_watchdog",
        "--job-id",
        str(state["job_id"]),
        "--namespace",
        str(state["namespace"]),
        "--deadline",
        str(deadline),
    ]
    process = subprocess.Popen(
        command,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        start_new_session=True,
        close_fds=True,
    )
    state["cancel_deadline"] = deadline
    state["cancel_watchdog_pid"] = process.pid
    save_state(root, state)


def cmd_watchdog(args: argparse.Namespace) -> None:
    state = {"job_id": args.job_id, "namespace": args.namespace}
    remaining = args.deadline - time.time()
    if remaining > 0:
        time.sleep(remaining)

    # Network connectivity can disappear while a laptop sleeps. Keep retrying
    # until cancellation succeeds or the Job is already terminal.
    retry_deadline = time.time() + 86400
    while time.time() < retry_deadline:
        try:
            if not job_is_active(state):
                return
        except (HfDevError, json.JSONDecodeError):
            time.sleep(30)
            continue
        result = run(
            [hf_command(), "jobs", "cancel", job_ref(state)], check=False, capture=True
        )
        if result.returncode == 0:
            return
        time.sleep(30)


def bootstrap(state: dict[str, Any], config: dict[str, Any]) -> None:
    command = str(config.get("bootstrap", "")).strip()
    if not command:
        return
    python = str(state.get("python") or config.get("python", "python3"))
    command = command.replace("{python}", shlex.quote(python))
    remote = f"cd {shlex.quote(state['remote_root'])} && {command}"
    print("Bootstrapping the remote Python environment...", flush=True)
    ssh_exec(state, config, f"bash -lc {shlex.quote(remote)}")


def verify_remote_source(state: dict[str, Any], config: dict[str, Any]) -> None:
    expected = f"{state['remote_root']}/src/transformers"
    code = (
        "from pathlib import Path; import transformers; "
        f"expected=Path({expected!r}).resolve(); "
        "actual=Path(transformers.__file__).resolve().parent; "
        "print(f'transformers source: {actual}'); "
        "assert actual == expected, f'expected editable source {expected}, got {actual}'"
    )
    python = str(state.get("python") or config.get("python", "python3"))
    remote = (
        f"cd {shlex.quote(state['remote_root'])} && {shlex.join([python, '-c', code])}"
    )
    ssh_exec(state, config, f"bash -lc {shlex.quote(remote)}")


def ensure_remote_notebook_packages(
    state: dict[str, Any], config: dict[str, Any]
) -> None:
    python = str(state.get("python") or config.get("python", "python3"))
    probe = shlex.join([python, "-c", "import ipykernel, jupyter_server"])
    result = ssh_exec(state, config, probe, check=False, capture=True)
    if result.returncode != 0:
        print("Installing the remote notebook runtime...", flush=True)
        install = (
            f"uv pip install --python {shlex.quote(python)} jupyter-server ipykernel "
            f"|| {shlex.join([python, '-m', 'pip', 'install', '-q', 'jupyter-server', 'ipykernel'])}"
        )
        ssh_exec(state, config, f"bash -lc {shlex.quote(install)}")
    kernel = [
        python,
        "-m",
        "ipykernel",
        "install",
        "--user",
        "--name",
        "hf-dev",
        "--display-name",
        "HF Dev GPU",
    ]
    ssh_exec(state, config, shlex.join(kernel), capture=True)


def wait_for_remote_notebook_marker(
    process: subprocess.Popen[str],
    *,
    timeout: int = 150,
) -> dict[str, Any]:
    assert process.stdout is not None
    selector = selectors.DefaultSelector()
    selector.register(process.stdout, selectors.EVENT_READ)
    deadline = time.monotonic() + timeout
    try:
        while time.monotonic() < deadline:
            if process.poll() is not None:
                raise HfDevError(
                    f"Remote Jupyter process exited during startup with code {process.returncode}"
                )
            if not selector.select(timeout=0.5):
                continue
            line = process.stdout.readline()
            if not line:
                continue
            if line.startswith("HF_DEV_REMOTE_JUPYTER_READY "):
                try:
                    return json.loads(line.removeprefix("HF_DEV_REMOTE_JUPYTER_READY "))
                except json.JSONDecodeError as error:
                    raise HfDevError(
                        "Remote Jupyter returned an invalid readiness marker"
                    ) from error
            print(line, end="", flush=True)
    finally:
        selector.close()
    raise HfDevError("Timed out waiting for the remote Jupyter server")


def wait_for_local_forward(
    process: subprocess.Popen[Any], port: int, *, timeout: int = 30
) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if process.poll() is not None:
            raise HfDevError(
                f"Jupyter SSH tunnel exited during startup with code {process.returncode}"
            )
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=0.5):
                return
        except OSError:
            time.sleep(0.1)
    raise HfDevError("Timed out waiting for the local Jupyter SSH tunnel")


def resolve_state(root: Path, config: dict[str, Any]) -> dict[str, Any]:
    state = load_state(root)
    assert state is not None
    if not job_is_active(state):
        details = inspect_job(state)
        raise HfDevError(
            f"The recorded Job is no longer active: {details.get('status')}"
        )
    return state


def cmd_up(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    require_current_hf()
    created = False
    state = load_state(root, required=False)
    if state is not None:
        if job_is_active(state):
            print(f"Reusing active Job {job_ref(state)}")
        else:
            state = None
    if state is None:
        prepare_cache(config, create=True)
        state = create_job(root, config)
        created = True
    try:
        wait_for_ssh(state, config)
        discover_remote_python(root, state, config)
        start_mutagen_session(root, state, config)
        if not args.no_bootstrap:
            bootstrap(state, config)
            verify_remote_source(state, config)
    except BaseException:
        if created:
            if state.get("mutagen_session"):
                stop_mutagen_session(str(state["mutagen_session"]))
            print(
                f"Setup failed; canceling newly created Job {job_ref(state)}...",
                file=sys.stderr,
                flush=True,
            )
            result = run(
                [hf_command(), "jobs", "cancel", job_ref(state)],
                check=False,
                capture=True,
            )
            if result.returncode == 0:
                remove_state(root)
            else:
                details = result.stderr.strip() or result.stdout.strip()
                print(
                    f"WARNING: automatic Job cancellation failed: {details}",
                    file=sys.stderr,
                    flush=True,
                )
        raise
    print(json.dumps(state, indent=2, sort_keys=True))


def cmd_plan(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    prepare_cache(config, create=False)
    print(shlex.join(launch_command(root, config)))
    print(f"local workspace:  {root}")
    print(
        f"remote workspace: /workspace/hf-dev/{repo_key(root)}-{safe_name(root.name)}"
    )


def cmd_info(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    state = load_state(root, required=False)
    if state is None:
        payload: dict[str, Any] = {"active": False, "repo_root": str(root)}
    else:
        payload = dict(state)
        if args.remote:
            try:
                details = inspect_job(state)
                payload["status"] = details.get("status")
                stage = str(details.get("status", {}).get("stage", "")).upper()
                payload["active"] = bool(stage) and stage not in TERMINAL_STAGES
            except HfDevError as error:
                payload["active"] = False
                payload["error"] = str(error)
        else:
            payload["active"] = True
    print(json.dumps(payload, indent=2, sort_keys=True))


def cmd_status(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    state = load_state(root)
    assert state is not None
    print(json.dumps(inspect_job(state), indent=2, sort_keys=True))


def cmd_sync(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    ensure_mutagen_sync(root, state, config)


def cmd_sync_file(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = load_state(root)
    assert state is not None
    if state.get("mutagen_session"):
        ensure_mutagen_sync(root, state, config)
    else:
        sync_file(root, state, config, args.path)


def cmd_shell(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    remote = f"cd {shlex.quote(state['remote_root'])} && exec bash -l"
    os.execvp(ssh_base(state, config)[0], ssh_base(state, config) + [remote])


def cmd_run(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    if not args.command:
        raise HfDevError("A remote command is required")
    command = args.command[0] if len(args.command) == 1 else shlex.join(args.command)
    remote = f"cd {shlex.quote(state['remote_root'])} && exec bash -lc {shlex.quote(command)}"
    result = ssh_exec(state, config, remote, check=False)
    raise SystemExit(result.returncode)


def cmd_debug(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    discover_remote_python(root, state, config)
    ensure_mutagen_sync(root, state, config)
    verify_remote_source(state, config)
    if not args.target:
        raise HfDevError("A script or module target is required")
    python = str(state.get("python") or config.get("python", "python3"))
    if args.module:
        kind = "module"
        target = args.module
        target_args = args.target
    else:
        kind = "script"
        target = args.target[0]
        target_args = args.target[1:]
    owner_token = uuid.uuid4().hex
    debug = [
        python,
        "-c",
        REMOTE_DEBUG_LAUNCHER,
        remote_debug_pid_path(root),
        owner_token,
        str(state["remote_root"]),
        str(args.port),
        kind,
        target,
        *target_args,
    ]
    remote = f"cd {shlex.quote(state['remote_root'])} && exec {shlex.join(debug)}"
    ssh = ssh_base(state, config)
    ssh[1:1] = [
        "-o",
        "ExitOnForwardFailure=yes",
        "-L",
        f"127.0.0.1:{args.port}:127.0.0.1:{args.port}",
    ]
    process = subprocess.Popen(ssh + [remote])
    try:
        returncode = process.wait()
    except KeyboardInterrupt:
        returncode = 130
    finally:
        stop_remote_debug(root, state, config, owner_token=owner_token)
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
    raise SystemExit(returncode)


def cmd_debug_stop(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    discover_remote_python(root, state, config)
    stop_remote_debug(root, state, config)
    print("Remote HF Dev debugger stopped.")


def cmd_notebook(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    discover_remote_python(root, state, config)
    ensure_mutagen_sync(root, state, config)
    verify_remote_source(state, config)
    ensure_remote_notebook_packages(state, config)

    if not 1 <= args.port <= 65535:
        raise HfDevError("The local notebook port must be between 1 and 65535")

    python = str(state.get("python") or config.get("python", "python3"))
    owner_token = uuid.uuid4().hex
    launch = [
        python,
        "-c",
        REMOTE_NOTEBOOK_LAUNCHER,
        remote_notebook_pid_path(root),
        owner_token,
        python,
        str(state["remote_root"]),
    ]
    remote = f"cd {shlex.quote(state['remote_root'])} && exec {shlex.join(launch)}"
    server = subprocess.Popen(
        ssh_base(state, config) + [remote],
        stdout=subprocess.PIPE,
        text=True,
    )
    tunnel: subprocess.Popen[Any] | None = None
    returncode = 0
    try:
        readiness = wait_for_remote_notebook_marker(server)
        remote_port = int(readiness["port"])
        auth_token = str(readiness["token"])
        kernel_name = str(readiness.get("kernel_name") or "hf-dev")

        tunnel_command = ssh_base(state, config)
        tunnel_command[1:1] = [
            "-o",
            "ExitOnForwardFailure=yes",
            "-N",
            "-L",
            f"127.0.0.1:{args.port}:127.0.0.1:{remote_port}",
        ]
        tunnel = subprocess.Popen(tunnel_command, stdin=subprocess.DEVNULL)
        wait_for_local_forward(tunnel, args.port)
        print(
            "HF_DEV_JUPYTER_READY "
            + json.dumps(
                {
                    "base_url": f"http://127.0.0.1:{args.port}",
                    "token": auth_token,
                    "kernel_name": kernel_name,
                }
            ),
            flush=True,
        )

        while server.poll() is None and tunnel.poll() is None:
            time.sleep(0.25)
        returncode = (
            server.returncode if server.poll() is not None else tunnel.returncode
        )
        returncode = returncode or 0
    except KeyboardInterrupt:
        returncode = 130
    except (KeyError, TypeError, ValueError) as error:
        raise HfDevError(
            "Remote Jupyter returned incomplete readiness information"
        ) from error
    finally:
        if tunnel is not None and tunnel.poll() is None:
            tunnel.terminate()
            try:
                tunnel.wait(timeout=5)
            except subprocess.TimeoutExpired:
                tunnel.kill()
                tunnel.wait()
        stop_remote_notebook(root, state, config, owner_token=owner_token)
        if server.poll() is None:
            server.terminate()
            try:
                server.wait(timeout=5)
            except subprocess.TimeoutExpired:
                server.kill()
                server.wait()
    raise SystemExit(returncode)


def cmd_notebook_stop(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    config = configuration(root, args)
    state = resolve_state(root, config)
    stop_remote_notebook(root, state, config)
    print("Remote HF Dev notebook server stopped.")


def cmd_logs(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    state = load_state(root)
    assert state is not None
    command = [hf_command(), "jobs", "logs"]
    if args.follow:
        command.append("--follow")
    command.append(job_ref(state))
    result = run(command, check=False)
    raise SystemExit(result.returncode)


def cmd_down(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    state = load_state(root)
    assert state is not None
    config = configuration(root, args)
    try:
        active = job_is_active(state)
    except HfDevError:
        active = True
    if active:
        stop_remote_notebook(root, state, config)
        stop_remote_debug(root, state, config)
    if state.get("mutagen_session"):
        stop_mutagen_session(str(state["mutagen_session"]))
    result = run([hf_command(), "jobs", "cancel", job_ref(state)], check=False)
    if result.returncode == 0:
        remove_state(root)
    raise SystemExit(result.returncode)


def cmd_doctor(args: argparse.Namespace) -> None:
    root = repo_root(args.repo)
    checks = {
        "repo_root": str(root),
        "hf": hf_command(),
        "hf_version": ".".join(str(part) for part in hf_version()),
        "ssh": command_path("HF_DEV_SSH_BIN", "ssh"),
        "scp": command_path("HF_DEV_SCP_BIN", "scp"),
        "mutagen": mutagen_command(),
        "state_file": str(state_path(root)),
    }
    checks["hf_version_ok"] = hf_version() >= MIN_HF_VERSION
    print(json.dumps(checks, indent=2, sort_keys=True))


def common_options(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--repo", help="Repository path (defaults to the current Git repository)"
    )
    parser.add_argument(
        "--identity-file", help="SSH private key registered on the Hugging Face Hub"
    )


def launch_options(parser: argparse.ArgumentParser) -> None:
    common_options(parser)
    parser.add_argument("--image")
    parser.add_argument("--flavor")
    parser.add_argument("--timeout")
    parser.add_argument("--namespace")
    parser.add_argument("--resource-group-id")
    parser.add_argument("--bootstrap")
    parser.add_argument("--python")
    parser.add_argument(
        "--cache-bucket", help="Opt-in private Bucket mounted as persistent HF_HOME"
    )
    parser.add_argument(
        "--no-cache",
        action="store_true",
        help="Use the default ephemeral Job-local Hub cache",
    )
    parser.add_argument("--no-forward-hf-token", action="store_true")


def parser() -> argparse.ArgumentParser:
    root = argparse.ArgumentParser(prog="hf-dev", description=__doc__)
    root.add_argument("--version", action="version", version=f"%(prog)s {VERSION}")
    commands = root.add_subparsers(dest="subcommand", required=True)

    up = commands.add_parser(
        "up", help="Create/reuse a Job, sync the repo, and prepare the environment"
    )
    launch_options(up)
    up.add_argument("--no-bootstrap", action="store_true")
    up.set_defaults(handler=cmd_up)

    plan = commands.add_parser(
        "plan", help="Print the Job launch plan without changing remote state"
    )
    launch_options(plan)
    plan.set_defaults(handler=cmd_plan)

    info = commands.add_parser(
        "info", help="Print locally recorded Job information as JSON"
    )
    common_options(info)
    info.add_argument(
        "--remote", action="store_true", help="Also query the remote Job status"
    )
    info.set_defaults(handler=cmd_info)

    status = commands.add_parser("status", help="Inspect the remote Job")
    common_options(status)
    status.set_defaults(handler=cmd_status)

    sync = commands.add_parser(
        "sync", help="Replace the remote workspace with local tracked/unignored files"
    )
    common_options(sync)
    sync.set_defaults(handler=cmd_sync)

    one_file = commands.add_parser(
        "sync-file", help="Sync one saved file to the remote workspace"
    )
    common_options(one_file)
    one_file.add_argument("path")
    one_file.set_defaults(handler=cmd_sync_file)

    shell = commands.add_parser(
        "shell", help="Open an SSH shell in the remote workspace"
    )
    common_options(shell)
    shell.set_defaults(handler=cmd_shell)

    remote_run = commands.add_parser(
        "run", help="Run a command in the remote workspace"
    )
    common_options(remote_run)
    remote_run.add_argument("command", nargs=argparse.REMAINDER)
    remote_run.set_defaults(handler=cmd_run)

    debug = commands.add_parser(
        "debug", help="Start debugpy remotely through an SSH tunnel"
    )
    common_options(debug)
    debug.add_argument("--port", type=int, default=5678)
    debug.add_argument("--module", help="Run a Python module, e.g. pytest")
    debug.add_argument("target", nargs=argparse.REMAINDER)
    debug.set_defaults(handler=cmd_debug)

    debug_stop = commands.add_parser(
        "debug-stop", help="Stop remote hf-dev debugger processes for this repository"
    )
    common_options(debug_stop)
    debug_stop.set_defaults(handler=cmd_debug_stop)

    notebook = commands.add_parser(
        "notebook", help="Start remote Jupyter through a local SSH tunnel"
    )
    common_options(notebook)
    notebook.add_argument(
        "--port",
        type=int,
        required=True,
        help="Local loopback port for the Jupyter tunnel",
    )
    notebook.set_defaults(handler=cmd_notebook)

    notebook_stop = commands.add_parser(
        "notebook-stop", help="Stop the remote Jupyter server for this repository"
    )
    common_options(notebook_stop)
    notebook_stop.set_defaults(handler=cmd_notebook_stop)

    logs = commands.add_parser("logs", help="Show Job logs")
    common_options(logs)
    logs.add_argument("--follow", action="store_true")
    logs.set_defaults(handler=cmd_logs)

    down = commands.add_parser("down", help="Cancel the Job and remove its local state")
    common_options(down)
    down.set_defaults(handler=cmd_down)

    doctor = commands.add_parser(
        "doctor", help="Check local prerequisites without contacting Jobs"
    )
    common_options(doctor)
    doctor.set_defaults(handler=cmd_doctor)

    watchdog = commands.add_parser("_watchdog", help=argparse.SUPPRESS)
    watchdog.add_argument("--job-id", required=True)
    watchdog.add_argument("--namespace", required=True)
    watchdog.add_argument("--deadline", required=True, type=float)
    watchdog.set_defaults(handler=cmd_watchdog)
    return root


def main() -> None:
    args = parser().parse_args()
    try:
        args.handler(args)
    except HfDevError as error:
        print(f"hf-dev: {error}", file=sys.stderr)
        raise SystemExit(2) from error
    except subprocess.CalledProcessError as error:
        print(
            f"hf-dev: command failed ({error.returncode}): {shlex.join(error.cmd)}",
            file=sys.stderr,
        )
        raise SystemExit(error.returncode) from error


if __name__ == "__main__":
    main()