Skip to content

feat: remote training over SSH (no HF Jobs/account required) #76

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZouzouWP wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: remote training over SSH (no HF Jobs/account required) #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 101 additions & 2 deletions frontend/src/components/training/config/TargetCard.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,27 @@
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,
SelectItem,
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[];
Expand All @@ -36,17 +47,30 @@ const TargetCard: React.FC<TargetCardProps> = ({
}) => {
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 = <K extends keyof SshConnectionConfig>(key: K, value: SshConnectionConfig[K]) => {
updateConfig("target", {
runner: "ssh_remote",
ssh: { ...(target.ssh ?? DEFAULT_SSH), [key]: value },
});
};

return (
<Card className="bg-slate-800/50 border-slate-700 rounded-xl">
<CardHeader>
Expand All @@ -61,6 +85,7 @@ const TargetCard: React.FC<TargetCardProps> = ({
</SelectTrigger>
<SelectContent className="bg-slate-800 border-slate-600 text-white">
<SelectItem value="local">Local — your machine (free)</SelectItem>
<SelectItem value="ssh">Remote server (SSH) — your own machine</SelectItem>
{flavors.map((f) => (
<SelectItem
key={f.name}
Expand All @@ -82,6 +107,80 @@ const TargetCard: React.FC<TargetCardProps> = ({
account when training completes.
</p>
</div>

{target.runner === "ssh_remote" && (
<div className="space-y-4 pt-2 border-t border-slate-700">
<p className="text-xs text-slate-500">
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.
</p>
<div className="grid grid-cols-1 md:grid-cols-3 gap-3">
<div className="md:col-span-2">
<Label className="text-slate-300">Host</Label>
<Input
value={target.ssh?.host ?? ""}
onChange={(e) => 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"
/>
</div>
<div>
<Label className="text-slate-300">Port</Label>
<NumberInput
value={target.ssh?.port ?? 22}
onChange={(v) => v !== undefined && updateSsh("port", v)}
className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1"
/>
</div>
</div>
<div className="grid grid-cols-1 md:grid-cols-2 gap-3">
<div>
<Label className="text-slate-300">Username</Label>
<Input
value={target.ssh?.username ?? ""}
onChange={(e) => updateSsh("username", e.target.value)}
placeholder="noah"
className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1"
/>
</div>
<div>
<Label className="text-slate-300">SSH key path (optional)</Label>
<Input
value={target.ssh?.ssh_key_path ?? ""}
onChange={(e) => 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"
/>
</div>
</div>
<div>
<Label className="text-slate-300">Remote working directory</Label>
<Input
value={target.ssh?.remote_workdir ?? ""}
onChange={(e) => updateSsh("remote_workdir", e.target.value)}
placeholder="/home/noah/lelab-runs"
className="bg-slate-900 border-slate-600 text-white rounded-lg mt-1"
/>
<p className="text-xs text-slate-500 mt-1">
Datasets and outputs are staged under here — nothing outside it is touched.
</p>
</div>
<div>
<Label className="text-slate-300">Remote Python command</Label>
<Input
value={target.ssh?.remote_python_cmd ?? "python"}
onChange={(e) => 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"
/>
<p className="text-xs text-slate-500 mt-1">
Whatever it takes to reach an interpreter with lerobot installed —
lelab runs this verbatim, followed by the training command.
</p>
</div>
</div>
)}
</CardContent>
</Card>
);
Expand Down
15 changes: 14 additions & 1 deletion frontend/src/components/training/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
13 changes: 12 additions & 1 deletion frontend/src/lib/jobsApi.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
10 changes: 9 additions & 1 deletion frontend/src/pages/Training.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,20 +275,28 @@ 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 =
isStarting ||
!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 (
Expand Down
87 changes: 81 additions & 6 deletions lelab/jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Comment on lines +118 to +123

@nicolas-rabault nicolas-rabault Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment is right about checkpoint pulls, but the other half of restart handling is missing, and it loses a running job. _load_from_disk only reattaches runner == "local" and runner == "hf_cloud"; an ssh_remote record that was running falls into the final else, is marked interrupted, and gets no runner. From then on stop(job_id) raises JobNotRunningError, so the remote training keeps holding the GPU with nothing in the UI able to reach it, even though ssh_remote_dir is persisted right here and would be enough. A lelab --dev reload hits this too, not just a full restart.

Capturing the remote pid at start is the fix that covers the most ground: have the remote shell print it (echo LELAB_REMOTE_PID=$$), parse it in _on_line the way HfCloudJobRunner parses the HF job id, and persist it on the record. Then _load_from_disk can get an ssh_remote branch that checks that pid over ssh and either reattaches a runner able to issue the kill or finalises the record honestly, and stop() can signal exactly that process group instead of pattern matching.

# 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:
Expand Down Expand Up @@ -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$")


Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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)

Expand All @@ -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
Expand Down Expand Up @@ -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/<step>/ 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)

@nicolas-rabault nicolas-rabault Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two problems here, and together they are the reason I cannot take this yet.

Every exception is swallowed at info level, so a broken remote is indistinguishable from a run that has not saved a checkpoint yet: expired key, full remote disk, host gone, all render as "none yet" in the UI while the user waits for checkpoints that will never arrive. Count consecutive failures and put the last error on the record so the page can say the pull is failing, keeping the listing itself degrading gracefully as you have it.

It is also blocking network I/O on whichever thread asks for a checkpoint count, and that is not only the frontend poll: _count_checkpoints runs inside _tick on every watchdog pass and inside list() and get() on the request path. One slow remote parks the watchdog for the full 120s scp timeout, which stalls progress broadcasts and finalisation for every other job including local ones, or hangs a GET /jobs. Pull on a background thread and let readers see whatever landed last.

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.

Expand Down
Loading