mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +00:00
migration: make per-rank resume prefetch resilient to Hub timeouts
Retry the destination list_repo_files with backoff (and fall back to an empty resume set) so a single transient huggingface.co read timeout no longer fails an entire SLURM rank and skips its dataset slice.
This commit is contained in:
@@ -73,6 +73,7 @@ class MigrateShard(PipelineStep):
|
|||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
import sys
|
import sys
|
||||||
|
import time
|
||||||
import traceback
|
import traceback
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -101,11 +102,22 @@ class MigrateShard(PipelineStep):
|
|||||||
manifest = Path(self.manifest_dir) / f"manifest_{rank:05d}.csv"
|
manifest = Path(self.manifest_dir) / f"manifest_{rank:05d}.csv"
|
||||||
|
|
||||||
api = HfApi()
|
api = HfApi()
|
||||||
dst_files = (
|
# Resume prefetch. A single transient Hub read timeout here must NOT kill the whole
|
||||||
set()
|
# rank (and skip its entire dataset slice), so retry with backoff and, as a last
|
||||||
if (self.only_classify or self.no_push)
|
# resort, fall back to an empty set (already-present datasets are re-checked per item
|
||||||
else set(api.list_repo_files(self.dst_repo, repo_type="dataset"))
|
# and, for --source molmoact, were already filtered out on the submit node).
|
||||||
)
|
dst_files: set = set()
|
||||||
|
if not self.only_classify and not self.no_push:
|
||||||
|
for attempt in range(5):
|
||||||
|
try:
|
||||||
|
dst_files = set(api.list_repo_files(self.dst_repo, repo_type="dataset"))
|
||||||
|
break
|
||||||
|
except Exception as e:
|
||||||
|
if attempt == 4:
|
||||||
|
logging.warning(f"Rank {rank}: could not list {self.dst_repo} after 5 "
|
||||||
|
f"tries ({e}); proceeding without a resume set.")
|
||||||
|
else:
|
||||||
|
time.sleep(5 * (attempt + 1))
|
||||||
|
|
||||||
fieldnames = sorted(
|
fieldnames = sorted(
|
||||||
{"root", "robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action"}
|
{"root", "robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action"}
|
||||||
|
|||||||
Reference in New Issue
Block a user