diff --git a/frontend/src/components/training/config/AdvancedCard.tsx b/frontend/src/components/training/config/AdvancedCard.tsx index 0556303bf..fee8d443f 100644 --- a/frontend/src/components/training/config/AdvancedCard.tsx +++ b/frontend/src/components/training/config/AdvancedCard.tsx @@ -56,25 +56,31 @@ const AdvancedCard: React.FC = ({ config, updateConfig })
Policy
-
- - -
-
+ {/* An HF Jobs pod picks its own device, so this only applies locally. */} + {config.target.runner === 'local' && ( +
+ + +
+ )} +
{ wandb_enable: false, wandb_mode: "online", wandb_disable_artifact: false, - policy_device: "cuda", policy_use_amp: false, optimizer_type: "adam", use_policy_training_preset: true, diff --git a/lelab/rollout.py b/lelab/rollout.py index 09c2a8d5a..52db4ced2 100644 --- a/lelab/rollout.py +++ b/lelab/rollout.py @@ -37,6 +37,8 @@ from pydantic import BaseModel, field_validator +from lerobot.utils.device_utils import auto_select_torch_device + from .camera_config import lerobot_camera_settings, validate_camera_rotations from .utils.config import setup_follower_calibration_file @@ -102,20 +104,6 @@ def _pump_stdout(proc: subprocess.Popen, log_handle) -> None: log_handle.close() -def _detect_device() -> str: - """cuda → mps → cpu, picked once at start time.""" - try: - import torch - - if torch.cuda.is_available(): - return "cuda" - if torch.backends.mps.is_available(): - return "mps" - except Exception: - pass - return "cpu" - - def _resolve_policy_path(policy_ref: str) -> str: """Turn a checkpoints API ref into a local path that lerobot accepts. @@ -345,7 +333,7 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: "lerobot.scripts.lerobot_rollout", "--strategy.type=base", f"--policy.path={policy_path}", - f"--policy.device={_detect_device()}", + f"--policy.device={auto_select_torch_device().type}", "--robot.type=so101_follower", f"--robot.port={request.follower_port}", f"--robot.id={follower_id}", diff --git a/lelab/train.py b/lelab/train.py index 637e304dc..50419ad04 100644 --- a/lelab/train.py +++ b/lelab/train.py @@ -24,6 +24,8 @@ from pydantic import BaseModel +from lerobot.utils.device_utils import auto_select_torch_device + if TYPE_CHECKING: from lelab.jobs import JobTarget @@ -75,7 +77,7 @@ class TrainingRequest(BaseModel): eval_use_async_envs: bool = False # Policy-specific - policy_device: str | None = "cuda" + policy_device: str | None = None policy_use_amp: bool = False # Hub upload (set by HfCloudJobRunner; not exposed in the form) policy_push_to_hub: bool = False @@ -145,15 +147,15 @@ def build_training_command( cmd.extend(["--seed", str(request.seed)]) # Policy device / AMP / hub - if request.policy_device: - cmd.extend(["--policy.device", request.policy_device]) cmd.extend(["--policy.use_amp", "true" if request.policy_use_amp else "false"]) # On HF Cloud, lerobot's submit_to_hf owns the model repo and sets push_to_hub on # the pod itself; _pod_forwarded_args drops any --policy.push_to_hub/--policy.repo_id # we'd pass, so we must not emit them. Local runs keep the existing behavior: # 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" if not is_cloud: + cmd.extend(["--policy.device", request.policy_device or auto_select_torch_device().type]) 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]) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 746ba90f3..a34196bc9 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -61,50 +61,6 @@ def test_inference_request_has_expected_defaults() -> None: assert req.duration_s == 60 -def test_detect_device_returns_cpu_when_neither_cuda_nor_mps(monkeypatch: pytest.MonkeyPatch) -> None: - import torch - - from lelab.rollout import _detect_device - - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False) - assert _detect_device() == "cpu" - - -def test_detect_device_prefers_cuda_over_mps(monkeypatch: pytest.MonkeyPatch) -> None: - import torch - - from lelab.rollout import _detect_device - - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) - assert _detect_device() == "cuda" - - -def test_detect_device_falls_back_to_mps_when_no_cuda(monkeypatch: pytest.MonkeyPatch) -> None: - import torch - - from lelab.rollout import _detect_device - - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) - assert _detect_device() == "mps" - - -def test_detect_device_returns_cpu_when_torch_probe_raises(monkeypatch: pytest.MonkeyPatch) -> None: - """The function wraps both probes in a broad try/except — if torch is - broken at runtime we still need a sensible fallback.""" - import torch - - from lelab.rollout import _detect_device - - def _boom() -> bool: - raise RuntimeError("simulated torch.cuda failure") - - monkeypatch.setattr(torch.cuda, "is_available", _boom) - assert _detect_device() == "cpu" - - def test_resolve_policy_path_returns_local_dir_unchanged(tmp_path) -> None: from lelab.rollout import _resolve_policy_path diff --git a/tests/test_train.py b/tests/test_train.py index be8a65377..3d99988f9 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -137,7 +137,9 @@ def test_cloud_target_emits_job_flags_and_skips_push_to_hub() -> None: # push_to_hub is requested, but for a cloud target lerobot's submit_to_hf owns the # model repo and _pod_forwarded_args drops --policy.* — so we must NOT emit them. - req = TrainingRequest(dataset_repo_id="x", policy_push_to_hub=True, policy_repo_id="me/x") + req = TrainingRequest( + dataset_repo_id="x", policy_push_to_hub=True, policy_repo_id="me/x", policy_device="mps" + ) cmd = build_training_command( req, "/tmp/out", job_target=JobTarget(runner="hf_cloud", flavor="a10g-small") ) @@ -145,6 +147,8 @@ def test_cloud_target_emits_job_flags_and_skips_push_to_hub() -> None: assert _arg_value(cmd, "--job.tags") == '["lelab"]' assert "--policy.push_to_hub" not in cmd assert "--policy.repo_id" not in cmd + # submit_to_hf resets policy.device so the pod picks its own device. + assert "--policy.device" not in cmd # An absolute host output_dir would be baked into the staged config and crash the # pod (mkdir /Users ...); checkpoints go to the Hub repo, so it must be omitted. assert "--output_dir" not in cmd @@ -177,6 +181,23 @@ def test_local_target_keeps_push_to_hub() -> None: assert "--job.target" not in cmd +def test_local_run_uses_detected_device_by_default(monkeypatch: pytest.MonkeyPatch) -> None: + import torch + + import lelab.train as train + + monkeypatch.setattr(train, "auto_select_torch_device", lambda: torch.device("mps")) + cmd = train.build_training_command(train.TrainingRequest(dataset_repo_id="x"), "/tmp/out") + assert _arg_value(cmd, "--policy.device") == "mps" + + +def test_local_run_keeps_explicit_device() -> None: + from lelab.train import TrainingRequest, build_training_command + + cmd = build_training_command(TrainingRequest(dataset_repo_id="x", policy_device="cpu"), "/tmp/out") + assert _arg_value(cmd, "--policy.device") == "cpu" + + def test_groot_policy_flags_are_emitted_only_for_groot() -> None: from lelab.train import TrainingRequest, build_training_command