fork mindXtrain from GitHub (Professor-Codephreak/mindXtrain@661bd41) as the mindX-specific line

dfb775d verified about 1 month ago

9.36 kB

"""Existing-droplet sync builder (rsync + ssh + scp argv arrays).
Pure builder. The orchestrator runs each step via
`mindxtrain.operator.runs.spawn_subprocess_streaming` so every command's
stdout streams back over the SSE pipeline.
Argv is always list-form — never `shell=True` — so user-supplied env values
can't be turned into shell injection. The remote shell snippets are
single-string command bodies (because `ssh` itself joins them on the wire),
but the *arguments* to ssh are still argv elements.
"""

from __future__ import annotations

import os
import shutil
from dataclasses import dataclass
from pathlib import Path

from mindxtrain.deploy.github_push import Step

DEFAULT_REMOTE_PATH = "/workspace/mindxtrain"
DEFAULT_CONTAINER = "rocm/primus:v26.2"
DEFAULT_SSH_KEY = "~/.ssh/id_ed25519"

# Common ssh hardening flags applied everywhere:
#  - BatchMode=yes        : never prompt for a password (would deadlock)
#  - StrictHostKeyChecking=accept-new : trust on first use, refuse changes
#  - ServerAliveInterval=30 : keep the socket open during long bench runs
_SSH_OPTS = (
    "-o", "BatchMode=yes",
    "-o", "StrictHostKeyChecking=accept-new",
    "-o", "ServerAliveInterval=30",
)

# rsync exclude rules. .git intentionally excluded — the cloud-init path
# clones from GitHub; this rsync path syncs working-tree state for fast
# iteration on a droplet that was provisioned manually.
_DEFAULT_EXCLUDES = (
    ".git",
    "__pycache__",
    ".venv",
    ".pytest_cache",
    ".mypy_cache",
    ".ruff_cache",
    "node_modules",
    "out/",
    "dist/",
    "build/",
    "*.egg-info",
)


@dataclass(frozen=True)
class DropletConfig:
    host: str
    user: str = "root"
    ssh_key: str = DEFAULT_SSH_KEY
    remote_path: str = DEFAULT_REMOTE_PATH
    container: str = DEFAULT_CONTAINER
    extras: str = "ml,eval,data,obs"


def required_env() -> tuple[str, ...]:
    """Env vars that must be set for /api/droplet/sync."""
    return ("DROPLET_HOST", "DROPLET_USER")


def missing_env(env: dict[str, str] | None = None) -> list[str]:
    src = env if env is not None else os.environ
    return [k for k in required_env() if not src.get(k)]


def _which_missing() -> list[str]:
    out = []
    for binary in ("rsync", "ssh", "scp"):
        if shutil.which(binary) is None:
            out.append(binary)
    return out


def status_missing(env: dict[str, str] | None = None) -> list[str]:
    return missing_env(env) + _which_missing()


def status_target(env: dict[str, str] | None = None) -> str:
    src = env if env is not None else os.environ
    user = src.get("DROPLET_USER", "")
    host = src.get("DROPLET_HOST", "")
    path = src.get("DROPLET_REMOTE_PATH", DEFAULT_REMOTE_PATH)
    if not (user and host):
        return ""
    return f"{user}@{host}:{path}"


def from_env(env: dict[str, str] | None = None) -> DropletConfig:
    src = env if env is not None else os.environ
    missing = missing_env(env)
    if missing:
        msg = f"droplet config missing env: {', '.join(missing)}"
        raise RuntimeError(msg)
    return DropletConfig(
        host=src["DROPLET_HOST"],
        user=src.get("DROPLET_USER", "root"),
        ssh_key=src.get("DROPLET_SSH_KEY", DEFAULT_SSH_KEY),
        remote_path=src.get("DROPLET_REMOTE_PATH", DEFAULT_REMOTE_PATH),
        container=src.get("DROPLET_CONTAINER", DEFAULT_CONTAINER),
    )


# ---- argv builders -------------------------------------------------------

