Skip to content

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

Open
ZouzouWP wants to merge 1 commit into
huggingface:mainfrom
ZouzouWP:feat/ssh-remote-training
Open

feat: remote training over SSH (no HF Jobs/account required)#76
ZouzouWP wants to merge 1 commit into
huggingface:mainfrom
ZouzouWP:feat/ssh-remote-training

Conversation

@ZouzouWP

Copy link
Copy Markdown

The only non-local training option was HF Jobs — cloud, paid, and tied to a Hugging Face account. This adds a way to train on any machine you can already SSH into (a spare GPU box, a rented cloud instance), with no account or extra service involved.

What changed

  • New ssh_remote job target, alongside local and hf_cloud. The connection form asks for host/port/username, an optional private key path (falls back to the local ssh-agent/default identity — LeLab never handles or stores a password), a remote working directory, and a free-form "remote Python command" (e.g. python, or source ~/venv/bin/activate && python), since LeLab has no way to know how the remote server is set up.
  • New lelab/runners/ssh_remote.py (SshRemoteJobRunner): stages the dataset via scp, launches training over ssh, and tracks the local ssh client's PID as process_pid.
  • Checkpoints are pulled back with a rate-limited scp (at most once every 30s, since the frontend polls checkpoints roughly every 5s) into the same local output_dir a local job would use — so checkpoint listing and inference reuse the exact same local code path, no separate "remote checkpoint" representation needed anywhere.
sequenceDiagram
    participant U as User
    participant T as TargetCard (Training page)
    participant API as LeLab backend
    participant R as SshRemoteJobRunner
    participant S as Remote server

    U->>T: Fill host/user/key/workdir/python cmd, start training
    T->>API: POST /jobs/training {target: {runner: ssh_remote, ssh: {...}}}
    API->>R: start()
    R->>S: scp dataset -> remote_workdir
    R->>S: ssh: run training command
    loop every ~30s at most (rate-limited)
        API->>S: scp checkpoints/ -> local output_dir
    end
    API-->>T: same checkpoint/metrics UI as a local job
Loading

Testing

Implemented and exercised locally (dataset staging, command construction, checkpoint-pull rate limiting), but not yet run end-to-end against a real remote host — needs a pass against a reachable GPU box before merge.

The only non-local training option was HF Jobs — cloud, paid, and tied to a Hugging Face account. This adds a way to train on any machine you can already SSH into (a spare GPU box, a rented cloud instance), with no account or extra service involved.

## What changed

- New `ssh_remote` job target, alongside `local` and `hf_cloud`. The connection form asks for host/port/username, an optional private key path (falls back to the local ssh-agent/default identity — LeLab never handles or stores a password), a remote working directory, and a free-form "remote Python command" (e.g. `python`, or `source ~/venv/bin/activate && python`), since LeLab has no way to know how the remote server is set up.
- New `lelab/runners/ssh_remote.py` (`SshRemoteJobRunner`): stages the dataset via `scp`, launches training over `ssh`, and tracks the local `ssh` client's PID as `process_pid`.
- Checkpoints are pulled back with a rate-limited `scp` (at most once every 30s, since the frontend polls checkpoints roughly every 5s) into the same local `output_dir` a local job would use — so checkpoint listing and inference reuse the exact same local code path, no separate "remote checkpoint" representation needed anywhere.

```mermaid
sequenceDiagram
    participant U as User
    participant T as TargetCard (Training page)
    participant API as LeLab backend
    participant R as SshRemoteJobRunner
    participant S as Remote server

    U->>T: Fill host/user/key/workdir/python cmd, start training
    T->>API: POST /jobs/training {target: {runner: ssh_remote, ssh: {...}}}
    API->>R: start()
    R->>S: scp dataset -> remote_workdir
    R->>S: ssh: run training command
    loop every ~30s at most (rate-limited)
        API->>S: scp checkpoints/ -> local output_dir
    end
    API-->>T: same checkpoint/metrics UI as a local job
```

## Testing

Implemented and exercised locally (dataset staging, command construction, checkpoint-pull rate limiting), but **not yet run end-to-end against a real remote host** — needs a pass against a reachable GPU box before merge.
@nicolas-rabault
nicolas-rabault self-requested a review August 3, 2026 15:15

@nicolas-rabault nicolas-rabault left a comment

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.

Thanks for this, the shape is right and the trust story is the one I care about: no passwords, key path only, BatchMode so it fails fast, and accept-new still rejecting a changed host key. Reusing SubprocessJobRunner with the local ssh client as the tailed process, and landing pulled checkpoints in the same local output_dir so listing and inference need no changes at all, is the good version of this feature.

Three things block merge, all in the review. A lelab restart orphans the remote training (no ssh_remote branch in _load_from_disk, so the record goes to "interrupted" and stop() then refuses, while the run keeps the GPU). A broken remote renders as "no checkpoints yet" because the pull swallows every error, and that same pull blocks the watchdog thread. And remote_python_cmd coming from the request body means any page the user has open can reach it, since the server has no auth and CORS is open.

Smaller things I am not holding the PR for: pkill -f also matches the shell running it,
no ConnectTimeout or ServerAliveInterval so a half-open connection sits at "running", a ~/... remote_workdir uploads to one path and tells the trainer another, and a preflight import lerobot would save a long scp on a typo.

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}"

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.

remote_python_cmd comes straight from the POST body and is spliced in unquoted, which makes /jobs/training effectively "run this shell string on that host with my ssh keys". Two things make that reachable from outside the UI: the server has no auth, and server.py sets CORS allow_origins=["*"] with allow_methods=["*"], so the preflight passes and any page the user has open can POST to 127.0.0.1:8000 and get a command run on any box the local ssh agent reaches. The same page can GET /jobs and read ssh_config back, since JobRecord is returned as is (host, port, username, key path).

Nothing here is wrong for a user driving their own UI, but it turns a local tool into a proxy for their ssh credentials, which no existing endpoint is. Blocking until we pick one: connection details come from a server side file under lelab/utils/config.py, the way saved ports and configs already work, with the request only naming a saved profile, or the training endpoints get a same origin check. That call is mine, not yours, tell me which you would rather build and I will decide.

Comment thread lelab/jobs.py
Comment on lines +118 to +123
# 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

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.

Comment thread lelab/jobs.py
from .runners.ssh_remote import pull_checkpoints # lazy: avoid circular import

try:
pull_checkpoints(record.ssh_config, record.ssh_remote_dir, record.output_dir)

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants