dejavu-client 0.1.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
.visualizer-credentials
|
|
2
|
+
node_modules/
|
|
3
|
+
.next/
|
|
4
|
+
episode_index_out/
|
|
5
|
+
manifests_out/
|
|
6
|
+
|
|
7
|
+
# Local Netlify folder
|
|
8
|
+
.netlify
|
|
9
|
+
|
|
10
|
+
# Python
|
|
11
|
+
__pycache__/
|
|
12
|
+
*.py[cod]
|
|
13
|
+
*$py.class
|
|
14
|
+
.venv/
|
|
15
|
+
venv/
|
|
16
|
+
.pytest_cache/
|
|
17
|
+
.ruff_cache/
|
|
18
|
+
.mypy_cache/
|
|
19
|
+
*.egg-info/
|
|
20
|
+
dist/
|
|
21
|
+
build/
|
|
22
|
+
|
|
23
|
+
# Python build artifacts
|
|
24
|
+
client/dist/
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dejavu-client
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: Client for the Deja Vu robot-data platform: curation manifests, checkpoints, artifacts, and rollout results
|
|
5
|
+
Project-URL: Homepage, https://github.com/tabtabtabai/data-visualizer
|
|
6
|
+
Project-URL: Issues, https://github.com/tabtabtabai/data-visualizer/issues
|
|
7
|
+
Author: TabTabTab
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
Keywords: checkpoints,datasets,lerobot,machine-learning,robotics
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
14
|
+
Requires-Python: >=3.10
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
|
|
17
|
+
# dejavu-client
|
|
18
|
+
|
|
19
|
+
Python client for the **Deja Vu** robot-data platform — curation manifests,
|
|
20
|
+
checkpoints, cached artifacts, and inference/sim rollout results.
|
|
21
|
+
|
|
22
|
+
Stdlib only, no dependencies: it installs identically on a laptop, a GCP or
|
|
23
|
+
Azure VM, Baseten, or inside a Dockerfile with no credentials configured.
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
pip install dejavu-client
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
## Use
|
|
30
|
+
|
|
31
|
+
```python
|
|
32
|
+
from dejavu_client import DejaVu
|
|
33
|
+
|
|
34
|
+
dv = DejaVu() # DEJAVU_API_KEY, DEJAVU_API_URL from env
|
|
35
|
+
|
|
36
|
+
# what exists, and how curated it is
|
|
37
|
+
for d in dv.datasets():
|
|
38
|
+
print(d["dataset"], d["episodes"], d["flagged"], d["bad"])
|
|
39
|
+
|
|
40
|
+
# episodes this run should skip — step one of any dataloader
|
|
41
|
+
m = dv.manifest("G1_Dex1_FoldTowel")
|
|
42
|
+
skip = set(m["drop"]) # record m["version"] alongside the run
|
|
43
|
+
|
|
44
|
+
# cache something expensive to recompute
|
|
45
|
+
key = dv.artifact_key("G1_Dex1_FoldTowel", 42,
|
|
46
|
+
transform="crop224", encoder="siglip-v1")
|
|
47
|
+
blob = dv.artifact_get(key)
|
|
48
|
+
if blob is None:
|
|
49
|
+
blob = expensive_encode(...)
|
|
50
|
+
dv.artifact_put(key, blob)
|
|
51
|
+
|
|
52
|
+
# save a checkpoint — a set of files, not one
|
|
53
|
+
dv.upload_checkpoint(
|
|
54
|
+
["step1000.pt", "config.yaml", "dataset_statistics.json"],
|
|
55
|
+
wandb_run_id=run.id, step=1000,
|
|
56
|
+
metrics={"val_loss": 0.42},
|
|
57
|
+
manifest={"schema": "unibot-checkpoint/v0", "control_space": "ee"},
|
|
58
|
+
manifest_version=m["version"],
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
# inspect stats without pulling gigabytes of weights
|
|
62
|
+
dv.download_checkpoint(ckpt_id, "./out", only=["dataset_statistics.json"])
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
## Things worth knowing
|
|
66
|
+
|
|
67
|
+
**Bytes never pass through the API.** It hands out presigned URLs and you
|
|
68
|
+
transfer straight to object storage, so a 19 GB checkpoint isn't limited by
|
|
69
|
+
request size or timeouts. Uploads over 5 GB go multipart automatically.
|
|
70
|
+
|
|
71
|
+
**Artifact reads bypass the API entirely** and hit the CDN, so a cache miss
|
|
72
|
+
costs one HTTP request to an edge node rather than a round trip to the
|
|
73
|
+
server. That is also what makes the cache shared across GPU boxes.
|
|
74
|
+
|
|
75
|
+
**Artifact keys are content-addressed** — `hash(dataset, episode, transform,
|
|
76
|
+
encoder)`. Change the recipe and you get a different key, so nothing needs
|
|
77
|
+
invalidating; stale entries are simply never requested again.
|
|
78
|
+
|
|
79
|
+
**A checkpoint is a set of files.** Weights plus `config.yaml` plus
|
|
80
|
+
`dataset_statistics.json`, because a serving stack often needs the small
|
|
81
|
+
ones and shouldn't have to download the weights to read them.
|
|
82
|
+
|
|
83
|
+
**There is no authorization.** A valid key may read and write anything; keys
|
|
84
|
+
exist for attribution and revocation. Keep them out of source control.
|
|
85
|
+
|
|
86
|
+
## Configuration
|
|
87
|
+
|
|
88
|
+
| Variable | Default |
|
|
89
|
+
|---|---|
|
|
90
|
+
| `DEJAVU_API_KEY` | *(required)* |
|
|
91
|
+
| `DEJAVU_API_URL` | `https://visualizer.tabtabtab.ai` |
|
|
92
|
+
|
|
93
|
+
Mint a key from the Deja Vu UI; it is shown once and stored only as a hash.
|
|
94
|
+
|
|
95
|
+
## Note
|
|
96
|
+
|
|
97
|
+
This is a client for a specific hosted service. It is published so that
|
|
98
|
+
ephemeral training boxes across clouds can install it without credentials —
|
|
99
|
+
it is not useful without access to a Deja Vu deployment.
|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+
# dejavu-client
|
|
2
|
+
|
|
3
|
+
Python client for the **Deja Vu** robot-data platform — curation manifests,
|
|
4
|
+
checkpoints, cached artifacts, and inference/sim rollout results.
|
|
5
|
+
|
|
6
|
+
Stdlib only, no dependencies: it installs identically on a laptop, a GCP or
|
|
7
|
+
Azure VM, Baseten, or inside a Dockerfile with no credentials configured.
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install dejavu-client
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Use
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
from dejavu_client import DejaVu
|
|
17
|
+
|
|
18
|
+
dv = DejaVu() # DEJAVU_API_KEY, DEJAVU_API_URL from env
|
|
19
|
+
|
|
20
|
+
# what exists, and how curated it is
|
|
21
|
+
for d in dv.datasets():
|
|
22
|
+
print(d["dataset"], d["episodes"], d["flagged"], d["bad"])
|
|
23
|
+
|
|
24
|
+
# episodes this run should skip — step one of any dataloader
|
|
25
|
+
m = dv.manifest("G1_Dex1_FoldTowel")
|
|
26
|
+
skip = set(m["drop"]) # record m["version"] alongside the run
|
|
27
|
+
|
|
28
|
+
# cache something expensive to recompute
|
|
29
|
+
key = dv.artifact_key("G1_Dex1_FoldTowel", 42,
|
|
30
|
+
transform="crop224", encoder="siglip-v1")
|
|
31
|
+
blob = dv.artifact_get(key)
|
|
32
|
+
if blob is None:
|
|
33
|
+
blob = expensive_encode(...)
|
|
34
|
+
dv.artifact_put(key, blob)
|
|
35
|
+
|
|
36
|
+
# save a checkpoint — a set of files, not one
|
|
37
|
+
dv.upload_checkpoint(
|
|
38
|
+
["step1000.pt", "config.yaml", "dataset_statistics.json"],
|
|
39
|
+
wandb_run_id=run.id, step=1000,
|
|
40
|
+
metrics={"val_loss": 0.42},
|
|
41
|
+
manifest={"schema": "unibot-checkpoint/v0", "control_space": "ee"},
|
|
42
|
+
manifest_version=m["version"],
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
# inspect stats without pulling gigabytes of weights
|
|
46
|
+
dv.download_checkpoint(ckpt_id, "./out", only=["dataset_statistics.json"])
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
## Things worth knowing
|
|
50
|
+
|
|
51
|
+
**Bytes never pass through the API.** It hands out presigned URLs and you
|
|
52
|
+
transfer straight to object storage, so a 19 GB checkpoint isn't limited by
|
|
53
|
+
request size or timeouts. Uploads over 5 GB go multipart automatically.
|
|
54
|
+
|
|
55
|
+
**Artifact reads bypass the API entirely** and hit the CDN, so a cache miss
|
|
56
|
+
costs one HTTP request to an edge node rather than a round trip to the
|
|
57
|
+
server. That is also what makes the cache shared across GPU boxes.
|
|
58
|
+
|
|
59
|
+
**Artifact keys are content-addressed** — `hash(dataset, episode, transform,
|
|
60
|
+
encoder)`. Change the recipe and you get a different key, so nothing needs
|
|
61
|
+
invalidating; stale entries are simply never requested again.
|
|
62
|
+
|
|
63
|
+
**A checkpoint is a set of files.** Weights plus `config.yaml` plus
|
|
64
|
+
`dataset_statistics.json`, because a serving stack often needs the small
|
|
65
|
+
ones and shouldn't have to download the weights to read them.
|
|
66
|
+
|
|
67
|
+
**There is no authorization.** A valid key may read and write anything; keys
|
|
68
|
+
exist for attribution and revocation. Keep them out of source control.
|
|
69
|
+
|
|
70
|
+
## Configuration
|
|
71
|
+
|
|
72
|
+
| Variable | Default |
|
|
73
|
+
|---|---|
|
|
74
|
+
| `DEJAVU_API_KEY` | *(required)* |
|
|
75
|
+
| `DEJAVU_API_URL` | `https://visualizer.tabtabtab.ai` |
|
|
76
|
+
|
|
77
|
+
Mint a key from the Deja Vu UI; it is shown once and stored only as a hash.
|
|
78
|
+
|
|
79
|
+
## Note
|
|
80
|
+
|
|
81
|
+
This is a client for a specific hosted service. It is published so that
|
|
82
|
+
ephemeral training boxes across clouds can install it without credentials —
|
|
83
|
+
it is not useful without access to a Deja Vu deployment.
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "dejavu-client"
|
|
3
|
+
version = "0.1.1"
|
|
4
|
+
description = "Client for the Deja Vu robot-data platform: curation manifests, checkpoints, artifacts, and rollout results"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10"
|
|
7
|
+
license = "MIT"
|
|
8
|
+
authors = [{ name = "TabTabTab" }]
|
|
9
|
+
keywords = ["robotics", "lerobot", "machine-learning", "checkpoints", "datasets"]
|
|
10
|
+
classifiers = [
|
|
11
|
+
"Programming Language :: Python :: 3",
|
|
12
|
+
"License :: OSI Approved :: MIT License",
|
|
13
|
+
"Intended Audience :: Science/Research",
|
|
14
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
15
|
+
]
|
|
16
|
+
|
|
17
|
+
# Deliberately none. Training boxes are ephemeral and span GCP, Azure and
|
|
18
|
+
# Baseten; a stdlib-only client installs identically everywhere with no
|
|
19
|
+
# wheels to build and nothing to conflict with a trainer's pinned stack.
|
|
20
|
+
dependencies = []
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Homepage = "https://github.com/tabtabtabai/data-visualizer"
|
|
24
|
+
Issues = "https://github.com/tabtabtabai/data-visualizer/issues"
|
|
25
|
+
|
|
26
|
+
[build-system]
|
|
27
|
+
requires = ["hatchling"]
|
|
28
|
+
build-backend = "hatchling.build"
|
|
29
|
+
|
|
30
|
+
[tool.hatch.build.targets.wheel]
|
|
31
|
+
packages = ["src/dejavu_client"]
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
"""Client for the Deja Vu robot-data platform.
|
|
2
|
+
|
|
3
|
+
from dejavu_client import DejaVu
|
|
4
|
+
|
|
5
|
+
dv = DejaVu() # DEJAVU_API_KEY=dvk_...
|
|
6
|
+
skip = set(dv.manifest("G1_Dex1_FoldTowel")["drop"])
|
|
7
|
+
|
|
8
|
+
See the DejaVu class docstring for the full contract.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
12
|
+
|
|
13
|
+
from .client import DejaVu
|
|
14
|
+
|
|
15
|
+
__all__ = ["DejaVu", "__version__"]
|
|
16
|
+
|
|
17
|
+
try:
|
|
18
|
+
# Read from installed metadata so the version lives in exactly one place
|
|
19
|
+
# (pyproject.toml) and a release bump cannot leave the two disagreeing.
|
|
20
|
+
__version__ = version("dejavu-client")
|
|
21
|
+
except PackageNotFoundError: # running from a source checkout
|
|
22
|
+
__version__ = "0.0.0.dev0"
|
|
@@ -0,0 +1,279 @@
|
|
|
1
|
+
"""Client for the Deja Vu platform API.
|
|
2
|
+
|
|
3
|
+
Everything a trainer needs, with no dependencies beyond the standard
|
|
4
|
+
library — it installs identically on GCP, Azure, Baseten or a laptop.
|
|
5
|
+
|
|
6
|
+
pip install dejavu-client
|
|
7
|
+
|
|
8
|
+
export DEJAVU_API_KEY=dvk_... # mint one in the Deja Vu UI
|
|
9
|
+
export DEJAVU_API_URL=https://visualizer.tabtabtab.ai
|
|
10
|
+
|
|
11
|
+
from dejavu_client import DejaVu
|
|
12
|
+
dv = DejaVu()
|
|
13
|
+
|
|
14
|
+
# 1. which episodes should this run skip?
|
|
15
|
+
m = dv.manifest("G1_Dex1_FoldTowel")
|
|
16
|
+
skip = set(m["drop"]) # record m["version"] with your run
|
|
17
|
+
|
|
18
|
+
# 2. cache an expensive artifact (latents, preprocessed tensors)
|
|
19
|
+
key = dv.artifact_key("G1_Dex1_FoldTowel", 42, transform="crop224", encoder="siglip-v1")
|
|
20
|
+
blob = dv.artifact_get(key) # None on miss — reads bypass the API
|
|
21
|
+
if blob is None:
|
|
22
|
+
blob = expensive_encode(...)
|
|
23
|
+
dv.artifact_put(key, blob)
|
|
24
|
+
|
|
25
|
+
# 3. save a checkpoint (streams straight to R2, not through the API)
|
|
26
|
+
dv.upload_checkpoint(["step1000.pt", "config.yaml", "dataset_statistics.json"],
|
|
27
|
+
wandb_run_id=run.id, step=1000,
|
|
28
|
+
metrics={"val_loss": 0.42},
|
|
29
|
+
manifest={"schema": "unibot-checkpoint/v0", ...},
|
|
30
|
+
manifest_version=m["version"])
|
|
31
|
+
|
|
32
|
+
Design notes worth knowing:
|
|
33
|
+
* Bytes never pass through the API service — it hands out presigned R2
|
|
34
|
+
URLs and you transfer directly. Cloud Run caps requests at 32 MB.
|
|
35
|
+
* Artifact READS go straight to the CDN with no auth and no API call, so
|
|
36
|
+
a cache miss check costs one HTTP HEAD against an edge node.
|
|
37
|
+
* There is no authorization: any valid key can do anything. Keys exist so
|
|
38
|
+
the request log can attribute writes and so they can be revoked.
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
from __future__ import annotations
|
|
42
|
+
|
|
43
|
+
import hashlib
|
|
44
|
+
import json
|
|
45
|
+
import os
|
|
46
|
+
import urllib.error
|
|
47
|
+
import urllib.request
|
|
48
|
+
from pathlib import Path
|
|
49
|
+
|
|
50
|
+
DEFAULT_URL = "https://visualizer.tabtabtab.ai" # the Python backend
|
|
51
|
+
PUBLIC_CDN = "https://data.tabtabtab.page"
|
|
52
|
+
|
|
53
|
+
# R2 caps a single PUT at 5 GB; our base checkpoints are ~19 GB.
|
|
54
|
+
PART_SIZE = 256 * 1024 * 1024
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class DejaVu:
|
|
58
|
+
def __init__(self, api_key: str | None = None, url: str | None = None):
|
|
59
|
+
self.key = api_key or os.environ["DEJAVU_API_KEY"]
|
|
60
|
+
self.url = (url or os.environ.get("DEJAVU_API_URL", DEFAULT_URL)).rstrip("/")
|
|
61
|
+
|
|
62
|
+
# ---------------------------------------------------------------- http
|
|
63
|
+
def _call(self, method: str, path: str, body: dict | None = None) -> dict:
|
|
64
|
+
req = urllib.request.Request(
|
|
65
|
+
f"{self.url}{path}",
|
|
66
|
+
method=method,
|
|
67
|
+
data=json.dumps(body).encode() if body is not None else None,
|
|
68
|
+
headers={
|
|
69
|
+
"Authorization": f"Bearer {self.key}",
|
|
70
|
+
**({"Content-Type": "application/json"} if body is not None else {}),
|
|
71
|
+
},
|
|
72
|
+
)
|
|
73
|
+
with urllib.request.urlopen(req, timeout=120) as r:
|
|
74
|
+
return json.loads(r.read() or "{}")
|
|
75
|
+
|
|
76
|
+
# ------------------------------------------------------------ curation
|
|
77
|
+
def datasets(self) -> list[dict]:
|
|
78
|
+
"""Datasets the platform knows about, with their curation state.
|
|
79
|
+
|
|
80
|
+
Derived server-side from Postgres, so it cannot drift from reality
|
|
81
|
+
the way a hardcoded list does. Each entry carries episode, flagged,
|
|
82
|
+
reviewed and bad counts.
|
|
83
|
+
"""
|
|
84
|
+
return self._call("GET", "/api/datasets")["datasets"]
|
|
85
|
+
|
|
86
|
+
def manifest(self, dataset: str, strict: bool = False) -> dict:
|
|
87
|
+
"""Episodes to skip. `strict` also drops flagged-but-unreviewed ones."""
|
|
88
|
+
return self._call("GET", f"/api/manifest/{dataset}" + ("?strict=1" if strict else ""))
|
|
89
|
+
|
|
90
|
+
def scores(self, dataset: str) -> dict:
|
|
91
|
+
return self._call("GET", f"/api/scores?dataset={dataset}")
|
|
92
|
+
|
|
93
|
+
def set_verdict(self, dataset: str, episode: int, verdict: str | None, note: str = "") -> dict:
|
|
94
|
+
return self._call("POST", "/api/verdict", {
|
|
95
|
+
"dataset": dataset, "episode": episode, "verdict": verdict, "note": note,
|
|
96
|
+
})
|
|
97
|
+
|
|
98
|
+
# ------------------------------------------------------------ artifacts
|
|
99
|
+
@staticmethod
|
|
100
|
+
def artifact_key(dataset: str, episode: int, transform: str, encoder: str) -> str:
|
|
101
|
+
"""Content-addressed key: change the recipe, get a different key.
|
|
102
|
+
|
|
103
|
+
That is the whole invalidation story — stale entries are simply never
|
|
104
|
+
requested again, so nothing has to be purged.
|
|
105
|
+
"""
|
|
106
|
+
digest = hashlib.sha256(
|
|
107
|
+
f"{dataset}|{episode}|{transform}|{encoder}".encode()
|
|
108
|
+
).hexdigest()[:32]
|
|
109
|
+
return f"latents/{digest}"
|
|
110
|
+
|
|
111
|
+
def artifact_get(self, key: str) -> bytes | None:
|
|
112
|
+
"""Fetch an artifact, or None on a miss.
|
|
113
|
+
|
|
114
|
+
Deliberately reads the CDN directly: no API round trip, and the edge
|
|
115
|
+
caches it, which is what makes this shared across GPU boxes.
|
|
116
|
+
|
|
117
|
+
Not read-after-write consistent — a GET issued immediately after a PUT
|
|
118
|
+
can still miss while the write propagates to the edge. That is fine
|
|
119
|
+
for the intended pattern (on a miss you compute the value and already
|
|
120
|
+
hold it in memory); just don't write a artifact and then read it back
|
|
121
|
+
to obtain it.
|
|
122
|
+
"""
|
|
123
|
+
try:
|
|
124
|
+
with urllib.request.urlopen(f"{PUBLIC_CDN}/derived/artifacts/{key}", timeout=60) as r:
|
|
125
|
+
return r.read()
|
|
126
|
+
except urllib.error.HTTPError as e:
|
|
127
|
+
# R2's public domain answers 403 for a key that isn't there (it
|
|
128
|
+
# will not confirm absence), so a miss is 403 or 404. Treating
|
|
129
|
+
# only 404 as a miss makes every cold cache raise.
|
|
130
|
+
if e.code in (403, 404):
|
|
131
|
+
return None
|
|
132
|
+
raise
|
|
133
|
+
|
|
134
|
+
def artifact_list(self, namespace: str | None = None, limit: int = 100) -> dict:
|
|
135
|
+
"""What's cached. Objects are world-readable but not enumerable from
|
|
136
|
+
outside (no directory index, keys are hashes) — this is our view."""
|
|
137
|
+
q = f"?limit={limit}" + (f"&namespace={namespace}" if namespace else "")
|
|
138
|
+
return self._call("GET", f"/api/artifacts{q}")
|
|
139
|
+
|
|
140
|
+
def artifact_put(self, key: str, blob: bytes) -> str:
|
|
141
|
+
resp = self._call("POST", f"/api/artifacts/{key}")
|
|
142
|
+
_put(resp["upload_url"], blob)
|
|
143
|
+
return resp["read_url"]
|
|
144
|
+
|
|
145
|
+
# ---------------------------------------------------------- checkpoints
|
|
146
|
+
def list_checkpoints(self, wandb_run_id: str | None = None) -> list[dict]:
|
|
147
|
+
q = f"?run={wandb_run_id}" if wandb_run_id else ""
|
|
148
|
+
return self._call("GET", f"/api/checkpoints{q}")["checkpoints"]
|
|
149
|
+
|
|
150
|
+
def download_checkpoint(self, checkpoint_id: str, dest_dir: str | Path,
|
|
151
|
+
only: list[str] | None = None) -> list[Path]:
|
|
152
|
+
"""Download a checkpoint's files. `only` fetches a subset by filename —
|
|
153
|
+
e.g. ["dataset_statistics.json"] to inspect stats without the weights."""
|
|
154
|
+
info = self._call("GET", f"/api/checkpoints/{checkpoint_id}")
|
|
155
|
+
dest_dir = Path(dest_dir)
|
|
156
|
+
dest_dir.mkdir(parents=True, exist_ok=True)
|
|
157
|
+
written = []
|
|
158
|
+
for f in info["files"]:
|
|
159
|
+
if only and f["filename"] not in only:
|
|
160
|
+
continue
|
|
161
|
+
out = dest_dir / f["filename"]
|
|
162
|
+
with urllib.request.urlopen(f["download_url"], timeout=3600) as r, out.open("wb") as fh:
|
|
163
|
+
while chunk := r.read(8 * 1024 * 1024):
|
|
164
|
+
fh.write(chunk)
|
|
165
|
+
written.append(out)
|
|
166
|
+
return written
|
|
167
|
+
|
|
168
|
+
def upload_checkpoint(
|
|
169
|
+
self,
|
|
170
|
+
paths: str | Path | list[str | Path],
|
|
171
|
+
wandb_run_id: str,
|
|
172
|
+
step: int,
|
|
173
|
+
metrics: dict | None = None,
|
|
174
|
+
manifest: dict | None = None,
|
|
175
|
+
manifest_version: str | None = None,
|
|
176
|
+
wandb_url: str | None = None,
|
|
177
|
+
) -> str:
|
|
178
|
+
"""Upload a checkpoint straight to R2 and register it. Returns the id.
|
|
179
|
+
|
|
180
|
+
A checkpoint is a SET of files — weights plus config.yaml plus
|
|
181
|
+
dataset_statistics.json — because the serving stack needs the small
|
|
182
|
+
ones and should not have to pull 19 GB of weights to read them.
|
|
183
|
+
|
|
184
|
+
`manifest` is the unibot-checkpoint/v0 document describing how to load
|
|
185
|
+
this checkpoint (regime, base_model, control_space, serving_env). It
|
|
186
|
+
is stored as its own queryable column, separate from `metrics`, which
|
|
187
|
+
is only for scalars.
|
|
188
|
+
"""
|
|
189
|
+
files = [Path(p) for p in ([paths] if isinstance(paths, (str, Path)) else paths)]
|
|
190
|
+
|
|
191
|
+
reg = self._call("POST", "/api/checkpoints", {
|
|
192
|
+
"wandb_run_id": wandb_run_id,
|
|
193
|
+
"wandb_url": wandb_url,
|
|
194
|
+
"step": step,
|
|
195
|
+
"files": [{"filename": f.name, "size": f.stat().st_size} for f in files],
|
|
196
|
+
"metrics": metrics or {},
|
|
197
|
+
"manifest": manifest or {},
|
|
198
|
+
"manifest_version": manifest_version,
|
|
199
|
+
})
|
|
200
|
+
|
|
201
|
+
by_name = {f.name: f for f in files}
|
|
202
|
+
completed = []
|
|
203
|
+
for entry in reg["files"]:
|
|
204
|
+
src = by_name[entry["filename"]]
|
|
205
|
+
if "part_urls" in entry:
|
|
206
|
+
# >5 GB: a single presigned PUT cannot take it.
|
|
207
|
+
etags = []
|
|
208
|
+
with src.open("rb") as fh:
|
|
209
|
+
for i, url in enumerate(entry["part_urls"], start=1):
|
|
210
|
+
etags.append({"PartNumber": i, "ETag": _put(url, fh.read(PART_SIZE))})
|
|
211
|
+
completed.append({
|
|
212
|
+
"filename": entry["filename"],
|
|
213
|
+
"upload_id": entry["upload_id"],
|
|
214
|
+
"parts": etags,
|
|
215
|
+
})
|
|
216
|
+
else:
|
|
217
|
+
_put(entry["upload_url"], src.read_bytes())
|
|
218
|
+
|
|
219
|
+
# Completing verifies every file actually landed, so a crashed upload
|
|
220
|
+
# never appears as a usable checkpoint.
|
|
221
|
+
self._call("POST", f"/api/checkpoints/{reg['id']}/complete",
|
|
222
|
+
{"files": completed})
|
|
223
|
+
return reg["id"]
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
# --------------------------------------------------------------- results
|
|
227
|
+
def create_result_set(self, wandb_run_id: str, name: str, kind: str = "eval") -> dict:
|
|
228
|
+
"""Open a result set. Returns {id, prefix}."""
|
|
229
|
+
return self._call("POST", "/api/results", {
|
|
230
|
+
"wandb_run_id": wandb_run_id, "name": name, "kind": kind,
|
|
231
|
+
})
|
|
232
|
+
|
|
233
|
+
def upload_result_file(self, result_id: str, rel_path: str, data: bytes,
|
|
234
|
+
content_type: str | None = None) -> str:
|
|
235
|
+
"""Upload one file into the set, keeping LeRobot's directory layout.
|
|
236
|
+
|
|
237
|
+
e.g. rel_path='meta/info.json', 'videos/observation.images.cam/…mp4',
|
|
238
|
+
'data/chunk-000/file-000.parquet' — the same shape as a training
|
|
239
|
+
dataset, which is what lets the viewer play these back.
|
|
240
|
+
"""
|
|
241
|
+
resp = self._call("POST", f"/api/results/{result_id}",
|
|
242
|
+
{"path": rel_path, "content_type": content_type})
|
|
243
|
+
_put(resp["upload_url"], data)
|
|
244
|
+
return resp["key"]
|
|
245
|
+
|
|
246
|
+
def upload_result_dir(self, result_id: str, root: str | Path) -> int:
|
|
247
|
+
root = Path(root)
|
|
248
|
+
n = 0
|
|
249
|
+
for f in sorted(root.rglob("*")):
|
|
250
|
+
if f.is_file():
|
|
251
|
+
self.upload_result_file(result_id, str(f.relative_to(root)), f.read_bytes())
|
|
252
|
+
n += 1
|
|
253
|
+
return n
|
|
254
|
+
|
|
255
|
+
def complete_result_set(self, result_id: str, episode_count: int | None = None) -> dict:
|
|
256
|
+
"""Close the set. Fails if nothing was uploaded."""
|
|
257
|
+
return self._call("POST", f"/api/results/{result_id}/complete",
|
|
258
|
+
{"episode_count": episode_count})
|
|
259
|
+
|
|
260
|
+
def result_files(self, result_id: str) -> dict:
|
|
261
|
+
"""Set metadata plus presigned URLs — playable directly, no proxy."""
|
|
262
|
+
return self._call("GET", f"/api/results/{result_id}")
|
|
263
|
+
|
|
264
|
+
def list_result_sets(self, wandb_run_id: str | None = None) -> list[dict]:
|
|
265
|
+
q = f"?run={wandb_run_id}" if wandb_run_id else ""
|
|
266
|
+
return self._call("GET", f"/api/results{q}")["results"]
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
def _put(url: str, data: bytes) -> str:
|
|
270
|
+
req = urllib.request.Request(url, method="PUT", data=data)
|
|
271
|
+
with urllib.request.urlopen(req, timeout=3600) as r:
|
|
272
|
+
return r.headers.get("ETag", "")
|
|
273
|
+
|
|
274
|
+
|
|
275
|
+
if __name__ == "__main__":
|
|
276
|
+
dv = DejaVu()
|
|
277
|
+
m = dv.manifest("G1_Dex1_PlacePhotoFrame")
|
|
278
|
+
print(f"manifest version={m['version']} drop={len(m['drop'])} flagged={len(m['flagged'])}")
|
|
279
|
+
print(f"checkpoints: {len(dv.list_checkpoints())}")
|