def build_rsync(cfg: DropletConfig, repo_root: Path) -> list[str]:
    excludes: list[str] = []
    for pat in _DEFAULT_EXCLUDES:
        excludes += ["--exclude", pat]
    ssh_inline = (
        f"ssh -i {cfg.ssh_key} "
        f"-o BatchMode=yes "
        f"-o StrictHostKeyChecking=accept-new"
    )
    # Trailing slashes matter — copy contents of repo_root into remote_path/.
    return [
        "rsync", "-az", "--delete", "--info=progress2",
        *excludes,
        "-e", ssh_inline,
        f"{repo_root.rstrip('/') if isinstance(repo_root, str) else str(repo_root).rstrip('/')}/",
        f"{cfg.user}@{cfg.host}:{cfg.remote_path}/",
    ]


def build_provision_ssh(cfg: DropletConfig) -> list[str]:
    """Idempotent: install podman + pull the container image, skip if already done."""
    remote = (
        "command -v podman >/dev/null 2>&1 || "
        "(sudo apt-get update && sudo apt-get install -y podman); "
        f"podman image exists {cfg.container} || podman pull {cfg.container}"
    )
    return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


def build_bench_ssh(cfg: DropletConfig) -> list[str]:
    """Run `mindxtrain bench --gpu 0 --out plan.json` inside the container.
    `-tt` forces a remote pty so SIGINT from the local ssh client propagates
    through the SSH channel to `podman` and on to the GPU process. Without
    `-tt`, the existing `_REGISTRY.cancel()` SIGINT/SIGTERM in
    `mindxtrain.operator.runs:331` reaches only the local ssh, leaving the
    remote GPU still spinning.
    """
    remote = (
        f"cd {cfg.remote_path} && podman run --rm "
        f"--device /dev/kfd --device /dev/dri "
        f"-v {cfg.remote_path}:{cfg.remote_path} -w {cfg.remote_path} "
        f"{cfg.container} bash -lc "
        f"'pip install -e .[{cfg.extras}] && "
        f"mindxtrain bench --gpu 0 --out plan.json'"
    )
    return ["ssh", "-tt", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


def build_scp_plan_back(cfg: DropletConfig, local_dest: Path) -> list[str]:
    return [
        "scp", "-i", cfg.ssh_key, *_SSH_OPTS,
        f"{cfg.user}@{cfg.host}:{cfg.remote_path}/plan.json",
        str(local_dest),
    ]


def build_ssh_probe(cfg: DropletConfig) -> list[str]:
    """Quick liveness probe — `ssh user@host true`. Used to wait for the box
    to come up after provisioning."""
    return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, "-o", "ConnectTimeout=5",
            f"{cfg.user}@{cfg.host}", "true"]


def build_tail_cloud_init(cfg: DropletConfig) -> list[str]:
    """Tail cloud-init's combined output until the bootstrap sentinel exists.
    Streams the log live; exits 0 once the sentinel appears. The remote shell
    here is a small loop, run via ssh's argv (not /bin/sh -c locally).
    """
    from mindxtrain.deploy.cloud_init import BOOTSTRAP_SENTINEL, CLOUD_INIT_LOG

    remote = (
        f"touch {CLOUD_INIT_LOG}; "
        f"tail -n +1 -F {CLOUD_INIT_LOG} & TAIL_PID=$!; "
        f"while [ ! -f {BOOTSTRAP_SENTINEL} ]; do sleep 5; done; "
        f"sleep 2; kill $TAIL_PID 2>/dev/null; true"
    )
    return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


def build_tail_training_log(cfg: DropletConfig) -> list[str]:
    """Tail the droplet's `mindxtrain train` combined log until done sentinel.
    Streams every stdout line from the in-progress training run; exits with
    the captured training exit code once `.train-done` appears. Used by the
    orchestrator's post-bootstrap step to bridge remote training events into
    the operator's run registry (where Coach's SSE picks them up).
    The training command tees its output to a stable path
    (`out/runs/_combined_train.log`) so we don't have to guess the per-recipe
    run_name from the operator side.
    """
    from mindxtrain.deploy.cloud_init import (
        TRAIN_DONE_SENTINEL,
        TRAIN_EXIT_SENTINEL,
    )

    combined_log = f"{cfg.remote_path}/out/runs/_combined_train.log"
    remote = (
        # Make sure the log exists so tail -F doesn't error before training
        # writes its first line.
        f"mkdir -p {cfg.remote_path}/out/runs; "
        f"touch {combined_log}; "
        f"tail -n +1 -F {combined_log} & TAIL_PID=$!; "
        # Block until training is terminal.
        f"while [ ! -f {TRAIN_DONE_SENTINEL} ]; do sleep 5; done; "
        # Give tail one more pass to flush.
        f"sleep 2; kill $TAIL_PID 2>/dev/null; "
        # Propagate the training exit code so the orchestrator can decide
        # succeeded vs failed. Defaults to 1 if the sentinel is missing.
        f"exit $(cat {TRAIN_EXIT_SENTINEL} 2>/dev/null || echo 1)"
    )
    return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


