| """Bucket persistence for the zero-GPU training Space. |
| |
| Only *weights and optimizer/trainer state* are persisted to an Hugging Face |
| Storage Bucket. Training data (the uploaded "trace") is never stored — it is |
| parsed into an in-memory ``datasets`` object and discarded after the run. |
| |
| The bucket layout for one experiment (one model + one method) is:: |
| |
| <checkpoint_id>/ |
| config.json, tokenizer.json, ... # model + LoRA adapter weights |
| trainer_state.json # global step, LR schedule, log |
| optimizer.pt # Adam / AdamW moments <-- the key |
| scheduler.pt # LR scheduler state |
| rng.pt # RNG seed state |
| scheduler.pt / trainer_state.json ... # everything HF Trainer saves |
| |
| ``trainer_state.json`` records ``global_step`` so the next call knows exactly |
| where to resume. ``state.json`` is a tiny index of the most recent checkpoint |
| so the UI can show progress and resume without scanning the bucket. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import tempfile |
| from pathlib import Path |
|
|
| from huggingface_hub import ( |
| batch_bucket_files, |
| create_bucket, |
| download_bucket_files, |
| list_bucket_tree, |
| ) |
|
|
| |
| |
| PERSIST_GLOB = ( |
| "config.json", |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "special_tokens_map.json", |
| "tokenizer.model", |
| "generation_config.json", |
| "adapter_config.json", |
| "adapter_model.safetensors", |
| "adapter_model.bin", |
| "model.safetensors", |
| "model.safetensors.index.json", |
| "pytorch_model.bin", |
| "trainer_state.json", |
| "optimizer.pt", |
| "scheduler.pt", |
| "rng.pt", |
| "training_args.bin", |
| ) |
|
|
|
|
| def _token() -> str | None: |
| return os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") |
|
|
|
|
| def ensure_bucket(bucket_id: str) -> None: |
| """Create the bucket if it does not exist. Namespace-qualified id is fine |
| (``username/my-bucket``); a bare name is created in the current user's |
| namespace.""" |
| try: |
| create_bucket(bucket_id=bucket_id, exist_ok=True, token=_token()) |
| except Exception: |
| |
| pass |
|
|
|
|
| def checkpoint_paths(bucket_id: str) -> list[str]: |
| """Return the sorted list of ``<checkpoint_id>`` prefixes currently in the |
| bucket (i.e. every completed, persisted step).""" |
| paths: list[str] = set() |
| try: |
| for entry in list_bucket_tree(bucket_id=bucket_id, prefix="", token=_token()): |
| p = getattr(entry, "path", "") or "" |
| |
| |
| first = p.split("/", 1)[0] |
| if first.startswith("checkpoint-"): |
| paths.add(first) |
| elif p == "checkpoint-1" or first == "checkpoint-1": |
| paths.add(p) |
| except Exception: |
| pass |
| return sorted(paths, key=_step_of) |
|
|
|
|
| def _step_of(prefix: str) -> int: |
| try: |
| return int(prefix.split("-")[-1]) |
| except Exception: |
| return -1 |
|
|
|
|
| def list_latest_checkpoint(bucket_id: str) -> str | None: |
| paths = checkpoint_paths(bucket_id) |
| return paths[-1] if paths else None |
|
|
|
|
| def push_checkpoint(bucket_id: str, local_dir: Path, checkpoint_id: str) -> bool: |
| """Upload every weight/state file under ``local_dir`` into |
| ``<bucket_id>/<checkpoint_id>/``. Returns True on success.""" |
| add: list[tuple[Path, str]] = [] |
| for path in local_dir.rglob("*"): |
| if path.is_file(): |
| dest = f"{checkpoint_id}/{path.relative_to(local_dir).as_posix()}" |
| add.append((path, dest)) |
| if not add: |
| return False |
| try: |
| batch_bucket_files(bucket_id, add=add, token=_token()) |
| return True |
| except Exception: |
| return False |
|
|
|
|
| def pull_checkpoint(bucket_id: str, checkpoint_id: str, local_dir: Path) -> bool: |
| """Download a persisted checkpoint from the bucket back into ``local_dir``. |
| |
| The *source* paths passed to ``download_bucket_files`` must be the full |
| bucket-relative paths (``<checkpoint_id>/<relative>``); the local *dest* |
| strips the prefix. Returns True if at least one file was retrieved. |
| |
| Xet-backed buckets can have read-after-write latency, so the whole |
| list→download→verify cycle is retried a few times before giving up. |
| """ |
| import time |
|
|
| for _attempt in range(8): |
| try: |
| entries = list( |
| list_bucket_tree(bucket_id=bucket_id, prefix=f"{checkpoint_id}/", token=_token()) |
| ) |
| except Exception: |
| entries = [] |
| files: list[tuple[str, str]] = [] |
| for e in entries: |
| full = getattr(e, "path", "") or "" |
| if full.startswith(f"{checkpoint_id}/"): |
| rel = full[len(f"{checkpoint_id}/"):] |
| files.append((full, f"{local_dir}/{rel}")) |
| if files: |
| try: |
| |
| local_dir.mkdir(parents=True, exist_ok=True) |
| for _src, dest in files: |
| Path(dest).parent.mkdir(parents=True, exist_ok=True) |
| download_bucket_files(bucket_id, files=files, token=_token()) |
| except Exception: |
| files = [] |
| |
| for _src, dest in files: |
| if Path(dest).exists(): |
| return True |
| |
| time.sleep(1.5) |
| return False |
|
|
|
|
| def save_state_index(bucket_id: str, state: dict) -> None: |
| """Persist a tiny JSON index of the run so the UI can resume without a full |
| tree scan. ``state`` e.g. ``{"global_step": 12, "loss": 0.42, "method": "dpo"}``.""" |
| content = json.dumps(state, indent=2).encode() |
| batch_bucket_files(bucket_id, add=[(content, "state.json")], token=_token()) |
|
|
|
|
| def load_state_index(bucket_id: str) -> dict: |
| local = Path(tempfile.gettempdir()) / "zero_gpu_state.json" |
| try: |
| download_bucket_files(bucket_id, files=[("state.json", str(local))], token=_token()) |
| except Exception: |
| return {} |
| try: |
| return json.loads(local.read_text()) |
| except Exception: |
| return {} |
|
|
|
|
| def save_job(bucket_id: str, job) -> None: |
| """Persist a queued job to the bucket (opt-in queue persistence). |
| |
| Writes ``queue/<job_id>.json`` so a queued request survives a Space |
| restart. ``job`` is a ``queue.Job`` (duck-typed for ``job_id`` / ``method`` |
| / ``rows``). Best-effort: never raises into the caller. |
| """ |
| try: |
| content = json.dumps( |
| { |
| "job_id": getattr(job, "job_id", ""), |
| "method": getattr(job, "method", ""), |
| "rows": getattr(job, "rows", []), |
| "submitted_at": getattr(job, "submitted_at", 0.0), |
| }, |
| indent=2, |
| ).encode() |
| batch_bucket_files( |
| bucket_id, |
| add=[(content, f"queue/{getattr(job, 'job_id', 'job')}.json")], |
| token=_token(), |
| ) |
| except Exception: |
| pass |
|
|
|
|
| def delete_job(bucket_id: str, job_id: str) -> None: |
| """Remove a completed job from the bucket queue (opt-in). Best-effort.""" |
| try: |
| from huggingface_hub import delete_bucket_file |
|
|
| delete_bucket_file(bucket_id, f"queue/{job_id}.json", token=_token()) |
| except Exception: |
| pass |
|
|
