diff --git a/frontend/src/components/training/config/TargetCard.tsx b/frontend/src/components/training/config/TargetCard.tsx index 9da59f0c..5e2131c2 100644 --- a/frontend/src/components/training/config/TargetCard.tsx +++ b/frontend/src/components/training/config/TargetCard.tsx @@ -1,6 +1,8 @@ import React from "react"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Label } from "@/components/ui/label"; +import { Input } from "@/components/ui/input"; +import { NumberInput } from "@/components/ui/number-input"; import { Select, SelectContent, @@ -8,9 +10,18 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; -import { ConfigComponentProps } from "../types"; +import { ConfigComponentProps, SshConnectionConfig } from "../types"; import { RunnerFlavor } from "@/lib/jobsApi"; +const DEFAULT_SSH: SshConnectionConfig = { + host: "", + port: 22, + username: "", + ssh_key_path: "", + remote_workdir: "", + remote_python_cmd: "python", +}; + interface TargetCardProps extends ConfigComponentProps { authenticated: boolean; flavors: RunnerFlavor[]; @@ -36,17 +47,30 @@ const TargetCard: React.FC = ({ }) => { const target = config.target; const value = - target.runner === "local" ? "local" : `hf:${target.flavor ?? ""}`; + target.runner === "local" + ? "local" + : target.runner === "ssh_remote" + ? "ssh" + : `hf:${target.flavor ?? ""}`; const handleChange = (v: string) => { if (v === "local") { updateConfig("target", { runner: "local" }); + } else if (v === "ssh") { + updateConfig("target", { runner: "ssh_remote", ssh: target.ssh ?? DEFAULT_SSH }); } else if (v.startsWith("hf:")) { const flavor = v.slice("hf:".length); updateConfig("target", { runner: "hf_cloud", flavor }); } }; + const updateSsh = (key: K, value: SshConnectionConfig[K]) => { + updateConfig("target", { + runner: "ssh_remote", + ssh: { ...(target.ssh ?? DEFAULT_SSH), [key]: value }, + }); + }; + return ( @@ -61,6 +85,7 @@ const TargetCard: React.FC = ({ Local — your machine (free) + Remote server (SSH) — your own machine {flavors.map((f) => ( = ({ account when training completes.

+ + {target.runner === "ssh_remote" && ( +
+

+ lelab makes no assumption about what's installed on this server — it + copies the dataset over via scp and runs the command below as-is. + Make sure lerobot is already set up there. +

+
+
+ + updateSsh("host", e.target.value)} + placeholder="gpu.example.com or 192.168.1.10" + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1" + /> +
+
+ + v !== undefined && updateSsh("port", v)} + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1" + /> +
+
+
+
+ + updateSsh("username", e.target.value)} + placeholder="noah" + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1" + /> +
+
+ + updateSsh("ssh_key_path", e.target.value)} + placeholder="Leave empty to use ssh-agent / default identity" + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1" + /> +
+
+
+ + updateSsh("remote_workdir", e.target.value)} + placeholder="/home/noah/lelab-runs" + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1" + /> +

+ Datasets and outputs are staged under here — nothing outside it is touched. +

+
+
+ + updateSsh("remote_python_cmd", e.target.value)} + placeholder="python, or e.g. source ~/venv/bin/activate && python" + className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1 font-mono text-sm" + /> +

+ Whatever it takes to reach an interpreter with lerobot installed — + lelab runs this verbatim, followed by the training command. +

+
+
+ )} ); diff --git a/frontend/src/components/training/types.ts b/frontend/src/components/training/types.ts index d0ee5f44..b2343bf9 100644 --- a/frontend/src/components/training/types.ts +++ b/frontend/src/components/training/types.ts @@ -1,5 +1,18 @@ +export interface SshConnectionConfig { + host: string; + port: number; + username: string; + ssh_key_path?: string; + remote_workdir: string; + remote_python_cmd: string; +} + export interface TrainingConfig { - target: { runner: "local" | "hf_cloud"; flavor?: string }; + target: { + runner: "local" | "hf_cloud" | "ssh_remote"; + flavor?: string; + ssh?: SshConnectionConfig; + }; // Dataset configuration dataset_repo_id: string; diff --git a/frontend/src/lib/jobsApi.ts b/frontend/src/lib/jobsApi.ts index 088b727b..6ad6709e 100644 --- a/frontend/src/lib/jobsApi.ts +++ b/frontend/src/lib/jobsApi.ts @@ -50,7 +50,18 @@ export interface TrainingRequest { optimizer_grad_clip_norm?: number; use_policy_training_preset: boolean; // Optional target for runner dispatch; omitted ⇒ local. - target?: { runner: "local" | "hf_cloud"; flavor?: string }; + target?: { + runner: "local" | "hf_cloud" | "ssh_remote"; + flavor?: string; + ssh?: { + host: string; + port: number; + username: string; + ssh_key_path?: string; + remote_workdir: string; + remote_python_cmd: string; + }; + }; } export interface JobRecord { diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx index a361d042..2cda04b1 100644 --- a/frontend/src/pages/Training.tsx +++ b/frontend/src/pages/Training.tsx @@ -275,6 +275,11 @@ const ConfigurationMode: React.FC = () => { const targetRequiresAuth = trainingConfig.target.runner === "hf_cloud"; const targetMissingFlavor = trainingConfig.target.runner === "hf_cloud" && !trainingConfig.target.flavor; + const targetMissingSshFields = + trainingConfig.target.runner === "ssh_remote" && + (!trainingConfig.target.ssh?.host.trim() || + !trainingConfig.target.ssh?.username.trim() || + !trainingConfig.target.ssh?.remote_workdir.trim()); const localBlocked = trainingConfig.target.runner === "local" && localJobRunning; const startDisabled = @@ -282,13 +287,16 @@ const ConfigurationMode: React.FC = () => { !trainingConfig.dataset_repo_id.trim() || localBlocked || (targetRequiresAuth && !authenticated) || - targetMissingFlavor; + targetMissingFlavor || + targetMissingSshFields; const startTooltip = localBlocked ? "Another local training is already running" : targetRequiresAuth && !authenticated ? "Log in to Hugging Face to use cloud compute" : targetMissingFlavor ? "Select a hardware flavor" + : targetMissingSshFields + ? "Fill in host, username, and remote working directory" : undefined; return ( diff --git a/lelab/jobs.py b/lelab/jobs.py index 10810a16..088b1d6b 100644 --- a/lelab/jobs.py +++ b/lelab/jobs.py @@ -46,12 +46,37 @@ JobState = Literal["running", "done", "failed", "interrupted"] +class SshConnectionConfig(BaseModel): + """Connection details for a user-owned remote training server. + + Deliberately makes no assumption about what's installed on the remote + host: `remote_python_cmd` is a free-form shell fragment (e.g. "python", + "/opt/venv/bin/python", "source ~/venv/bin/activate && python") rather + than a hardcoded interpreter path, since lelab has no way to know how + the user's server is set up. + """ + + host: str + port: int = 22 + username: str + # Path to a private key file. None ⇒ rely on the local ssh agent / + # default identity files (~/.ssh/id_rsa, etc.) — never a password: lelab + # doesn't handle interactive auth or store secrets. + ssh_key_path: str | None = None + # Absolute path on the remote host under which datasets/outputs are + # staged for this job (e.g. "/home/user/lelab-runs"). + remote_workdir: str + remote_python_cmd: str = "python" + + class JobTarget(BaseModel): """Where a job should run. `local` ⇒ LocalJobRunner. `hf_cloud` requires - a non-empty `flavor` from HfApi.list_jobs_hardware().""" + a non-empty `flavor` from HfApi.list_jobs_hardware(). `ssh_remote` + requires `ssh`.""" - runner: Literal["local", "hf_cloud"] = "local" + runner: Literal["local", "hf_cloud", "ssh_remote"] = "local" flavor: str | None = None + ssh: SshConnectionConfig | None = None class TrainingMetrics(BaseModel): @@ -79,15 +104,23 @@ class JobRecord(BaseModel): exit_code: int | None = None error_message: str | None = None metrics: TrainingMetrics = TrainingMetrics() - runner: Literal["local", "hf_cloud", "imported"] = "local" - # PID of the detached subprocess (local runner only); survives uvicorn - # --reload so a fresh registry can re-attach by tailing the log file. + runner: Literal["local", "hf_cloud", "ssh_remote", "imported"] = "local" + # PID of the detached subprocess (local + ssh_remote runners: for + # ssh_remote this is the local `ssh` client's pid, not a remote pid). + # Survives uvicorn --reload so a fresh registry can re-attach by tailing + # the log file. process_pid: int | None = None # HF Jobs identifiers (hf_cloud runner only) hf_job_id: str | None = None hf_flavor: str | None = None hf_repo_id: str | None = None hf_job_url: str | None = None + # ssh_remote runner only: connection used to start this job (so checkpoint + # pulls still work after a lelab restart) and the remote output directory + # the runner picked (outside record.output_dir, which stays a local path + # for _list_local_checkpoints to scan after a pull). + ssh_config: SshConnectionConfig | None = None + ssh_remote_dir: str | None = None # Captured from training stdout the first time wandb prints the run URL. wandb_run_url: str | None = None # Number of checkpoints currently visible (local: filesystem; cloud: @@ -523,6 +556,11 @@ def _list_local_checkpoints(output_dir: str) -> list[JobCheckpoint]: _CLOUD_CKPT_TTL_SECONDS = 30.0 + +# Min seconds between scp checkpoint pulls for a single ssh_remote job. A full +# recursive scp of the checkpoints/ tree isn't cheap, and the frontend polls +# checkpoints every ~5s — this keeps that from hammering the remote host. +_SSH_CHECKPOINT_PULL_INTERVAL_S = 30.0 _CKPT_PATH_RE = re.compile(r"^checkpoints/(\d+)/pretrained_model/config\.json$") @@ -672,6 +710,9 @@ def __init__(self, output_root: Path) -> None: # repo_id -> (expires_at_epoch, checkpoint list) self._cloud_ckpt_cache: dict[str, tuple[float, list[JobCheckpoint]]] = {} + # job_id -> last scp-pull epoch, for _maybe_pull_ssh_checkpoints' rate limit. + self._ssh_pull_cache: dict[str, float] = {} + # Fired (best-effort) on every state change: new job, stop initiated, # watchdog finalisation, delete. Server wires this to a WebSocket # broadcast so the frontend can refetch on-event instead of polling. @@ -797,10 +838,13 @@ def get(self, job_id: str) -> JobRecord: def start(self, config: TrainingRequest, target: JobTarget | None = None) -> JobRecord: from .runners.hf_cloud import HfCloudJobRunner # lazy import to avoid circular import + from .runners.ssh_remote import SshRemoteJobRunner # lazy import to avoid circular import target = target or JobTarget() if target.runner == "hf_cloud" and not target.flavor: raise ValueError("flavor is required when runner is hf_cloud") + if target.runner == "ssh_remote" and not target.ssh: + raise ValueError("ssh connection details are required when runner is ssh_remote") with self._lock: # Local trainings are bounded by this machine's GPU/USB resources, @@ -824,6 +868,7 @@ def start(self, config: TrainingRequest, target: JobTarget | None = None) -> Job started_at=time.time(), runner=target.runner, hf_flavor=target.flavor, + ssh_config=target.ssh, ) job_dir.mkdir(parents=True, exist_ok=True) @@ -833,6 +878,8 @@ def start(self, config: TrainingRequest, target: JobTarget | None = None) -> Job log_path = _job_log_path(self._output_root, job_id) if target.runner == "local": runner = LocalJobRunner(record.metrics, log_file_path=log_path) + elif target.runner == "ssh_remote": + runner = SshRemoteJobRunner(record.metrics, log_path, target.ssh) else: runner = HfCloudJobRunner(record.metrics, log_path, target.flavor) @@ -850,8 +897,10 @@ def start(self, config: TrainingRequest, target: JobTarget | None = None) -> Job # / page URL / model repo are printed by lerobot's submit_to_hf and # only appear in stdout a few seconds after start, so they're None # here; the watchdog (_tick) parses and persists them once they land. - if target.runner == "local": + if target.runner in ("local", "ssh_remote"): record.process_pid = runner.pid() + if target.runner == "ssh_remote": + record.ssh_remote_dir = runner.remote_output_dir() self._persist(record, force=True) self._runners[job_id] = runner @@ -1021,11 +1070,37 @@ def _checkpoints_for(self, record: JobRecord) -> builtins.list[JobCheckpoint]: return _list_imported_local(record.output_dir) if record.runner == "local": return _list_local_checkpoints(record.output_dir) + if record.runner == "ssh_remote": + # Pull the remote checkpoints/ tree into the same local output_dir + # a local job would use, then reuse the exact same local listing — + # no separate "remote" checkpoint representation needed. Rate + # limited so the ~5s poll from the frontend doesn't scp on every tick. + self._maybe_pull_ssh_checkpoints(record) + return _list_local_checkpoints(record.output_dir) # Cloud: _list_imported_hub prefers the checkpoints// tree (pushed when # save_checkpoint_to_hub is on) and falls back to the final model at the repo # root, so a finished run is always reachable even with no per-step tree. return self._list_cloud_cached(record.hf_repo_id) + def _maybe_pull_ssh_checkpoints(self, record: JobRecord) -> None: + """Best-effort scp pull of the remote checkpoints/ dir, rate-limited + per job. Silently no-ops on any failure (offline server, nothing + saved yet) — checkpoint listing degrades to "none yet" rather than + surfacing a transient network error to the whole jobs page.""" + if not record.ssh_config or not record.ssh_remote_dir: + return + now = time.time() + last = self._ssh_pull_cache.get(record.id, 0.0) + if now - last < _SSH_CHECKPOINT_PULL_INTERVAL_S: + return + self._ssh_pull_cache[record.id] = now + from .runners.ssh_remote import pull_checkpoints # lazy: avoid circular import + + try: + pull_checkpoints(record.ssh_config, record.ssh_remote_dir, record.output_dir) + except Exception as exc: + logger.info("Checkpoint pull skipped for ssh_remote job %s: %s", record.id, exc) + def list_checkpoints(self, job_id: str) -> builtins.list[JobCheckpoint]: """Return checkpoints saved for this job, ascending by step. diff --git a/lelab/runners/ssh_remote.py b/lelab/runners/ssh_remote.py new file mode 100644 index 00000000..a5842610 --- /dev/null +++ b/lelab/runners/ssh_remote.py @@ -0,0 +1,212 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SSH runner — runs a training on a server the user owns and controls +directly over SSH, rather than a HF-managed cloud container. + +Deliberately makes no assumption about what's already set up on the remote +host (lerobot install, venv, CUDA...) — `SshConnectionConfig.remote_python_cmd` +is a free-form shell fragment the user provides, and lelab just uses SSH/scp +as a plain transport: sync the dataset up, run the training command remotely, +tail its output, pull checkpoints back down on request. + +Known limitation: the dataset sync (`sync_dataset_up`) runs synchronously +before the training subprocess is spawned, so job start blocks the request +thread for as long as the scp takes — fine for a quick test dataset, but a +large video dataset could make the "start training" request hang for +minutes. Moving the sync into the background (state="staging" before +"running") is a reasonable follow-up once this has a real server to test +against. +""" + +from __future__ import annotations + +import logging +import shlex +import subprocess +from pathlib import Path + +from ..jobs import SshConnectionConfig, SubprocessJobRunner, TrainingMetrics +from ..train import TrainingRequest, build_training_command + +logger = logging.getLogger(__name__) + +# No TTY is attached to these subprocess calls, so an interactive +# "authenticity of host ... can't be established, continue connecting?" +# prompt would hang forever. accept-new trusts (and pins) an unseen host key +# automatically but still rejects a *changed* one — same safety net `ssh` +# gives you interactively, minus the prompt. +_SSH_OPTS = ["-o", "StrictHostKeyChecking=accept-new", "-o", "BatchMode=yes"] + + +def _resolve_local_dataset_root(config: TrainingRequest) -> Path: + """Same default-local-directory logic as record.py's resume fix: an + explicit dataset_root wins, otherwise fall back to the standard local + cache location for this repo_id.""" + if config.dataset_root: + return Path(config.dataset_root).expanduser().resolve() + from lerobot.utils.constants import HF_LEROBOT_HOME + + return (Path(HF_LEROBOT_HOME) / config.dataset_repo_id).resolve() + + +def _ssh_base_args(ssh: SshConnectionConfig) -> list[str]: + args = ["ssh", *_SSH_OPTS, "-p", str(ssh.port)] + if ssh.ssh_key_path: + args += ["-i", ssh.ssh_key_path] + args.append(f"{ssh.username}@{ssh.host}") + return args + + +def _scp_base_args(ssh: SshConnectionConfig) -> list[str]: + # scp's port flag is -P (capital), unlike ssh's -p — easy to typo. + args = ["scp", "-r", *_SSH_OPTS, "-P", str(ssh.port)] + if ssh.ssh_key_path: + args += ["-i", ssh.ssh_key_path] + return args + + +def _run_ssh_command(ssh: SshConnectionConfig, remote_command: str, timeout: float = 20.0) -> str: + """Run one short remote command over a fresh SSH connection and return + its stdout. Raises RuntimeError with stderr on a non-zero exit.""" + result = subprocess.run( + [*_ssh_base_args(ssh), remote_command], + capture_output=True, + text=True, + timeout=timeout, + ) + if result.returncode != 0: + raise RuntimeError(f"ssh command failed ({result.returncode}): {result.stderr.strip()}") + return result.stdout + + +def sync_dataset_up(ssh: SshConnectionConfig, local_root: Path, remote_datasets_dir: str) -> str: + """scp -r the local dataset directory under remote_datasets_dir on the + remote host. Returns the resulting remote dataset directory (mirrors + local_root's own directory name, so this never has to guess the + / split of a repo_id). + + No timeout: dataset video directories can be many GB, and this call is + expected to block for a while — see the module docstring's caveat about + that blocking the start-job request. + """ + if not local_root.is_dir(): + raise FileNotFoundError( + f"Dataset not found locally at {local_root} — record or download it before " + "training on a remote server (lelab only syncs what's already on this machine)." + ) + _run_ssh_command(ssh, f"mkdir -p {shlex.quote(remote_datasets_dir)}") + scp_cmd = [*_scp_base_args(ssh), str(local_root), f"{ssh.username}@{ssh.host}:{remote_datasets_dir}/"] + logger.info("Syncing dataset to remote host: %s", " ".join(scp_cmd)) + result = subprocess.run(scp_cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"scp dataset sync failed ({result.returncode}): {result.stderr.strip()}") + return f"{remote_datasets_dir}/{local_root.name}" + + +def pull_checkpoints(ssh: SshConnectionConfig, remote_output_dir: str, local_output_dir: str) -> None: + """scp -r the remote checkpoints/ tree down into local_output_dir, so the + existing local-checkpoint listing code (which just scans the filesystem) + picks them up unchanged. Standalone function (not a method) so a fresh + JobRegistry can pull checkpoints for a job whose SshRemoteJobRunner + instance didn't survive a lelab restart — record.ssh_config / + record.ssh_remote_dir are enough on their own. + + Best-effort: raises on any failure (missing remote dir before the first + checkpoint is saved, network hiccup, ...) and the caller is expected to + swallow it — see JobRegistry._maybe_pull_ssh_checkpoints. + """ + Path(local_output_dir).mkdir(parents=True, exist_ok=True) + remote_checkpoints_dir = f"{remote_output_dir}/checkpoints" + scp_cmd = [ + *_scp_base_args(ssh), + f"{ssh.username}@{ssh.host}:{remote_checkpoints_dir}", + str(local_output_dir), + ] + result = subprocess.run(scp_cmd, capture_output=True, text=True, timeout=120) + if result.returncode != 0: + raise RuntimeError(f"scp checkpoint pull failed ({result.returncode}): {result.stderr.strip()}") + + +class SshRemoteJobRunner(SubprocessJobRunner): + """Run a training on a user-owned remote server over SSH. + + Reuses SubprocessJobRunner's spawn/pump/parse pipeline exactly like + HfCloudJobRunner does, but the tailed subprocess here is the local `ssh` + client itself: its stdout mirrors the remote training command's stdout + for as long as the SSH session stays open, so no separate "tail -f" + round-trip is needed. `start_new_session=True` (set in `_spawn`) lets + that local ssh client — and therefore the log stream — survive a uvicorn + --reload the same way a local job's subprocess does. + + stop() has to reach the *remote* process explicitly: killing the local + ssh client doesn't reliably propagate a signal to whatever it's running + remotely. It pattern-matches the remote output dir (unique per job, since + it's passed as --output_dir) via `pkill -f` over a fresh connection. + """ + + def __init__( + self, + metrics: TrainingMetrics, + log_file_path: Path, + ssh: SshConnectionConfig, + ) -> None: + super().__init__(metrics, log_file_path) + self._ssh = ssh + self._remote_output_dir: str | None = None + + def start(self, job_id: str, config: TrainingRequest, output_dir: str) -> None: + # `output_dir` is the local path JobRegistry reserves for this job + # (used later to land pulled-back checkpoints) — it's not a path on + # the remote host, so the actual --output_dir the remote process gets + # is a fresh one under the user's configured remote_workdir instead. + local_dataset_root = _resolve_local_dataset_root(config) + remote_datasets_dir = f"{self._ssh.remote_workdir}/datasets" + remote_dataset_root = sync_dataset_up(self._ssh, local_dataset_root, remote_datasets_dir) + + remote_output_dir = f"{self._ssh.remote_workdir}/outputs/{job_id}" + self._remote_output_dir = remote_output_dir + + remote_config = config.model_copy(update={"dataset_root": remote_dataset_root}) + + # Build with a placeholder interpreter, then splice in + # remote_python_cmd unquoted — it may be a shell fragment like + # "source ~/venv/bin/activate && python", which shlex.quote would + # otherwise mangle into a single inert token. + cmd = build_training_command(remote_config, remote_output_dir, "__LELAB_PYEXEC__") + assert cmd[0] == "__LELAB_PYEXEC__" + quoted_rest = " ".join(shlex.quote(a) for a in cmd[1:]) + remote_train_cmd = f"{self._ssh.remote_python_cmd} {quoted_rest}" + + remote_shell = ( + f"mkdir -p {shlex.quote(remote_output_dir)} && " + f"cd {shlex.quote(self._ssh.remote_workdir)} && " + f"{remote_train_cmd}" + ) + ssh_cmd = [*_ssh_base_args(self._ssh), remote_shell] + logger.info("Starting ssh_remote job %s on %s: %s", job_id, self._ssh.host, remote_shell) + self._spawn(ssh_cmd, thread_name=f"job-{job_id}-ssh") + + def remote_output_dir(self) -> str: + assert self._remote_output_dir is not None, "remote_output_dir() called before start()" + return self._remote_output_dir + + def stop(self) -> None: + if self._remote_output_dir: + try: + _run_ssh_command(self._ssh, f"pkill -f {shlex.quote(self._remote_output_dir)}") + except Exception as exc: + # Already finished, or pkill matched nothing — fine either way. + logger.info("Remote pkill for %s ignored: %s", self._remote_output_dir, exc) + super().stop()