# ---- pipeline assembly ---------------------------------------------------

def sync_steps(
    cfg: DropletConfig,
    repo_root: Path,
    *,
    run_bench: bool = True,
    fetch_plan: bool = True,
    plan_dest: Path | None = None,
) -> list[Step]:
    steps: list[Step] = [
        Step(label="rsync", cmd=build_rsync(cfg, repo_root)),
        Step(label="provision", cmd=build_provision_ssh(cfg)),
    ]
    if run_bench:
        steps.append(Step(label="bench", cmd=build_bench_ssh(cfg)))
    if fetch_plan and run_bench:
        dest = plan_dest or repo_root / "out" / "plan.remote.json"
        dest.parent.mkdir(parents=True, exist_ok=True)
        steps.append(Step(label="scp-plan", cmd=build_scp_plan_back(cfg, dest)))
    return steps


__all__ = [
    "DEFAULT_CONTAINER",
    "DEFAULT_REMOTE_PATH",
    "DEFAULT_SSH_KEY",
    "DropletConfig",
    "build_bench_ssh",
    "build_provision_ssh",
    "build_rsync",
    "build_scp_plan_back",
    "build_ssh_probe",
    "build_tail_cloud_init",
    "from_env",
    "missing_env",
    "required_env",
    "status_missing",
    "status_target",
    "sync_steps",
]

	"""Existing-droplet sync builder (rsync + ssh + scp argv arrays).

	Pure builder. The orchestrator runs each step via
	`mindxtrain.operator.runs.spawn_subprocess_streaming` so every command's
	stdout streams back over the SSE pipeline.

	Argv is always list-form — never `shell=True` — so user-supplied env values
	can't be turned into shell injection. The remote shell snippets are
	single-string command bodies (because `ssh` itself joins them on the wire),
	but the arguments to ssh are still argv elements.
	"""

	from __future__ import annotations

	import os
	import shutil
	from dataclasses import dataclass
	from pathlib import Path

	from mindxtrain.deploy.github_push import Step

	DEFAULT_REMOTE_PATH = "/workspace/mindxtrain"
	DEFAULT_CONTAINER = "rocm/primus:v26.2"
	DEFAULT_SSH_KEY = "~/.ssh/id_ed25519"

	# Common ssh hardening flags applied everywhere:
	# - BatchMode=yes : never prompt for a password (would deadlock)
	# - StrictHostKeyChecking=accept-new : trust on first use, refuse changes
	# - ServerAliveInterval=30 : keep the socket open during long bench runs
	_SSH_OPTS = (
	"-o", "BatchMode=yes",
	"-o", "StrictHostKeyChecking=accept-new",
	"-o", "ServerAliveInterval=30",
	)

	# rsync exclude rules. .git intentionally excluded — the cloud-init path
	# clones from GitHub; this rsync path syncs working-tree state for fast
	# iteration on a droplet that was provisioned manually.
	_DEFAULT_EXCLUDES = (
	".git",
	"__pycache__",
	".venv",
	".pytest_cache",
	".mypy_cache",
	".ruff_cache",
	"node_modules",
	"out/",
	"dist/",
	"build/",
	"*.egg-info",
	)


	@dataclass(frozen=True)
	class DropletConfig:
	host: str
	user: str = "root"
	ssh_key: str = DEFAULT_SSH_KEY
	remote_path: str = DEFAULT_REMOTE_PATH
	container: str = DEFAULT_CONTAINER
	extras: str = "ml,eval,data,obs"


	def required_env() -> tuple[str, ...]:
	"""Env vars that must be set for /api/droplet/sync."""
	return ("DROPLET_HOST", "DROPLET_USER")


	def missing_env(env: dict[str, str] \| None = None) -> list[str]:
	src = env if env is not None else os.environ
	return [k for k in required_env() if not src.get(k)]


	def _which_missing() -> list[str]:
	out = []
	for binary in ("rsync", "ssh", "scp"):
	if shutil.which(binary) is None:
	out.append(binary)
	return out


	def status_missing(env: dict[str, str] \| None = None) -> list[str]:
	return missing_env(env) + _which_missing()


	def status_target(env: dict[str, str] \| None = None) -> str:
	src = env if env is not None else os.environ
	user = src.get("DROPLET_USER", "")
	host = src.get("DROPLET_HOST", "")
	path = src.get("DROPLET_REMOTE_PATH", DEFAULT_REMOTE_PATH)
	if not (user and host):
	return ""
	return f"{user}@{host}:{path}"


	def from_env(env: dict[str, str] \| None = None) -> DropletConfig:
	src = env if env is not None else os.environ
	missing = missing_env(env)
	if missing:
	msg = f"droplet config missing env: {', '.join(missing)}"
	raise RuntimeError(msg)
	return DropletConfig(
	host=src["DROPLET_HOST"],
	user=src.get("DROPLET_USER", "root"),
	ssh_key=src.get("DROPLET_SSH_KEY", DEFAULT_SSH_KEY),
	remote_path=src.get("DROPLET_REMOTE_PATH", DEFAULT_REMOTE_PATH),
	container=src.get("DROPLET_CONTAINER", DEFAULT_CONTAINER),
	)


	# ---- argv builders -------------------------------------------------------

	def build_rsync(cfg: DropletConfig, repo_root: Path) -> list[str]:
	excludes: list[str] = []
	for pat in _DEFAULT_EXCLUDES:
	excludes += ["--exclude", pat]
	ssh_inline = (
	f"ssh -i {cfg.ssh_key} "
	f"-o BatchMode=yes "
	f"-o StrictHostKeyChecking=accept-new"
	)
	# Trailing slashes matter — copy contents of repo_root into remote_path/.
	return [
	"rsync", "-az", "--delete", "--info=progress2",
	*excludes,
	"-e", ssh_inline,
	f"{repo_root.rstrip('/') if isinstance(repo_root, str) else str(repo_root).rstrip('/')}/",
	f"{cfg.user}@{cfg.host}:{cfg.remote_path}/",
	]


	def build_provision_ssh(cfg: DropletConfig) -> list[str]:
	"""Idempotent: install podman + pull the container image, skip if already done."""
	remote = (
	"command -v podman >/dev/null 2>&1 \|\| "
	"(sudo apt-get update && sudo apt-get install -y podman); "
	f"podman image exists {cfg.container} \|\| podman pull {cfg.container}"
	)
	return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


	def build_bench_ssh(cfg: DropletConfig) -> list[str]:
	"""Run `mindxtrain bench --gpu 0 --out plan.json` inside the container.

	`-tt` forces a remote pty so SIGINT from the local ssh client propagates
	through the SSH channel to `podman` and on to the GPU process. Without
	`-tt`, the existing `_REGISTRY.cancel()` SIGINT/SIGTERM in
	`mindxtrain.operator.runs:331` reaches only the local ssh, leaving the
	remote GPU still spinning.
	"""
	remote = (
	f"cd {cfg.remote_path} && podman run --rm "
	f"--device /dev/kfd --device /dev/dri "
	f"-v {cfg.remote_path}:{cfg.remote_path} -w {cfg.remote_path} "
	f"{cfg.container} bash -lc "
	f"'pip install -e .[{cfg.extras}] && "
	f"mindxtrain bench --gpu 0 --out plan.json'"
	)
	return ["ssh", "-tt", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


	def build_scp_plan_back(cfg: DropletConfig, local_dest: Path) -> list[str]:
	return [
	"scp", "-i", cfg.ssh_key, *_SSH_OPTS,
	f"{cfg.user}@{cfg.host}:{cfg.remote_path}/plan.json",
	str(local_dest),
	]


	def build_ssh_probe(cfg: DropletConfig) -> list[str]:
	"""Quick liveness probe — `ssh user@host true`. Used to wait for the box
	to come up after provisioning."""
	return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, "-o", "ConnectTimeout=5",
	f"{cfg.user}@{cfg.host}", "true"]


	def build_tail_cloud_init(cfg: DropletConfig) -> list[str]:
	"""Tail cloud-init's combined output until the bootstrap sentinel exists.

	Streams the log live; exits 0 once the sentinel appears. The remote shell
	here is a small loop, run via ssh's argv (not /bin/sh -c locally).
	"""
	from mindxtrain.deploy.cloud_init import BOOTSTRAP_SENTINEL, CLOUD_INIT_LOG

	remote = (
	f"touch {CLOUD_INIT_LOG}; "
	f"tail -n +1 -F {CLOUD_INIT_LOG} & TAIL_PID=$!; "
	f"while [ ! -f {BOOTSTRAP_SENTINEL} ]; do sleep 5; done; "
	f"sleep 2; kill $TAIL_PID 2>/dev/null; true"
	)
	return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


	def build_tail_training_log(cfg: DropletConfig) -> list[str]:
	"""Tail the droplet's `mindxtrain train` combined log until done sentinel.

	Streams every stdout line from the in-progress training run; exits with
	the captured training exit code once `.train-done` appears. Used by the
	orchestrator's post-bootstrap step to bridge remote training events into
	the operator's run registry (where Coach's SSE picks them up).

	The training command tees its output to a stable path
	(`out/runs/_combined_train.log`) so we don't have to guess the per-recipe
	run_name from the operator side.
	"""
	from mindxtrain.deploy.cloud_init import (
	TRAIN_DONE_SENTINEL,
	TRAIN_EXIT_SENTINEL,
	)

	combined_log = f"{cfg.remote_path}/out/runs/_combined_train.log"
	remote = (
	# Make sure the log exists so tail -F doesn't error before training
	# writes its first line.
	f"mkdir -p {cfg.remote_path}/out/runs; "
	f"touch {combined_log}; "
	f"tail -n +1 -F {combined_log} & TAIL_PID=$!; "
	# Block until training is terminal.
	f"while [ ! -f {TRAIN_DONE_SENTINEL} ]; do sleep 5; done; "
	# Give tail one more pass to flush.
	f"sleep 2; kill $TAIL_PID 2>/dev/null; "
	# Propagate the training exit code so the orchestrator can decide
	# succeeded vs failed. Defaults to 1 if the sentinel is missing.
	f"exit $(cat {TRAIN_EXIT_SENTINEL} 2>/dev/null \|\| echo 1)"
	)
	return ["ssh", "-i", cfg.ssh_key, *_SSH_OPTS, f"{cfg.user}@{cfg.host}", remote]


	# ---- pipeline assembly ---------------------------------------------------

	def sync_steps(
	cfg: DropletConfig,
	repo_root: Path,
	*,
	run_bench: bool = True,
	fetch_plan: bool = True,
	plan_dest: Path \| None = None,
	) -> list[Step]:
	steps: list[Step] = [
	Step(label="rsync", cmd=build_rsync(cfg, repo_root)),
	Step(label="provision", cmd=build_provision_ssh(cfg)),
	]
	if run_bench:
	steps.append(Step(label="bench", cmd=build_bench_ssh(cfg)))
	if fetch_plan and run_bench:
	dest = plan_dest or repo_root / "out" / "plan.remote.json"
	dest.parent.mkdir(parents=True, exist_ok=True)
	steps.append(Step(label="scp-plan", cmd=build_scp_plan_back(cfg, dest)))
	return steps


	__all__ = [
	"DEFAULT_CONTAINER",
	"DEFAULT_REMOTE_PATH",
	"DEFAULT_SSH_KEY",
	"DropletConfig",
	"build_bench_ssh",
	"build_provision_ssh",
	"build_rsync",
	"build_scp_plan_back",
	"build_ssh_probe",
	"build_tail_cloud_init",
	"from_env",
	"missing_env",
	"required_env",
	"status_missing",
	"status_target",
	"sync_steps",
	]