diff --git a/CLAUDE.md b/CLAUDE.md index 079894fbe..58c5e2dd7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -16,7 +16,7 @@ Install (editable, requires Python ≥3.10): pip install -e . ``` -`lerobot` is pinned to a release tag on `huggingface/lerobot` (currently `v0.6.1`, see [pyproject.toml](pyproject.toml)) — installed from git rather than PyPI because the published wheels don't expose the `lerobot-rollout` script that [lelab/rollout.py](lelab/rollout.py) shells out to. Bump the tag when you want a newer upstream release; expect import-path drift (e.g. `lerobot.configs.dataset`, `lerobot.utils.feature_utils`, `lerobot.common.control_utils`) and adjust call sites accordingly. +`lerobot` is pinned to a ref on `huggingface/lerobot` (currently `v0.7.0`, see [pyproject.toml](pyproject.toml)), installed from git rather than PyPI because the published wheels don't expose the `lerobot-rollout` script that [lelab/rollout.py](lelab/rollout.py) shells out to. Bump the ref when you want a newer upstream release; expect import-path drift (e.g. `lerobot.configs.dataset`, `lerobot.utils.feature_utils`, `lerobot.common.control_utils`) and CLI flag renames (v0.7.0 replaced `--wandb.*` with `--trackio.*`). [tests/test_lerobot_cli.py](tests/test_lerobot_cli.py) runs the training and rollout argv LeLab builds through the installed LeRobot's parsers, so run the tests right after a bump. Since v0.7.0 the arms' motors go through LeRobot's rustypot-backed `SerialMotorsBus`: use its public methods (`set_operating_mode`, `set_half_turn_homings`, `disconnect`), not Feetech SDK internals. Run servers (entry point defined in [pyproject.toml](pyproject.toml)): @@ -35,8 +35,8 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml); [server.py](lelab/server.py) is a thin FastAPI router. Each feature lives in its own module that owns its global state and exposes `handle_*` functions plus a Pydantic request model: -- [record.py](lelab/record.py) — dataset recording (wraps `lerobot.record.record`); patches `lerobot.common.utils.control_utils` keyboard listener so frontend buttons replace arrow-key controls. The worker is a daemon thread, so the server's shutdown hook calls `stop_recording_and_wait()` to let an active session finalize on Ctrl+C. -- [teleoperate.py](lelab/teleoperate.py) — leader→follower teleoperation (wraps `lerobot.teleoperate`). +- [record.py](lelab/record.py): dataset recording around LeRobot's `record_loop`; frontend buttons set the same `events` dict LeRobot's keyboard listener would. The worker is a daemon thread, so the server's shutdown hook calls `stop_recording_and_wait()` to let an active session finalize on Ctrl+C. +- [teleoperate.py](lelab/teleoperate.py): leader→follower teleoperation loop over LeRobot's SO-101 leader and follower. - [replay.py](lelab/replay.py): plays a recorded episode's actions on the follower arm, an in-process take on `lerobot-replay` so the dataset browser can stop it. - [calibrate.py](lelab/calibrate.py) — step-by-step web calibration with a `CalibrationManager` singleton and `_step_complete` threading.Event. - [train.py](lelab/train.py) — wraps the LeRobot training CLI as a subprocess (psutil for lifecycle, queue for log streaming). diff --git a/frontend/src/components/jobs/JobsSection.tsx b/frontend/src/components/jobs/JobsSection.tsx index 042446475..2e189eb30 100644 --- a/frontend/src/components/jobs/JobsSection.tsx +++ b/frontend/src/components/jobs/JobsSection.tsx @@ -104,7 +104,7 @@ const JobsSection: React.FC = () => { ...j, state: s.state, metrics: s.metrics, - wandb_run_url: s.wandb_run_url, + tracking_url: s.tracking_url, checkpoint_count: s.checkpoint_count, }; }); diff --git a/frontend/src/components/training/PolicyExtraDialog.tsx b/frontend/src/components/training/PolicyExtraDialog.tsx index ff97982a2..7719548d9 100644 --- a/frontend/src/components/training/PolicyExtraDialog.tsx +++ b/frontend/src/components/training/PolicyExtraDialog.tsx @@ -18,7 +18,7 @@ interface Props { open: boolean; onOpenChange: (open: boolean) => void; policyType: string; - packageName: string; // the probed module, e.g. "transformers" + packageName: string; // the requirements still unmet, e.g. "transformers<5.23.0,>=5.22.0" installTarget: string; // e.g. "lerobot[smolvla]" installHint: string; // e.g. "pip install 'lerobot[smolvla]'" } @@ -63,11 +63,11 @@ const PolicyExtraDialog: React.FC = ({ idleTitle={title} idleDescription={ <> - Training a {policyType} policy needs the{" "} + Training a {policyType} policy needs{" "} {packageName}{" "} - package (installed via{" "} + (installed via{" "} {installTarget}), - which isn't in this environment yet. Install it to train this policy. + which this environment doesn't have yet. Install it to train this policy. } doneDescription={} diff --git a/frontend/src/components/training/WandbInstallDialog.tsx b/frontend/src/components/training/WandbInstallDialog.tsx deleted file mode 100644 index 3ead2f07c..000000000 --- a/frontend/src/components/training/WandbInstallDialog.tsx +++ /dev/null @@ -1,68 +0,0 @@ -import React from "react"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogHeader, - DialogTitle, -} from "@/components/ui/dialog"; -import { useInstallExtra } from "@/hooks/useInstallExtra"; -import { - InstallProgress, - InstallTitleIcon, - RestartInstructions, - installTitle, -} from "./InstallProgress"; - -interface Props { - open: boolean; - onOpenChange: (open: boolean) => void; - installHint: string; -} - -const WandbInstallDialog: React.FC = ({ open, onOpenChange, installHint }) => { - const install = useInstallExtra("system/wandb-extra", open); - - return ( - - - - - - {installTitle(install.state, "Weights & Biases Not Installed")} - - - Install the wandb package to enable W&B logging. - - - -
- - Enabling W&B logging requires the{" "} - - wandb - {" "} - package, which isn't installed in this environment. Install it - to log this run to W&B. - - } - doneDescription={} - /> -
- - - ); -}; - -export default WandbInstallDialog; diff --git a/frontend/src/components/training/config/AdvancedCard.tsx b/frontend/src/components/training/config/AdvancedCard.tsx index fee8d443f..70bd886fa 100644 --- a/frontend/src/components/training/config/AdvancedCard.tsx +++ b/frontend/src/components/training/config/AdvancedCard.tsx @@ -251,74 +251,33 @@ const AdvancedCard: React.FC = ({ config, updateConfig }) - {config.wandb_enable && ( + {config.trackio_enable && ( <>
- Weights & Biases + Trackio
-
-
- - - updateConfig('wandb_notes', e.target.value || undefined) - } - placeholder="Training run notes..." - className="bg-slate-900 border-slate-600 text-white rounded-lg" - /> -
-
- - -
-
- - updateConfig('wandb_disable_artifact', checked) - } - /> - +

+ Signed in, runs sync to a private lelab-trackio Space on your account by default. +

)} - {!config.wandb_enable && } + {!config.trackio_enable && } {/* Misc */}
diff --git a/frontend/src/components/training/config/EssentialsCard.tsx b/frontend/src/components/training/config/EssentialsCard.tsx index 0103668e2..ebb521ad5 100644 --- a/frontend/src/components/training/config/EssentialsCard.tsx +++ b/frontend/src/components/training/config/EssentialsCard.tsx @@ -1,4 +1,4 @@ -import React, { useState } from 'react'; +import React from 'react'; import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'; import { Input } from '@/components/ui/input'; import { NumberInput } from '@/components/ui/number-input'; @@ -14,8 +14,6 @@ import { import { ConfigComponentProps } from '../types'; import DatasetCombobox from '@/components/replay/DatasetCombobox'; import { DatasetItem } from '@/lib/replayApi'; -import WandbInstallDialog from '../WandbInstallDialog'; -import { useApi } from '@/contexts/ApiContext'; interface EssentialsCardProps extends ConfigComponentProps { datasets: DatasetItem[]; @@ -23,35 +21,6 @@ interface EssentialsCardProps extends ConfigComponentProps { } const EssentialsCard: React.FC = ({ config, updateConfig, datasets, datasetsLoading }) => { - const { baseUrl, fetchWithHeaders } = useApi(); - const [wandbDialogOpen, setWandbDialogOpen] = useState(false); - const [wandbInstallHint, setWandbInstallHint] = useState('pip install wandb'); - - const handleWandbToggle = async (checked: boolean) => { - if (!checked) { - updateConfig('wandb_enable', false); - return; - } - // Check availability before flipping the switch on. If wandb isn't - // importable in this lelab process, surface the same install flow used - // for the training extra (accelerate) instead of letting the user start - // a run that will fail. - try { - const r = await fetchWithHeaders(`${baseUrl}/system/wandb-extra`); - const data: { available: boolean; install_hint: string } = await r.json(); - if (data.available) { - updateConfig('wandb_enable', true); - } else { - setWandbInstallHint(data.install_hint); - setWandbDialogOpen(true); - } - } catch { - // Backend unreachable — let the user proceed; training start will - // surface the real error if wandb is genuinely missing. - updateConfig('wandb_enable', true); - } - }; - return ( @@ -140,35 +109,29 @@ const EssentialsCard: React.FC = ({ config, updateConfig, d
updateConfig('trackio_enable', checked)} className="data-[state=checked]:bg-green-500" /> -
- - - {config.wandb_enable && ( + {config.trackio_enable && (
-
diff --git a/frontend/src/components/training/types.ts b/frontend/src/components/training/types.ts index cd4882ddd..9bcc4498b 100644 --- a/frontend/src/components/training/types.ts +++ b/frontend/src/components/training/types.ts @@ -21,13 +21,10 @@ export interface TrainingConfig { // Output configuration resume: boolean; - // Weights & Biases - wandb_enable: boolean; - wandb_project?: string; - wandb_entity?: string; - wandb_notes?: string; - wandb_mode?: string; - wandb_disable_artifact: boolean; + // Trackio + trackio_enable: boolean; + trackio_project?: string; + trackio_space_id?: string; // Policy-specific parameters policy_device?: string; diff --git a/frontend/src/hooks/useInstallExtra.ts b/frontend/src/hooks/useInstallExtra.ts index f53ba137b..a7567d3b0 100644 --- a/frontend/src/hooks/useInstallExtra.ts +++ b/frontend/src/hooks/useInstallExtra.ts @@ -26,7 +26,7 @@ export interface UseInstallExtraResult { } /** - * Drives the backend extra-install flow (`accelerate`, `wandb`, …). Seeds state + * Drives the backend extra-install flow (`accelerate`, policy extras, …). Seeds state * from `${endpointPrefix}/install-status`, polls while installing, and exposes * install/retry handlers. Pass `enabled=false` to gate seeding on dialog open. */ diff --git a/frontend/src/lib/jobsApi.ts b/frontend/src/lib/jobsApi.ts index 8cd82f70e..0cfbfeda3 100644 --- a/frontend/src/lib/jobsApi.ts +++ b/frontend/src/lib/jobsApi.ts @@ -36,12 +36,9 @@ export interface TrainingRequest { save_freq: number; save_checkpoint: boolean; resume: boolean; - wandb_enable: boolean; - wandb_project?: string; - wandb_entity?: string; - wandb_notes?: string; - wandb_mode?: string; - wandb_disable_artifact: boolean; + trackio_enable: boolean; + trackio_project?: string; + trackio_space_id?: string; policy_device?: string; policy_use_amp: boolean; optimizer_type?: string; @@ -79,7 +76,7 @@ export interface JobRecord { hf_flavor: string | null; hf_repo_id: string | null; hf_job_url: string | null; - wandb_run_url: string | null; + tracking_url: string | null; checkpoint_count: number; } @@ -89,7 +86,7 @@ export interface JobProgressSnapshot { id: string; state: JobState; metrics: TrainingMetrics; - wandb_run_url: string | null; + tracking_url: string | null; checkpoint_count: number; } diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx index 2dd9a23f3..fafe4bcd1 100644 --- a/frontend/src/pages/Training.tsx +++ b/frontend/src/pages/Training.tsx @@ -80,12 +80,9 @@ function configToRequest(c: TrainingConfig): TrainingRequest { save_freq: c.save_freq, save_checkpoint: c.save_checkpoint, resume: c.resume, - wandb_enable: c.wandb_enable, - wandb_project: c.wandb_project, - wandb_entity: c.wandb_entity, - wandb_notes: c.wandb_notes, - wandb_mode: c.wandb_mode, - wandb_disable_artifact: c.wandb_disable_artifact, + trackio_enable: c.trackio_enable, + trackio_project: c.trackio_project, + trackio_space_id: c.trackio_space_id, policy_device: c.policy_device, policy_use_amp: c.policy_use_amp, optimizer_type: c.optimizer_type, @@ -131,9 +128,7 @@ const ConfigurationMode: React.FC = () => { save_freq: 1000, save_checkpoint: true, resume: false, - wandb_enable: false, - wandb_mode: "online", - wandb_disable_artifact: false, + trackio_enable: false, policy_use_amp: false, optimizer_type: "adam", use_policy_training_preset: true, @@ -586,14 +581,14 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => { View on Hub ↗ )} - {job.wandb_run_url && ( + {job.tracking_url && ( - View on W&B ↗ + View metrics ↗ )} diff --git a/lelab/calibrate.py b/lelab/calibrate.py index bc4e241d7..a3f40620c 100644 --- a/lelab/calibrate.py +++ b/lelab/calibrate.py @@ -22,13 +22,11 @@ import logging import threading import time -import traceback from copy import deepcopy from dataclasses import dataclass from typing import Any, Literal from lerobot.motors import MotorCalibration -from lerobot.motors.feetech import OperatingMode from lerobot.robots import ( Robot, make_robot_from_config, @@ -314,8 +312,7 @@ def _calibration_worker(self, request: CalibrationRequest): self._update_status(error=str(e)) self._cleanup_and_finish(str(e), status="error") except Exception as e: - logger.error(f"Calibration error: {e}") - logger.error(traceback.format_exc()) + logger.exception(f"Calibration error: {e}") # Ensure cleanup happens even on error self._cleanup_and_finish(f"Calibration failed: {e}", status="error") finally: @@ -333,19 +330,12 @@ def _step_homing(self): # Disable torque to allow manual movement during recording self.device.bus.disable_torque() - for motor in self.device.bus.motors: - self.device.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value) + self.device.bus.set_operating_mode("position") - self.device.bus.reset_calibration() - actual_positions = self.device.bus.sync_read("Present_Position", normalize=False) - logger.info(f"Current positions for homing: {actual_positions}") - - self._homing_offsets = self.device.bus._get_half_turn_homings(actual_positions) + # Resets the calibration, then writes the offsets that center each present position. + self._homing_offsets = self.device.bus.set_half_turn_homings() logger.info(f"Calculated homing offsets: {self._homing_offsets}") - for motor, offset in self._homing_offsets.items(): - self.device.bus.write("Homing_Offset", motor, offset) - def _step_range_recording(self): """Record range of motion as the user moves all joints.""" logger.info("Starting range recording step") @@ -397,18 +387,7 @@ def _step_range_recording(self): # Record positions until user completes step while not self._step_complete.is_set() and not self.stop_calibration: try: - # Try reading positions with retry on port contention - positions = None - for attempt in range(3): # Try up to 3 times - try: - positions = self.device.bus.sync_read("Present_Position", normalize=False) - break # Success, exit retry loop - except Exception as read_error: - if "Port is in use" in str(read_error) and attempt < 2: - time.sleep(0.01) # Short delay before retry - continue - else: - raise read_error # Re-raise if not port contention or final attempt + positions = self.device.bus.sync_read("Present_Position", normalize=False) if positions: # Validate the readings - filter out invalid/zero values @@ -444,11 +423,7 @@ def _step_range_recording(self): except CalibrationDiscontinuityError: raise except Exception as e: - if "Port is in use" in str(e): - logger.debug(f"Port busy during position read: {e}") - else: - logger.warning(f"Error reading positions during recording: {e}") - # Increase sleep time on error to reduce port contention + logger.warning(f"Error reading positions during recording: {e}") time.sleep(0.2) if self.stop_calibration: diff --git a/lelab/jobs.py b/lelab/jobs.py index cca712970..2710381d9 100644 --- a/lelab/jobs.py +++ b/lelab/jobs.py @@ -36,7 +36,7 @@ from queue import Empty, Queue from typing import Literal, Protocol, runtime_checkable -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel, Field from .train import TrainingRequest @@ -88,8 +88,11 @@ class JobRecord(BaseModel): hf_flavor: str | None = None hf_repo_id: str | None = None hf_job_url: str | None = None - # Captured from training stdout the first time wandb prints the run URL. - wandb_run_url: str | None = None + # Captured from training stdout the first time Trackio prints the run's Space URL. + # Records persisted before the W&B to Trackio switch stored it as wandb_run_url. + tracking_url: str | None = Field( + default=None, validation_alias=AliasChoices("tracking_url", "wandb_run_url") + ) # Number of checkpoints currently visible (local: filesystem; cloud: # Hub repo). Filled in by JobRegistry.list/get; persisted as zero. checkpoint_count: int = 0 @@ -141,20 +144,20 @@ def stop(self) -> None: ... def is_running(self) -> bool: ... def returncode(self) -> int | None: ... def stream_log_lines(self) -> list[LogLine]: ... - def wandb_run_url(self) -> str | None: ... + def tracking_url(self) -> str | None: ... # tqdm progress: "Training: 1%|▏ | 125/10000 [02:02<2:36:10, 1.05step/s]" _TQDM_RE = re.compile(r"Training:\s*\d+%[^|]*\|[^|]*\|\s*(\d+)/(\d+)\s*\[(?:[\d:]+)<([\d:]+)") -# Wandb prints something like "wandb: 🚀 View run at https://wandb.ai///runs/" -# when it boots. We capture the first URL of that shape we see. -_WANDB_URL_RE = re.compile(r"https://wandb\.ai/[^\s/]+/[^\s/]+/runs/[A-Za-z0-9]+") +# LeRobot logs "Track this run --> https://huggingface.co/spaces//" when the run +# syncs to a Space; a run kept in the local Trackio database prints a `trackio show` command. +_TRACKING_URL_RE = re.compile(r"Track this run --> (https://\S+)") -def extract_wandb_run_url(line: str) -> str | None: - match = _WANDB_URL_RE.search(line) - return match.group(0) if match else None +def extract_tracking_url(line: str) -> str | None: + match = _TRACKING_URL_RE.search(line) + return match.group(1) if match else None def _parse_duration(s: str) -> float | None: @@ -230,7 +233,7 @@ def __init__( self._stop_event = threading.Event() self._log_file_path = log_file_path self._log_file = None # type: ignore[assignment] - self._wandb_run_url: str | None = None + self._tracking_url: str | None = None def _open_log_file(self) -> None: """Open the persistent log sink (one JSON line per consumed line). @@ -306,8 +309,8 @@ def stream_log_lines(self) -> list[LogLine]: pass return out - def wandb_run_url(self) -> str | None: - return self._wandb_run_url + def tracking_url(self) -> str | None: + return self._tracking_url # -- internals -- @@ -328,10 +331,10 @@ def _consume_lines(self, lines: Iterable[str]) -> None: if not stripped: continue parse_metrics_into(stripped, self._metrics) - if self._wandb_run_url is None: - url = extract_wandb_run_url(stripped) + if self._tracking_url is None: + url = extract_tracking_url(stripped) if url is not None: - self._wandb_run_url = url + self._tracking_url = url self._on_line(stripped) log_line = LogLine(timestamp=time.time(), message=stripped) if self._log_file is not None: @@ -402,7 +405,7 @@ def __init__( # Replay everything that's already on disk so the parser catches up # on metrics, then tail from the current EOF. self._tail_offset = 0 - self._wandb_run_url: str | None = None + self._tracking_url: str | None = None def start(self, job_id: str, config: TrainingRequest, output_dir: str) -> None: # Required by JobRunner Protocol but irrelevant here; the subprocess @@ -446,8 +449,8 @@ def stream_log_lines(self) -> list[LogLine]: def pid(self) -> int | None: return self._pid - def wandb_run_url(self) -> str | None: - return self._wandb_run_url + def tracking_url(self) -> str | None: + return self._tracking_url # -- internals -- @@ -476,10 +479,10 @@ def _tail_loop(self) -> None: except Exception: continue parse_metrics_into(log_line.message, self._metrics) - if self._wandb_run_url is None: - url = extract_wandb_run_url(log_line.message) + if self._tracking_url is None: + url = extract_tracking_url(log_line.message) if url is not None: - self._wandb_run_url = url + self._tracking_url = url if self._log_queue.qsize() >= 1000: with contextlib.suppress(Empty): self._log_queue.get_nowait() @@ -631,6 +634,23 @@ def _generate_job_id(policy_type: str, dataset_repo_id: str) -> str: return f"{policy_type}_{dataset_slug}_{timestamp}" +def _with_default_trackio_space(config: TrainingRequest) -> TrainingRequest: + """Sync a tracked run to the user's `lelab-trackio` Space when none is given. + + A Space is what gives the run a dashboard link, and LeRobot requires one for + HF Jobs, whose local Trackio database vanishes with the pod. Signed out, the + run stays in this machine's Trackio database. + """ + if not config.trackio_enable or config.trackio_space_id: + return config + from .utils.hf_auth import cached_whoami # lazy: keeps unit-test imports cheap + + info = cached_whoami() + if not info: + return config + return config.model_copy(update={"trackio_space_id": f"{info['name']}/lelab-trackio"}) + + def _job_dir(output_root: Path, job_id: str) -> Path: return output_root / job_id @@ -695,7 +715,7 @@ def __init__(self, output_root: Path) -> None: self._on_change: Callable[[], None] | None = None # Fired from the watchdog at ~1Hz with a compact snapshot of every - # running job (id, state, metrics, wandb url, checkpoint count) so + # running job (id, state, metrics, tracking url, checkpoint count) so # the dashboard keeps the progress bar live without refetching /jobs. self._on_progress: Callable[[builtins.list[dict]], None] | None = None @@ -818,6 +838,7 @@ def start(self, config: TrainingRequest, target: JobTarget | None = None) -> Job target = target or JobTarget() if target.runner == "hf_cloud" and not target.flavor: raise ValueError("flavor is required when runner is hf_cloud") + config = _with_default_trackio_space(config) with self._lock: # Local trainings are bounded by this machine's GPU/USB resources, @@ -1213,12 +1234,12 @@ def _tick(self) -> None: if runner is None or record is None: continue if runner.is_running(): - # Pull the wandb run URL once it appears in stdout. - if record.wandb_run_url is None: - url = runner.wandb_run_url() + # Pull the Trackio Space URL once it appears in stdout. + if record.tracking_url is None: + url = runner.tracking_url() if url is not None: with self._lock: - record.wandb_run_url = url + record.tracking_url = url self._persist(record, force=True) # Cloud jobs print their HF job id / page URL / model repo a few # seconds after start; capture them onto the record once parsed. @@ -1231,7 +1252,7 @@ def _tick(self) -> None: "id": record.id, "state": record.state, "metrics": record.metrics.model_dump(), - "wandb_run_url": record.wandb_run_url, + "tracking_url": record.tracking_url, "checkpoint_count": self._count_checkpoints(record), } ) @@ -1244,8 +1265,8 @@ def _tick(self) -> None: self._sync_cloud_ids(record, runner) rc = runner.returncode() with self._lock: - if record.wandb_run_url is None: - record.wandb_run_url = runner.wandb_run_url() + if record.tracking_url is None: + record.tracking_url = runner.tracking_url() record.state = "done" if rc == 0 else "failed" record.ended_at = time.time() record.exit_code = rc diff --git a/lelab/record.py b/lelab/record.py index 74044b7bb..f48724e2e 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -17,7 +17,6 @@ import shutil import threading import time -import traceback from datetime import datetime from typing import Any @@ -347,10 +346,7 @@ def recording_worker(): "robot_type": getattr(dataset.meta, "robot_type", "Unknown robot"), } except Exception as e: - # lerobot's init_logging() installs a root formatter that only - # renders the message, so logger.exception() would drop the - # traceback. Fold it into the message instead. - logger.error("Recording session failed: %s\n%s", e, traceback.format_exc()) + logger.exception("Recording session failed: %s", e) current_phase = "error" if recording_start_time: session_end_elapsed_seconds = int(time.time() - recording_start_time) @@ -760,8 +756,7 @@ def handle_upload_dataset(request: UploadRequest) -> dict[str, Any]: return {"success": False, "message": str(e)} except Exception as e: - logger.error(f"Error uploading dataset {request.dataset_repo_id}: {e}") - logger.error(f"Full traceback: {traceback.format_exc()}") + logger.exception(f"Error uploading dataset {request.dataset_repo_id}: {e}") err_text = str(e).lower() looks_like_auth = any( @@ -875,11 +870,13 @@ def record_with_web_events(cfg: RecordConfig, web_events: dict) -> LeRobotDatase logger.error(f"❌ TELEOP CONNECTION: Failed to connect teleoperator: {e}") raise - # Torque is still disabled at this point, which is what writing the - # calibration registers requires. + # Feetech motors only store calibration registers with the torque off, + # which also opens their EEPROM lock. logger.info("Writing calibration to motors...") + robot.bus.disable_torque() robot.bus.write_calibration(robot.calibration) if teleop is not None: + teleop.bus.disable_torque() teleop.bus.write_calibration(teleop.calibration) try: diff --git a/lelab/replay.py b/lelab/replay.py index f9d736424..2e3bf6d6d 100644 --- a/lelab/replay.py +++ b/lelab/replay.py @@ -88,6 +88,8 @@ def handle_start_replay(request: ReplayRequest) -> dict[str, Any]: f"Could not connect to the follower arm on {request.follower_port}. " "Make sure it's plugged in and powered on, then try again." ) from e + # Feetech motors only store calibration registers with the torque off. + robot.bus.disable_torque() robot.bus.write_calibration(robot.calibration) robot.configure() except Exception as e: diff --git a/lelab/rollout.py b/lelab/rollout.py index fed7dbabb..b5b04c89d 100644 --- a/lelab/rollout.py +++ b/lelab/rollout.py @@ -75,6 +75,13 @@ def validate_rotations(cls, cameras: dict[str, dict[str, Any]]) -> dict[str, dic # UI can present a "rollout time" separate from the multi-second policy # load + bus connect + camera connect setup overhead. _ROLLOUT_START_MARKER = "Rollout setup complete" +# How long Stop waits for a graceful exit before killing. Once the loop runs, +# teardown (back to the start pose, then disconnect) takes about 4s. During +# setup LeRobot only notes the signal and keeps loading the policy, then +# connecting the arm (huggingface/lerobot#4866), before tearing down: killing it +# mid-connect would leave the arm powered. +_STOP_GRACE_S = 5 +_SETUP_STOP_GRACE_S = 60 def _pump_stdout(proc: subprocess.Popen, log_handle) -> None: @@ -344,7 +351,8 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: "--robot.type=so101_follower", f"--robot.port={request.follower_port}", f"--robot.id={follower_id}", - f"--task={request.task}", + # JSON-quoted: draccus decodes a string as YAML, so "step 1: grab it" became a dict. + f"--task={json.dumps(request.task)}", f"--duration={request.duration_s}", *_rollout_inference_args(policy_path), ] @@ -417,13 +425,14 @@ def handle_stop_inference() -> dict[str, Any]: if not inference_active or _inference_proc is None: return {"success": False, "status_code": 409, "message": "No inference is active"} proc = _inference_proc + grace_s = _STOP_GRACE_S if _inference_rollout_started_at is not None else _SETUP_STOP_GRACE_S try: proc.terminate() try: - proc.wait(timeout=5) + proc.wait(timeout=grace_s) except subprocess.TimeoutExpired: - logger.warning("Inference did not exit in 5s; killing") + logger.warning("Inference did not exit in %ss; killing", grace_s) proc.kill() proc.wait() except Exception as exc: diff --git a/lelab/server.py b/lelab/server.py index dbb69ba68..33a7442f9 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -109,13 +109,10 @@ handle_get_cuda_status, handle_get_policy_extra, handle_get_training_extra, - handle_get_wandb_extra, handle_install_policy_extra, handle_install_policy_extra_status, handle_install_training_extra, handle_install_training_extra_status, - handle_install_wandb_extra, - handle_install_wandb_extra_status, warn_if_cuda_mismatch, ) @@ -937,24 +934,6 @@ def install_training_extra_status(): return handle_install_training_extra_status() -@app.get("/system/wandb-extra") -def get_wandb_extra(): - """Return whether the `wandb` package is importable in this lelab process.""" - return handle_get_wandb_extra() - - -@app.post("/system/wandb-extra/install") -def install_wandb_extra(): - """Spawn `pip install wandb` as a background subprocess. No-op if already running.""" - return handle_install_wandb_extra() - - -@app.get("/system/wandb-extra/install-status") -def install_wandb_extra_status(): - """Return current wandb install state plus any pending log lines (drained on read).""" - return handle_install_wandb_extra_status() - - @app.get("/system/policy-extra/{policy_type}") def get_policy_extra(policy_type: str): """Whether the optional LeRobot extra a policy needs (e.g. transformers for diff --git a/lelab/teleoperate.py b/lelab/teleoperate.py index 2f9981851..581a3ab97 100644 --- a/lelab/teleoperate.py +++ b/lelab/teleoperate.py @@ -179,9 +179,12 @@ def handle_start_teleoperation(request: TeleoperateRequest, websocket_manager=No "Make sure it's plugged in and powered on, then try again." ) from e - # Write calibration to motors' memory + # Feetech motors only store calibration registers with the torque off, + # which also opens their EEPROM lock. logger.info("Writing calibration to motors...") + robot.bus.disable_torque() robot.bus.write_calibration(robot.calibration) + teleop_device.bus.disable_torque() teleop_device.bus.write_calibration(teleop_device.calibration) # Connect cameras and configure motors diff --git a/lelab/train.py b/lelab/train.py index 50419ad04..a370869e1 100644 --- a/lelab/train.py +++ b/lelab/train.py @@ -60,14 +60,11 @@ class TrainingRequest(BaseModel): resume: bool = False job_name: str | None = None - # Weights & Biases - wandb_enable: bool = False - wandb_project: str | None = None - wandb_entity: str | None = None - wandb_notes: str | None = None - wandb_run_id: str | None = None - wandb_mode: str | None = "online" - wandb_disable_artifact: bool = False + # Trackio. A run without a Space stays in this machine's Trackio database; + # JobRegistry.start defaults the Space for a signed-in user. + trackio_enable: bool = False + trackio_project: str | None = None + trackio_space_id: str | None = None # Environment / evaluation env_type: str | None = None @@ -154,11 +151,18 @@ def build_training_command( # LeRobot defaults push_to_hub=True and demands --policy.repo_id when so. # submit_to_hf also resets policy.device so the pod picks its own device. is_cloud = job_target is not None and job_target.runner == "hf_cloud" + device = None if is_cloud else request.policy_device or auto_select_torch_device().type if not is_cloud: - cmd.extend(["--policy.device", request.policy_device or auto_select_torch_device().type]) + cmd.extend(["--policy.device", device]) cmd.extend(["--policy.push_to_hub", "true" if request.policy_push_to_hub else "false"]) if request.policy_push_to_hub and request.policy_repo_id: cmd.extend(["--policy.repo_id", request.policy_repo_id]) + # Training takes mixed precision from accelerate alone (huggingface/lerobot#4010); use_amp + # above stays for the saved config, which evaluation and inference read. fp16 matches the + # torch.autocast default they run under on CUDA, and works on pre-Ampere GPUs such as the + # T4. HF Jobs pods are GPUs. + if request.policy_use_amp and (is_cloud or device == "cuda"): + cmd.extend(["--accelerator.mixed_precision", "fp16"]) # GR00T-specific policy flags. These options only exist on the groot config, # so emit them only for --policy.type=groot; sending them to another policy @@ -201,20 +205,14 @@ def build_training_command( if request.job_name: cmd.extend(["--job_name", request.job_name]) - # W&B - cmd.extend(["--wandb.enable", "true" if request.wandb_enable else "false"]) - if request.wandb_enable: - if request.wandb_project: - cmd.extend(["--wandb.project", request.wandb_project]) - if request.wandb_entity: - cmd.extend(["--wandb.entity", request.wandb_entity]) - if request.wandb_notes: - cmd.extend(["--wandb.notes", request.wandb_notes]) - if request.wandb_run_id: - cmd.extend(["--wandb.run_id", request.wandb_run_id]) - if request.wandb_mode: - cmd.extend(["--wandb.mode", request.wandb_mode]) - cmd.extend(["--wandb.disable_artifact", "true" if request.wandb_disable_artifact else "false"]) + # Trackio + cmd.extend(["--trackio.enable", "true" if request.trackio_enable else "false"]) + if request.trackio_enable: + if request.trackio_project: + cmd.extend(["--trackio.project", request.trackio_project]) + if request.trackio_space_id: + # Private only applies when Trackio creates the Space; an existing one keeps its visibility. + cmd.extend(["--trackio.space_id", request.trackio_space_id, "--trackio.private", "true"]) # Env if request.env_type: diff --git a/lelab/utils/devices.py b/lelab/utils/devices.py index 59cda0635..6ba985814 100644 --- a/lelab/utils/devices.py +++ b/lelab/utils/devices.py @@ -10,7 +10,6 @@ import logging import threading -from contextlib import suppress from typing import Any # Recording, teleoperation, inference and replay all drive the follower: each @@ -35,17 +34,13 @@ def safe_disconnect_device(device: Any, logger: logging.Logger, context: str = " def _force_close_device_resources(device: Any, logger: logging.Logger) -> None: """Best-effort release for serial/camera resources after disconnect fails.""" bus = getattr(device, "bus", None) - port_handler = getattr(bus, "port_handler", None) - if port_handler is not None: - with suppress(Exception): - port_handler.clearPort() - with suppress(Exception): - port_handler.is_using = False + if bus is not None: + # SerialMotorsBus.disconnect() closes the port even when disabling the torque fails. try: - port_handler.closePort() - logger.info("Force-closed serial port after disconnect failure") + bus.disconnect() + logger.info("Released serial port after disconnect failure") except Exception as exc: - logger.warning("Failed to force-close serial port after disconnect failure: %s", exc) + logger.warning("Serial port released, but disabling the torque failed: %s", exc) cameras = getattr(device, "cameras", None) if isinstance(cameras, dict): diff --git a/lelab/utils/system.py b/lelab/utils/system.py index ed4023c7e..d45768bd9 100644 --- a/lelab/utils/system.py +++ b/lelab/utils/system.py @@ -21,8 +21,10 @@ import sys import threading import time +from importlib import metadata from typing import Any +from packaging.requirements import Requirement from pydantic import BaseModel logger = logging.getLogger(__name__) @@ -33,9 +35,6 @@ TRAINING_AVAILABLE: bool = importlib.util.find_spec("accelerate") is not None TRAINING_INSTALL_HINT: str = "pip install accelerate" -WANDB_AVAILABLE: bool = importlib.util.find_spec("wandb") is not None -WANDB_INSTALL_HINT: str = "pip install wandb" - def _build_install_cmd(package: str) -> list[str]: """Pick the best installer for the running Python. @@ -159,7 +158,6 @@ def _drain_queue(self) -> None: training_install_manager = InstallManager("accelerate") -wandb_install_manager = InstallManager("wandb") def handle_get_training_extra() -> dict[str, Any]: @@ -174,57 +172,78 @@ def handle_install_training_extra_status() -> dict[str, Any]: return training_install_manager.get_status() -def handle_get_wandb_extra() -> dict[str, Any]: - return {"available": WANDB_AVAILABLE, "install_hint": WANDB_INSTALL_HINT} - - -def handle_install_wandb_extra() -> dict[str, Any]: - return wandb_install_manager.start() - - -def handle_install_wandb_extra_status() -> dict[str, Any]: - return wandb_install_manager.get_status() - - # --------------------------------------------------------------------------- # # Policy extras # --------------------------------------------------------------------------- # # Some LeRobot policies import an optional extra at construction time; training # (or inference) otherwise dies with a buried ImportError once the subprocess is -# already running. Map each such policy to the module we probe and the -# ``pip install lerobot[extra]`` target. Policies not listed (act, vqbet, tdmpc) -# need nothing extra. -POLICY_EXTRAS: dict[str, tuple[str, str]] = { - # policy_type: (probe_module, install_target) - "smolvla": ("transformers", "lerobot[smolvla]"), - "pi0": ("transformers", "lerobot[pi]"), - "pi0_fast": ("transformers", "lerobot[pi]"), - "pi05": ("transformers", "lerobot[pi]"), - "xvla": ("transformers", "lerobot[xvla]"), - "evo1": ("transformers", "lerobot[evo1]"), - "eo1": ("qwen_vl_utils", "lerobot[eo1]"), - "vla_jepa": ("qwen_vl_utils", "lerobot[vla_jepa]"), - "molmoact2": ("peft", "lerobot[molmoact2]"), - "wall_x": ("torchdiffeq", "lerobot[wallx]"), - "multi_task_dit": ("diffusers", "lerobot[multi_task_dit]"), - "fastwam": ("diffusers", "lerobot[fastwam]"), - "lingbot_va": ("diffusers", "lerobot[lingbot_va]"), - "diffusion": ("diffusers", "lerobot[diffusion]"), - # groot pulls in peft (LoRA), diffusers, timm, dm-tree and decord; peft is a - # reliable, cross-platform stand-in for "the lerobot[groot] extra is present". - "groot": ("peft", "lerobot[groot]"), +# already running. Map each such policy to its ``lerobot[extra]``; what the +# extra requires is read from lerobot's own metadata. Policies not listed (act, +# vqbet, tdmpc) need nothing extra. +POLICY_EXTRAS: dict[str, str] = { + "smolvla": "smolvla", + "pi0": "pi", + "pi0_fast": "pi", + "pi05": "pi", + "xvla": "xvla", + "evo1": "evo1", + "eo1": "eo1", + "vla_jepa": "vla_jepa", + "molmoact2": "molmoact2", + "wall_x": "wallx", + "multi_task_dit": "multi_task_dit", + "fastwam": "fastwam", + "lingbot_va": "lingbot_va", + "diffusion": "diffusion", + "groot": "groot", } + +def _unmet_requirements(extra: str) -> list[Requirement]: + """Requirements of ``lerobot[extra]`` this environment lacks or holds at a version out of range. + + A version check, not only an import probe: an environment upgraded from an older LeRobot + can still import a package the new extra pins to a different range. + """ + requires = [Requirement(r) for r in metadata.requires("lerobot") or []] + unmet: list[Requirement] = [] + pending, seen = [extra], set() + while pending: + name = pending.pop() + if name in seen: + continue + seen.add(name) + for req in requires: + # Skip core requirements, whose markers hold with or without an extra. + if ( + req.marker is None + or not req.marker.evaluate({"extra": name}) + or req.marker.evaluate({"extra": ""}) + ): + continue + if req.name == "lerobot": # an extra built from other extras + pending.extend(req.extras) + continue + try: + installed = metadata.version(req.name) + except metadata.PackageNotFoundError: + unmet.append(req) + continue + if not req.specifier.contains(installed, prereleases=True): + unmet.append(req) + return unmet + + # One install manager per install target (lerobot[smolvla] / lerobot[pi] / …), # created lazily so pi0 and pi0_fast share the lerobot[pi] install. _policy_install_managers: dict[str, InstallManager] = {} def _policy_install_manager(policy_type: str) -> InstallManager | None: - spec = POLICY_EXTRAS.get(policy_type) - if spec is None: + extra = POLICY_EXTRAS.get(policy_type) + if extra is None: return None - target = spec[1] + target = f"lerobot[{extra}]" mgr = _policy_install_managers.get(target) if mgr is None: mgr = InstallManager(target) @@ -233,14 +252,14 @@ def _policy_install_manager(policy_type: str) -> InstallManager | None: def handle_get_policy_extra(policy_type: str) -> dict[str, Any]: - """Whether the optional extra a policy needs is importable right now. + """Whether the optional extra a policy needs is installed right now. - Probed live (not cached at import) so a restart after installing is picked + Checked live (not cached at import) so a restart after installing is picked up. Policies that need nothing report ``available`` so the UI never blocks them. """ - spec = POLICY_EXTRAS.get(policy_type) - if spec is None: + extra = POLICY_EXTRAS.get(policy_type) + if extra is None: return { "policy_type": policy_type, "needs_extra": False, @@ -249,16 +268,13 @@ def handle_get_policy_extra(policy_type: str) -> dict[str, Any]: "install_target": "", "install_hint": "", } - probe, target = spec - try: - available = importlib.util.find_spec(probe) is not None - except (ImportError, ValueError): - available = False + target = f"lerobot[{extra}]" + unmet = _unmet_requirements(extra) return { "policy_type": policy_type, "needs_extra": True, - "available": available, - "package": probe, + "available": not unmet, + "package": ", ".join(f"{req.name}{req.specifier}" for req in unmet), "install_target": target, "install_hint": f"pip install '{target}'", } diff --git a/pyproject.toml b/pyproject.toml index 84edd1a84..ea356df5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "websockets>=15.0.1", "uvicorn>=0.24.0", "psutil>=5.9.0", - "lerobot[core_scripts,feetech,training] @ git+https://github.com/huggingface/lerobot.git@v0.6.1", + "lerobot[core_scripts,feetech,training] @ git+https://github.com/huggingface/lerobot.git@v0.7.0", # Windows-only: real DirectShow camera names for /available-cameras so the # frontend can match a camera to its browser deviceId (issues #12, #16). # Guarded by try/except at the call site, so its absence degrades gracefully. diff --git a/tests/test_calibrate.py b/tests/test_calibrate.py index 760ef9383..929c3a138 100644 --- a/tests/test_calibrate.py +++ b/tests/test_calibrate.py @@ -139,6 +139,25 @@ def test_get_status_returns_independent_recorded_ranges_snapshot() -> None: } +def test_step_homing_uses_the_serial_bus_api() -> None: + """Autospec checks every bus call against LeRobot's SerialMotorsBus signatures.""" + from unittest.mock import create_autospec + + from lelab.calibrate import CalibrationManager + from lerobot.motors import SerialMotorsBus + + bus = create_autospec(SerialMotorsBus, instance=True) + bus.set_half_turn_homings.return_value = {"shoulder_pan": -12} + mgr = CalibrationManager() + mgr.device = type("Device", (), {"bus": bus})() + + mgr._step_homing() + + bus.disable_torque.assert_called_once_with() + bus.set_operating_mode.assert_called_once_with("position") + assert mgr._homing_offsets == {"shoulder_pan": -12} + + def test_status_polling_does_not_wait_for_recording_bus_read() -> None: """A slow recording read does not block status polling.""" from lelab.calibrate import CalibrationManager @@ -311,16 +330,16 @@ def test_cleanup_device_force_releases_and_clears_when_disconnect_fails() -> Non device handle — otherwise the COM port stays busy and blocks the next run.""" from lelab.calibrate import CalibrationManager - class PortHandler: + class Bus: def __init__(self) -> None: - self.closed = False + self.released = False - def closePort(self) -> None: # noqa: N802 - mirrors LeRobot port handler API - self.closed = True + def disconnect(self) -> None: + self.released = True class Device: def __init__(self) -> None: - self.bus = type("Bus", (), {"port_handler": PortHandler()})() + self.bus = Bus() def disconnect(self) -> None: raise RuntimeError("Failed to write 'Torque_Enable' on id_=6") @@ -331,5 +350,5 @@ def disconnect(self) -> None: mgr._cleanup_device() - assert device.bus.port_handler.closed is True # force-released despite failure + assert device.bus.released is True # force-released despite failure assert mgr.device is None # handle cleared so a new calibration can start diff --git a/tests/test_devices.py b/tests/test_devices.py index eab3e209b..b00e1c1ed 100644 --- a/tests/test_devices.py +++ b/tests/test_devices.py @@ -2,21 +2,32 @@ import logging +from lerobot.motors import Motor, MotorNormMode, SerialMotorsBus -def test_safe_disconnect_force_closes_serial_port_after_disconnect_failure() -> None: - from lelab.utils.devices import safe_disconnect_device - class PortHandler: - def __init__(self) -> None: - self.cleared = False - self.closed = False - self.is_using = True +class _UnresponsiveRustypotBus: + """Stands in for an open rustypot.Bus whose motors no longer answer.""" + + def __init__(self) -> None: + self.closed = False + + def set_torque(self, ids, enabled, retries=0): + return dict.fromkeys(ids, "no status packet") + + def close(self) -> None: + self.closed = True + - def clearPort(self) -> None: # noqa: N802 - mirrors LeRobot port handler API - self.cleared = True +def _connected_bus() -> SerialMotorsBus: + bus = SerialMotorsBus( + port="/dev/null", motors={"gripper": Motor(6, "sts3215", MotorNormMode.RANGE_0_100)} + ) + bus._bus = _UnresponsiveRustypotBus() + return bus - def closePort(self) -> None: # noqa: N802 - mirrors LeRobot port handler API - self.closed = True + +def test_safe_disconnect_releases_serial_port_after_disconnect_failure() -> None: + from lelab.utils.devices import safe_disconnect_device class Camera: def __init__(self) -> None: @@ -27,34 +38,27 @@ def disconnect(self) -> None: class Device: def __init__(self) -> None: - self.bus = type("Bus", (), {"port_handler": PortHandler()})() + self.bus = _connected_bus() self.cameras = {"cam": Camera()} def disconnect(self) -> None: - raise RuntimeError("Failed to write 'Torque_Enable' on id_=6") + raise RuntimeError("camera is not connected") device = Device() + port = device.bus._bus safe_disconnect_device(device, logging.getLogger(__name__)) - assert device.bus.port_handler.cleared is True - assert device.bus.port_handler.is_using is False - assert device.bus.port_handler.closed is True + assert port.closed is True # released even though disabling the torque failed + assert device.bus.is_connected is False assert device.cameras["cam"].disconnected is True def test_safe_disconnect_uses_normal_disconnect_when_it_succeeds() -> None: from lelab.utils.devices import safe_disconnect_device - class PortHandler: - def __init__(self) -> None: - self.closed = False - - def closePort(self) -> None: # noqa: N802 - mirrors LeRobot port handler API - self.closed = True - class Device: def __init__(self) -> None: - self.bus = type("Bus", (), {"port_handler": PortHandler()})() + self.bus = _connected_bus() self.disconnected = False def disconnect(self) -> None: @@ -64,4 +68,4 @@ def disconnect(self) -> None: safe_disconnect_device(device, logging.getLogger(__name__)) assert device.disconnected is True - assert device.bus.port_handler.closed is False + assert device.bus.is_connected is True # the fallback did not run diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 22f42d1a2..bfc68138d 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -21,18 +21,66 @@ import pytest -def test_extract_wandb_run_url_finds_canonical_url() -> None: - from lelab.jobs import extract_wandb_run_url +def test_extract_tracking_url_finds_the_space_url() -> None: + from lelab.jobs import extract_tracking_url - line = "wandb: \U0001f680 View run at https://wandb.ai/me/myproj/runs/abc123 trailing text" - assert extract_wandb_run_url(line) == "https://wandb.ai/me/myproj/runs/abc123" + line = "INFO 2026-11-05 15:03:26 io_utils.py:101 Track this run --> https://huggingface.co/spaces/me/lelab-trackio" + assert extract_tracking_url(line) == "https://huggingface.co/spaces/me/lelab-trackio" -def test_extract_wandb_run_url_returns_none_when_absent() -> None: - from lelab.jobs import extract_wandb_run_url +def test_extract_tracking_url_ignores_a_local_database_run() -> None: + from lelab.jobs import extract_tracking_url - assert extract_wandb_run_url("nothing here") is None - assert extract_wandb_run_url("https://example.com/runs/abc") is None + assert extract_tracking_url('Track this run --> trackio show --project "lerobot"') is None + assert extract_tracking_url("https://huggingface.co/spaces/me/other") is None + + +def test_job_record_reads_a_legacy_wandb_run_url() -> None: + from lelab.jobs import JobRecord + from lelab.train import TrainingRequest + + record = JobRecord.model_validate( + { + "id": "j", + "name": "n", + "state": "done", + "config": TrainingRequest(dataset_repo_id="x").model_dump(), + "output_dir": "/tmp/out", + "started_at": 0.0, + "wandb_run_url": "https://wandb.ai/me/p/runs/abc", + } + ) + assert record.tracking_url == "https://wandb.ai/me/p/runs/abc" + assert "wandb_run_url" not in record.model_dump() + + +@pytest.mark.parametrize( + ("whoami", "space_id", "expected"), + [ + ({"name": "me"}, None, "me/lelab-trackio"), + ({"name": "me"}, "team/runs", "team/runs"), # an explicit Space wins + (None, None, None), # signed out: the run stays in the local database + ], +) +def test_tracked_runs_default_to_the_users_trackio_space(monkeypatch, whoami, space_id, expected) -> None: + from lelab import jobs + from lelab.train import TrainingRequest + from lelab.utils import hf_auth + + monkeypatch.setattr(hf_auth, "cached_whoami", lambda: whoami) + config = TrainingRequest(dataset_repo_id="x", trackio_enable=True, trackio_space_id=space_id) + + assert jobs._with_default_trackio_space(config).trackio_space_id == expected + + +def test_untracked_runs_get_no_trackio_space(monkeypatch) -> None: + from lelab import jobs + from lelab.train import TrainingRequest + from lelab.utils import hf_auth + + monkeypatch.setattr(hf_auth, "cached_whoami", lambda: {"name": "me"}) + + assert jobs._with_default_trackio_space(TrainingRequest(dataset_repo_id="x")).trackio_space_id is None def test_parse_duration_handles_mm_ss_and_hh_mm_ss() -> None: diff --git a/tests/test_lerobot_cli.py b/tests/test_lerobot_cli.py new file mode 100644 index 000000000..853cdf058 --- /dev/null +++ b/tests/test_lerobot_cli.py @@ -0,0 +1,132 @@ +# 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. +"""The argv LeLab builds, run through the installed LeRobot's own CLI parsers. + +The argv-shape tests elsewhere cannot see a flag LeRobot renamed or dropped: +they all passed while every training run died on `--wandb.enable`. +""" + +# No `from __future__ import annotations`: parser.wrap reads the config class off the annotation. +import re +import sys +from pathlib import Path + +import pytest + +import lerobot.scripts.lerobot_train # noqa: F401 registers the policy configs, as the CLI does +from lerobot.configs import parser +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.train import TrainPipelineConfig + +_ESSENTIALS_CARD = Path(__file__).parent.parent / "frontend/src/components/training/config/EssentialsCard.tsx" + + +@parser.wrap() +def _parse_training(cfg: TrainPipelineConfig) -> TrainPipelineConfig: + return cfg + + +def test_every_policy_the_ui_offers_is_a_lerobot_policy_type() -> None: + offered = set(re.findall(r'', _ESSENTIALS_CARD.read_text())) + assert offered, "policy list not found in EssentialsCard.tsx" + assert offered - set(PreTrainedConfig.get_known_choices()) == set() + + +@pytest.mark.parametrize( + ("policy_type", "runner"), [("act", "local"), ("act", "hf_cloud"), ("groot", "local")] +) +def test_training_argv_parses(monkeypatch, policy_type, runner) -> None: + from lelab.jobs import JobTarget + from lelab.train import TrainingRequest, build_training_command + + request = TrainingRequest( + dataset_repo_id="me/dataset", + policy_type=policy_type, + policy_device="cpu", + policy_use_amp=True, + trackio_enable=True, + trackio_project="lelab", + trackio_space_id="me/lelab-trackio", + optimizer_lr=1e-4, + **( + {"policy_embodiment_tag": "new_embodiment", "policy_relative_exclude_joints": ["gripper"]} + if policy_type == "groot" + else {} + ), + ) + target = JobTarget(runner=runner, flavor="a10g-small") if runner == "hf_cloud" else None + cmd = build_training_command(request, "/tmp/out", job_target=target) + monkeypatch.setattr(sys, "argv", ["lerobot-train", *cmd[3:]]) + + cfg = _parse_training() + + assert cfg.policy.type == policy_type + assert cfg.trackio.space_id == "me/lelab-trackio" + + +def test_rollout_argv_parses(monkeypatch, tmp_path) -> None: + import lerobot.scripts.lerobot_rollout as lerobot_rollout + from lelab import rollout + from lelab.rollout import InferenceRequest + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.act.configuration_act import ACTConfig + + checkpoint = tmp_path / "checkpoints" / "000050" / "pretrained_model" + checkpoint.mkdir(parents=True) + ACTConfig( + input_features={ + "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(6,)), + "observation.images.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 480, 640)), + }, + output_features={"action": PolicyFeature(type=FeatureType.ACTION, shape=(6,))}, + )._save_pretrained(checkpoint) + + commands: list[list[str]] = [] + + def capture(cmd, **kwargs): + commands.append(cmd) + raise OSError("stop before spawning") + + monkeypatch.setattr(rollout, "inference_active", False) + monkeypatch.setattr(rollout, "setup_follower_calibration_file", lambda name: "robot_a") + monkeypatch.setattr(rollout.subprocess, "Popen", capture) + rollout.handle_start_inference( + InferenceRequest( + follower_port="/dev/ttyUSB0", + follower_config="robot_a", + policy_ref=str(checkpoint), + task="step 1: grab the cube", + cameras={"front": {"type": "opencv", "camera_index": 0, "width": 640, "height": 480, "fps": 30}}, + ) + ) + + parsed = [] + + class _StopBeforeHardwareError(Exception): + pass + + def stop_before_hardware(cfg, _shutdown_event): + parsed.append(cfg) + raise _StopBeforeHardwareError + + monkeypatch.setattr(lerobot_rollout, "build_rollout_context", stop_before_hardware) + monkeypatch.setattr(sys, "argv", ["lerobot-rollout", *commands[0][3:]]) + with pytest.raises(_StopBeforeHardwareError): + lerobot_rollout.main() + + (cfg,) = parsed + assert cfg.robot.port == "/dev/ttyUSB0" + assert cfg.robot.id == "robot_a" + assert cfg.task == "step 1: grab the cube" + assert set(cfg.robot.cameras) == {"front"} diff --git a/tests/test_rollout.py b/tests/test_rollout.py index a34196bc9..c533f11ce 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -253,6 +253,58 @@ def boom(*args, **kwargs): assert rollout.inference_active is False +@pytest.mark.parametrize("task", ['put "it" in: the box', "step 1: grab the cube", "yes", "null", ""]) +def test_handle_start_inference_passes_the_task_verbatim(monkeypatch, tmp_path, task) -> None: + """draccus decodes a string flag as YAML, so the task must reach it quoted.""" + from dataclasses import dataclass + + import draccus + + from lelab import rollout + + monkeypatch.setattr(rollout, "setup_follower_calibration_file", lambda cfg: "robot_a") + monkeypatch.setattr(rollout, "_resolve_policy_path", lambda ref: str(tmp_path / "model")) + commands: list[list[str]] = [] + + def capture(cmd, **kwargs): + commands.append(cmd) + raise OSError("stop before spawning") + + monkeypatch.setattr(rollout.subprocess, "Popen", capture) + rollout.handle_start_inference(_stub_request().model_copy(update={"task": task})) + + @dataclass + class RolloutTask: + task: str = "" + + task_arg = next(arg for arg in commands[0] if arg.startswith("--task=")) + assert draccus.parse(RolloutTask, args=[task_arg]).task == task + + +@pytest.mark.parametrize( + ("rollout_started_at", "grace"), [(None, "_SETUP_STOP_GRACE_S"), (123.0, "_STOP_GRACE_S")] +) +def test_stop_gives_setup_longer_to_exit(monkeypatch, rollout_started_at, grace) -> None: + """Killed mid-setup, LeRobot may be connecting the arm and leave it powered.""" + from lelab import rollout + + waits: list[float] = [] + + class Proc: + def terminate(self) -> None: + pass + + def wait(self, timeout=None) -> None: + waits.append(timeout) + + monkeypatch.setattr(rollout, "inference_active", True) + monkeypatch.setattr(rollout, "_inference_proc", Proc()) + monkeypatch.setattr(rollout, "_inference_rollout_started_at", rollout_started_at) + + assert rollout.handle_stop_inference()["success"] is True + assert waits == [getattr(rollout, grace)] + + def test_classify_outcome_ok_warns_and_fails() -> None: from lelab.rollout import _classify_outcome diff --git a/tests/test_train.py b/tests/test_train.py index 3d99988f9..6aa4fb90f 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -60,27 +60,27 @@ def test_optional_dataset_fields_only_present_when_set() -> None: assert cmd2[idx + 1 : idx + 4] == ["0", "1", "2"] -def test_wandb_block_only_serialized_when_enabled() -> None: +def test_trackio_block_only_serialized_when_enabled() -> None: from lelab.train import TrainingRequest, build_training_command - off = build_training_command(TrainingRequest(dataset_repo_id="x", wandb_enable=False), "/tmp/out") - assert _arg_value(off, "--wandb.enable") == "false" - assert "--wandb.project" not in off + off = build_training_command(TrainingRequest(dataset_repo_id="x"), "/tmp/out") + assert _arg_value(off, "--trackio.enable") == "false" + assert "--trackio.project" not in off + assert "--trackio.space_id" not in off on = build_training_command( TrainingRequest( dataset_repo_id="x", - wandb_enable=True, - wandb_project="proj", - wandb_entity="me", - wandb_run_id="abc", + trackio_enable=True, + trackio_project="proj", + trackio_space_id="me/lelab-trackio", ), "/tmp/out", ) - assert _arg_value(on, "--wandb.enable") == "true" - assert _arg_value(on, "--wandb.project") == "proj" - assert _arg_value(on, "--wandb.entity") == "me" - assert _arg_value(on, "--wandb.run_id") == "abc" + assert _arg_value(on, "--trackio.enable") == "true" + assert _arg_value(on, "--trackio.project") == "proj" + assert _arg_value(on, "--trackio.space_id") == "me/lelab-trackio" + assert _arg_value(on, "--trackio.private") == "true" def test_push_to_hub_emits_repo_id_only_when_enabled() -> None: @@ -198,6 +198,25 @@ def test_local_run_keeps_explicit_device() -> None: assert _arg_value(cmd, "--policy.device") == "cpu" +@pytest.mark.parametrize( + ("device", "runner", "expected"), + [("cuda", "local", "fp16"), ("cpu", "local", None), ("mps", "local", None), (None, "hf_cloud", "fp16")], +) +def test_amp_trains_in_fp16_on_gpus(device, runner, expected) -> None: + from lelab.jobs import JobTarget + from lelab.train import TrainingRequest, build_training_command + + target = JobTarget(runner=runner, flavor="a10g-small") if runner == "hf_cloud" else None + req = TrainingRequest(dataset_repo_id="x", policy_device=device, policy_use_amp=True) + cmd = build_training_command(req, "/tmp/out", job_target=target) + + assert _arg_value(cmd, "--policy.use_amp") == "true" + if expected is None: + assert "--accelerator.mixed_precision" not in cmd + else: + assert _arg_value(cmd, "--accelerator.mixed_precision") == expected + + def test_groot_policy_flags_are_emitted_only_for_groot() -> None: from lelab.train import TrainingRequest, build_training_command diff --git a/tests/test_utils_system.py b/tests/test_utils_system.py index a6ffa5a43..29ae1e0f0 100644 --- a/tests/test_utils_system.py +++ b/tests/test_utils_system.py @@ -18,6 +18,8 @@ import logging import sys +import pytest + def test_build_install_cmd_contains_pip_and_package() -> None: from lelab.utils.system import _build_install_cmd @@ -130,20 +132,69 @@ def test_cuda_status_endpoint_returns_expected_shape(client) -> None: def test_policy_extra_maps_policies_to_install_targets() -> None: - """smolvla/pi0/pi0_fast/diffusion map to the right probe module + lerobot[extra].""" from lelab.utils.system import handle_get_policy_extra smol = handle_get_policy_extra("smolvla") assert smol["needs_extra"] is True - assert smol["package"] == "transformers" assert smol["install_target"] == "lerobot[smolvla]" assert "lerobot[smolvla]" in smol["install_hint"] - # pi0 and pi0_fast share the lerobot[pi] extra; diffusion uses diffusers. + # pi0, pi0_fast and pi05 share the lerobot[pi] extra. assert handle_get_policy_extra("pi0")["install_target"] == "lerobot[pi]" assert handle_get_policy_extra("pi0_fast")["install_target"] == "lerobot[pi]" - assert handle_get_policy_extra("diffusion")["package"] == "diffusers" - assert handle_get_policy_extra("diffusion")["install_target"] == "lerobot[diffusion]" + assert handle_get_policy_extra("wall_x")["install_target"] == "lerobot[wallx]" + + +def test_every_policy_extra_exists_in_the_installed_lerobot() -> None: + """Catches an extra LeRobot renames or drops on the next version bump.""" + from importlib import metadata + + from packaging.utils import canonicalize_name + + from lelab.utils.system import POLICY_EXTRAS + + provided = {canonicalize_name(e) for e in metadata.metadata("lerobot").get_all("Provides-Extra")} + assert {e for e in POLICY_EXTRAS.values() if canonicalize_name(e) not in provided} == set() + + +def _fake_lerobot_requires(monkeypatch, installed: dict[str, str]) -> None: + from lelab.utils import system + + requires = [ + 'torchcodec<0.18.0; sys_platform == "darwin" or sys_platform != "darwin"', # core requirement + 'transformers<5.23.0,>=5.22.0; extra == "transformers-dep"', + 'num2words<0.6.0,>=0.5.14; extra == "smolvla"', + 'lerobot[transformers-dep]; extra == "smolvla"', + ] + + def version(name: str) -> str: + if name not in installed: + raise system.metadata.PackageNotFoundError(name) + return installed[name] + + monkeypatch.setattr(system.metadata, "requires", lambda _dist: requires) + monkeypatch.setattr(system.metadata, "version", version) + + +@pytest.mark.parametrize( + ("installed", "available", "package"), + [ + ({"transformers": "5.22.0", "num2words": "0.5.14"}, True, ""), + ({"transformers": "5.5.0", "num2words": "0.5.14"}, False, "transformers<5.23.0,>=5.22.0"), + ({"transformers": "5.22.0"}, False, "num2words<0.6.0,>=0.5.14"), + ], +) +def test_policy_extra_checks_versions_through_nested_extras( + monkeypatch, installed, available, package +) -> None: + """A missing core requirement (torchcodec here) is not the extra's to report.""" + from lelab.utils import system + + _fake_lerobot_requires(monkeypatch, installed) + status = system.handle_get_policy_extra("smolvla") + + assert status["available"] is available + assert status["package"] == package def test_policy_extra_core_policy_needs_nothing() -> None: @@ -155,17 +206,6 @@ def test_policy_extra_core_policy_needs_nothing() -> None: assert act["install_target"] == "" -def test_policy_extra_available_reflects_find_spec(monkeypatch) -> None: - import importlib.util - - from lelab.utils import system - - monkeypatch.setattr(importlib.util, "find_spec", lambda name: object()) - assert system.handle_get_policy_extra("smolvla")["available"] is True - monkeypatch.setattr(importlib.util, "find_spec", lambda name: None) - assert system.handle_get_policy_extra("smolvla")["available"] is False - - def test_policy_extra_install_is_noop_for_core_policy() -> None: from lelab.utils.system import handle_install_policy_extra, handle_install_policy_extra_status