| """The whole Crello corpus, served from a Hugging Face bucket. |
| |
| FineEnvs/crello-bucket (cyberagent/crello 5.1.0, copied as is) |
| data/<split>-000NN-of-000MM.parquet 38 shards, 100-design row groups, 18 GB |
| resources/fonts.pickle |
| openenv/index/<snapshot>/catalog.sqlite built here, published next to the data |
| |
| **Index once.** `build_index` converts every design exactly as the pack builder |
| does, keeps the ones that make a playable task, and stores the finished task |
| record (ground truth, asset list, brief, difficulty features) plus where the |
| design lives (shard, row group, row). That is the whole catalog: listing tasks, |
| `get_task` and reset never touch the images. |
| |
| **Fetch on demand.** Asset pixels are only needed once an episode renders. The |
| first time a design is used its row group is read (from a mount at `/corpus` |
| on a Space, over HTTP locally), the design is converted again and checked |
| against the indexed record, and its assets land in a byte-bounded disk cache. |
| |
| **Pinned identity.** The snapshot id hashes the index version, the conversion |
| rules and each shard's content hash. A shard that changed since indexing is |
| refused rather than silently serving different ground truth. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import fcntl |
| import hashlib |
| import json |
| import os |
| import shutil |
| import sqlite3 |
| import threading |
| import time |
| import zlib |
| from collections import OrderedDict |
| from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor |
| from contextlib import closing |
| from pathlib import Path |
|
|
| from .difficulty import Scorer, features, shape |
| from .pack import SPLITS, TaskPack, accept, make_task, write_design |
|
|
| INDEX_VERSION = 1 |
| DEFAULT_BUCKET = "FineEnvs/crello-bucket" |
| DEFAULT_MAX_ELEMENTS = 30 |
| INDEX_PREFIX = "openenv/index" |
| FONTS_PATH = "resources/fonts.pickle" |
| SOURCE = {"repo": "cyberagent/crello", "revision": "5.1.0", |
| "license": "CDLA-Permissive-2.0; templates (c) VistaCreate, see dataset card"} |
|
|
| DDL = """ |
| CREATE TABLE tasks( |
| task_id TEXT PRIMARY KEY, split TEXT NOT NULL, position INTEGER NOT NULL, |
| crello_id TEXT NOT NULL UNIQUE, file TEXT NOT NULL, row_group INTEGER NOT NULL, |
| row INTEGER NOT NULL, public TEXT NOT NULL, payload BLOB NOT NULL, |
| features TEXT NOT NULL, score REAL NOT NULL, tier TEXT NOT NULL, shape TEXT NOT NULL, |
| doc_sha256 TEXT NOT NULL, UNIQUE(split, position)); |
| CREATE TABLE rejected(crello_id TEXT, split TEXT, reason TEXT); |
| CREATE TABLE meta(key TEXT PRIMARY KEY, value TEXT NOT NULL); |
| """ |
|
|
|
|
| def canonical(obj) -> str: |
| return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False) |
|
|
|
|
| def sha256_bytes(data: bytes) -> str: |
| return hashlib.sha256(data).hexdigest() |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| h = hashlib.sha256() |
| with open(path, "rb") as fh: |
| for chunk in iter(lambda: fh.read(1 << 22), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def doc_digest(doc: dict) -> str: |
| return sha256_bytes(canonical(doc).encode()) |
|
|
|
|
| |
|
|
|
|
| class BucketSource: |
| """Read-only access to the bucket, from a mounted/copied directory or over HTTP.""" |
|
|
| def __init__(self, bucket_id: str = DEFAULT_BUCKET, root: str | Path | None = None): |
| self.bucket_id = os.environ.get("DESIGN_CANVAS_BUCKET_ID") or bucket_id |
| self.root = Path(root) if root else None |
| self._fs = None |
|
|
| @property |
| def fs(self): |
| if self._fs is None: |
| from huggingface_hub import HfFileSystem |
|
|
| self._fs = HfFileSystem() |
| return self._fs |
|
|
| def url(self, path: str) -> str: |
| return f"hf://buckets/{self.bucket_id}/{path}" |
|
|
| def open(self, path: str): |
| if self.root: |
| return open(self.root / path, "rb") |
| return self.fs.open(self.url(path), "rb") |
|
|
| def inventory(self) -> list[dict]: |
| """Every parquet shard with its size and content hash (the bucket's xet hash).""" |
| if self.root: |
| return [{"path": str(p.relative_to(self.root)), "size": p.stat().st_size, |
| "sha256": sha256_file(p)} |
| for p in sorted((self.root / "data").glob("*.parquet"))] |
| from huggingface_hub import HfApi |
|
|
| return [{"path": f.path, "size": f.size, "xet_hash": f.xet_hash} |
| for f in sorted(HfApi().list_bucket_tree(self.bucket_id, prefix="data", recursive=True), |
| key=lambda f: f.path) |
| if f.path.endswith(".parquet")] |
|
|
| def verify(self, spec: dict) -> None: |
| """Refuse a shard whose content no longer matches the indexed one.""" |
| if self.root: |
| path = self.root / spec["path"] |
| if path.stat().st_size != spec["size"]: |
| raise ValueError(f"{spec['path']} changed since it was indexed (size)") |
| return |
| if not spec.get("xet_hash"): |
| raise ValueError(f"{spec['path']} was indexed from a local copy; serve it with a source root") |
| from huggingface_hub import HfApi |
|
|
| rows = list(HfApi().get_bucket_paths_info(self.bucket_id, [spec["path"]])) |
| if len(rows) != 1 or rows[0].xet_hash != spec["xet_hash"] or rows[0].size != spec["size"]: |
| raise ValueError(f"{spec['path']} changed in {self.bucket_id} since it was indexed; " |
| "rebuild the index rather than serve different ground truth") |
|
|
| def fetch(self, path: str, dest: Path) -> Path: |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| tmp = dest.with_suffix(dest.suffix + f".{os.getpid()}.part") |
| with self.open(path) as src, open(tmp, "wb") as out: |
| shutil.copyfileobj(src, out, 1 << 22) |
| tmp.replace(dest) |
| return dest |
|
|
|
|
| def split_of(path: str) -> str: |
| name = Path(path).name |
| for split in SPLITS: |
| if name.startswith(split + "-"): |
| return split |
| raise ValueError(f"cannot tell the split of {path}") |
|
|
|
|
| def crello_features(parquet_file): |
| import datasets |
|
|
| return datasets.Features.from_arrow_schema(parquet_file.schema_arrow) |
|
|
|
|
| |
|
|
|
|
| def _index_shard(args) -> list[dict]: |
| """Worker: convert every design in one shard; return one record per design.""" |
| import pyarrow.parquet as pq |
|
|
| from .crello import convert_row |
| from .fonts import FontRegistry |
|
|
| bucket_id, root, path, fonts_dir, max_elements = args |
| source = BucketSource(bucket_id, root) |
| fonts = FontRegistry(fonts_dir) |
| split = split_of(path) |
| out = [] |
| with source.open(path) as fh: |
| pf = pq.ParquetFile(fh, pre_buffer=True) |
| feats = crello_features(pf) |
| cols = [n for n in pf.schema_arrow.names if n != "preview"] |
| for rg in range(pf.num_row_groups): |
| for row_i, row in enumerate(pf.read_row_group(rg, columns=cols).to_pylist()): |
| design = convert_row(row, feats, fonts) |
| rec = {"crello_id": design.crello_id, "split": split, "file": path, |
| "row_group": rg, "row": row_i} |
| reason = accept(design, max_elements) |
| if reason: |
| out.append({**rec, "rejected": reason}) |
| continue |
| task = make_task(design, split, 0) |
| for a in task["assets"]: |
| img = design.assets[a["id"]] |
| a["sha256"] = sha256_bytes(img.tobytes() + repr(img.size).encode()) |
| f = features(design.doc) |
| out.append({**rec, "task": task, "features": f, |
| "cluster": row.get("cluster_index"), |
| "group": feats["group"].names[row["group"]] if row.get("group") is not None else None}) |
| return out |
|
|
|
|
| def build_index(out: Path | str, source: BucketSource, *, max_elements: int = DEFAULT_MAX_ELEMENTS, |
| workers: int = 8, shards: list[str] | None = None, log=print) -> dict: |
| """Convert the whole corpus into `out/catalog.sqlite` + `out/manifest.json`.""" |
| from .fonts import extract_crello_fonts |
|
|
| out = Path(out) |
| out.mkdir(parents=True, exist_ok=True) |
| t0 = time.time() |
| inventory = [s for s in source.inventory() if shards is None or s["path"] in shards] |
| if not inventory: |
| raise ValueError("no parquet shards found in the source") |
| fonts_dir = out / "fonts" |
| if not (fonts_dir / "manifest.json").exists(): |
| pickle_path = source.fetch(FONTS_PATH, out / "fonts.pickle") if not source.root \ |
| else source.root / FONTS_PATH |
| extract_crello_fonts(pickle_path, fonts_dir) |
| log(f"indexing {len(inventory)} shards with {workers} workers") |
| jobs = [(source.bucket_id, str(source.root) if source.root else None, s["path"], str(fonts_dir), |
| max_elements) for s in inventory] |
| records = [] |
| with ProcessPoolExecutor(workers) as ex: |
| for i, recs in enumerate(ex.map(_index_shard, jobs), 1): |
| records.extend(recs) |
| log(f" {i}/{len(jobs)} shards, {len(records)} designs, {time.time() - t0:.0f}s") |
|
|
| kept = [r for r in records if "task" in r] |
| scorer = Scorer([r["features"] for r in kept]) |
| counts, rejected = {}, {} |
| db_path = out / "catalog.sqlite" |
| db_path.unlink(missing_ok=True) |
| with closing(sqlite3.connect(db_path)) as db: |
| db.executescript(DDL) |
| for split in SPLITS: |
| mine = sorted((r for r in kept if r["split"] == split), |
| key=lambda r: (r["file"], r["row_group"], r["row"])) |
| counts[split] = len(mine) |
| for pos, r in enumerate(mine): |
| task = r["task"] |
| task.update(task_id=f"{split}-{pos:05d}", index=pos) |
| score = scorer.score(r["features"]) |
| tier = scorer.tier(score) |
| task["difficulty"] = tier |
| task["difficulty_score"] = score |
| task["meta"] = {"cluster": r["cluster"], "group": r["group"]} |
| pub = {**TaskPack.public(task), "difficulty_score": score, |
| "shape": shape(r["features"]["aspect"])} |
| db.execute("INSERT INTO tasks VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)", ( |
| task["task_id"], split, pos, r["crello_id"], r["file"], r["row_group"], r["row"], |
| canonical(pub), zlib.compress(canonical(task).encode(), 6), |
| canonical({**r["features"], "cluster": r["cluster"], "group": r["group"], |
| "category": task["category"], "format": task["format"]}), |
| score, tier, pub["shape"], doc_digest(task["doc"]))) |
| for r in records: |
| if "rejected" in r: |
| db.execute("INSERT INTO rejected VALUES(?,?,?)", (r["crello_id"], r["split"], r["rejected"])) |
| rejected[r["rejected"]] = rejected.get(r["rejected"], 0) + 1 |
| db.commit() |
| db.execute("VACUUM") |
| config = {"max_elements": max_elements, "index_version": INDEX_VERSION} |
| snapshot = sha256_bytes(canonical([INDEX_VERSION, config, scorer.describe(), |
| [[s["path"], s["size"], s.get("xet_hash") or s.get("sha256")] |
| for s in inventory]]).encode()) |
| fonts_sha = sha256_file(out / "fonts.pickle") if (out / "fonts.pickle").exists() else \ |
| sha256_file(source.root / FONTS_PATH) |
| manifest = { |
| "status": "ready", |
| "storage": "bucket-parquet", |
| "index_version": INDEX_VERSION, |
| "snapshot_id": snapshot, |
| "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "bucket_id": source.bucket_id, |
| "source": SOURCE, |
| "config": config, |
| "difficulty": scorer.describe(), |
| "splits": counts, |
| "rejected": dict(sorted(rejected.items(), key=lambda kv: -kv[1])), |
| "fonts": {"path": FONTS_PATH, "sha256": fonts_sha}, |
| "catalog": {"path": f"{INDEX_PREFIX}/{snapshot[:16]}/catalog.sqlite", |
| "size": db_path.stat().st_size, "sha256": sha256_file(db_path)}, |
| "source_files": inventory, |
| } |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| log(f"index ready: {counts}, rejected {sum(rejected.values())}, {time.time() - t0:.0f}s") |
| return manifest |
|
|
|
|
| def publish_index(out: Path | str, bucket_id: str | None = None) -> str: |
| """Upload catalog + manifest next to the data. Returns the manifest's hf:// URL.""" |
| from huggingface_hub import HfApi |
|
|
| out = Path(out) |
| manifest = json.loads((out / "manifest.json").read_text()) |
| bucket = bucket_id or manifest["bucket_id"] |
| prefix = manifest["catalog"]["path"].rsplit("/", 1)[0] |
| HfApi().batch_bucket_files(bucket, add=[ |
| (str(out / "catalog.sqlite"), manifest["catalog"]["path"]), |
| (str(out / "manifest.json"), f"{prefix}/manifest.json"), |
| ]) |
| return f"hf://buckets/{bucket}/{prefix}/manifest.json" |
|
|
|
|
| |
|
|
|
|
| class DiskLRU: |
| """Byte-bounded directory cache. Entries used in the last `min_age` seconds are never evicted. |
| |
| Safe to share between server replicas on one machine: an entry is filled in a private |
| temporary directory and renamed into place (the first finished copy wins), and "last |
| used" is the mtime of the entry's `.ok` marker, so a replica never evicts an entry that |
| another one touched recently. |
| """ |
|
|
| _TMP = ".tmp-" |
|
|
| def __init__(self, root: Path, budget: int, min_age: float = 900.0): |
| self.root = Path(root) |
| self.root.mkdir(parents=True, exist_ok=True) |
| self.budget = budget |
| self.min_age = min_age |
| self._lock = threading.Lock() |
| self._keys: dict[str, threading.Lock] = {} |
| self._sizes: OrderedDict[str, int] = OrderedDict() |
| for p in sorted(self.root.iterdir(), key=lambda p: p.stat().st_mtime): |
| if p.name.startswith(self._TMP): |
| if time.time() - p.stat().st_mtime > 3600: |
| shutil.rmtree(p, ignore_errors=True) |
| elif (p / ".ok").exists(): |
| self._sizes[p.name] = int((p / ".ok").read_text() or 0) |
| elif p.is_dir(): |
| shutil.rmtree(p, ignore_errors=True) |
|
|
| def _hit(self, key: str, path: Path) -> bool: |
| ok = path / ".ok" |
| try: |
| size = int(ok.read_text() or 0) |
| os.utime(ok) |
| except (OSError, ValueError): |
| return False |
| with self._lock: |
| self._sizes[key] = size |
| self._sizes.move_to_end(key) |
| return True |
|
|
| def get(self, key: str, fill) -> Path: |
| path = self.root / key |
| with self._lock: |
| lock = self._keys.setdefault(key, threading.Lock()) |
| with lock: |
| if self._hit(key, path): |
| return path |
| tmp = self.root / f"{self._TMP}{key}-{os.getpid()}-{threading.get_ident()}" |
| shutil.rmtree(tmp, ignore_errors=True) |
| tmp.mkdir(parents=True) |
| try: |
| fill(tmp) |
| size = sum(f.stat().st_size for f in tmp.rglob("*") if f.is_file()) |
| (tmp / ".ok").write_text(str(size)) |
| if path.exists() and not (path / ".ok").exists(): |
| shutil.rmtree(path, ignore_errors=True) |
| try: |
| tmp.rename(path) |
| except OSError: |
| if not self._hit(key, path): |
| raise |
| return path |
| finally: |
| shutil.rmtree(tmp, ignore_errors=True) |
| with self._lock: |
| self._sizes[key] = size |
| self._evict(keep=key) |
| return path |
|
|
| def _evict(self, keep: str) -> None: |
| now = time.time() |
| total = sum(self._sizes.values()) |
| if total <= self.budget: |
| return |
| used = {} |
| for key in list(self._sizes): |
| try: |
| used[key] = (self.root / key / ".ok").stat().st_mtime |
| except OSError: |
| total -= self._sizes.pop(key) |
| for key in sorted(used, key=used.get): |
| if total <= self.budget: |
| break |
| if key == keep or now - used[key] < self.min_age: |
| continue |
| total -= self._sizes.pop(key) |
| trash = self.root / f"{self._TMP}evict-{key}-{os.getpid()}" |
| try: |
| (self.root / key).rename(trash) |
| except OSError: |
| continue |
| shutil.rmtree(trash, ignore_errors=True) |
|
|
| def stats(self) -> dict: |
| with self._lock: |
| return {"entries": len(self._sizes), "bytes": sum(self._sizes.values()), "budget": self.budget} |
|
|
|
|
| |
|
|
|
|
| def _read_manifest(value: str | Path) -> tuple[dict, str]: |
| """Manifest + the directory (local or hf://) it came from.""" |
| value = str(value) |
| if value.startswith("hf://"): |
| from huggingface_hub import HfFileSystem |
|
|
| with HfFileSystem().open(value, "rb") as fh: |
| return json.loads(fh.read()), value.rsplit("/", 1)[0] |
| path = Path(value).expanduser() |
| if path.is_dir(): |
| path = path / "manifest.json" |
| return json.loads(path.read_text()), str(path.parent) |
|
|
|
|
| class CorpusPack(TaskPack): |
| """The TaskPack interface over the full bucket-backed corpus. |
| |
| Everything the server calls works unchanged: task records come from the |
| index, and any call that needs pixels (`asset_paths`, `preview_path`, |
| renders, grading) materialises the design first. |
| """ |
|
|
| def __init__(self, manifest: str | Path, cache_dir: str | Path | None = None, *, |
| source_root: str | Path | None = None, design_cache_bytes: int = 4 << 30, |
| group_cache_bytes: int = 2 << 30): |
| from .fonts import FontRegistry, extract_crello_fonts |
| from .render import Renderer |
|
|
| self.manifest, where = _read_manifest(manifest) |
| if self.manifest.get("storage") != "bucket-parquet" or \ |
| self.manifest.get("index_version") != INDEX_VERSION: |
| raise ValueError("not a supported corpus manifest; rebuild the index") |
| self.snapshot_id = self.manifest["snapshot_id"] |
| cache = Path(cache_dir or os.environ.get("DESIGN_CANVAS_CACHE") |
| or Path.home() / ".cache/design_canvas/corpus").expanduser() / self.snapshot_id[:16] |
| cache.mkdir(parents=True, exist_ok=True) |
| self.cache_dir = cache |
| root = source_root or os.environ.get("DESIGN_CANVAS_SOURCE_ROOT") |
| self.source = BucketSource(self.manifest["bucket_id"], root) |
| self._specs = {s["path"]: s for s in self.manifest["source_files"]} |
|
|
| fonts_dir = cache / "fonts" |
| |
| |
| with open(cache / ".setup.lock", "w") as lock: |
| fcntl.flock(lock, fcntl.LOCK_EX) |
| self._catalog = self._local_catalog(where) |
| if not (fonts_dir / "manifest.json").exists(): |
| pickle_path = self.source.root / FONTS_PATH if self.source.root else \ |
| self.source.fetch(FONTS_PATH, cache / "fonts.pickle") |
| if sha256_file(pickle_path) != self.manifest["fonts"]["sha256"]: |
| raise ValueError("fonts.pickle does not match the indexed corpus") |
| extract_crello_fonts(pickle_path, fonts_dir) |
| self.root = cache |
| self.fonts = FontRegistry(fonts_dir) |
| self.renderer = Renderer(self.fonts) |
| self._docs = OrderedDict() |
| self._renders = OrderedDict() |
| self._tasks_cache: OrderedDict[str, dict] = OrderedDict() |
| self._lock = threading.Lock() |
| self._db_local = threading.local() |
| self.designs = DiskLRU(cache / "designs", design_cache_bytes) |
| self.groups = DiskLRU(cache / "groups", group_cache_bytes, min_age=60.0) |
| self._features = None |
| self._prefetch = ThreadPoolExecutor(2, thread_name_prefix="design-canvas-prefetch") |
|
|
| def _local_catalog(self, where: str) -> Path: |
| """SQLite is copied to local disk and checksummed; it is never queried over a mount.""" |
| spec = self.manifest["catalog"] |
| local = self.cache_dir / "catalog.sqlite" |
| if local.exists() and local.stat().st_size == spec["size"]: |
| return local |
| name = spec["path"].rsplit("/", 1)[1] |
| if where.startswith("hf://"): |
| from huggingface_hub import HfFileSystem |
|
|
| tmp = local.with_suffix(".part") |
| HfFileSystem().get_file(f"{where}/{name}", str(tmp)) |
| else: |
| candidates = [Path(where) / name] |
| if self.source.root: |
| candidates.append(self.source.root / spec["path"]) |
| src = next((p for p in candidates if p.exists()), None) |
| if src is None: |
| return self._local_catalog(f"hf://buckets/{self.source.bucket_id}/{spec['path'].rsplit('/', 1)[0]}") |
| tmp = local.with_suffix(".part") |
| shutil.copyfile(src, tmp) |
| if sha256_file(tmp) != spec["sha256"]: |
| tmp.unlink() |
| raise ValueError("catalog.sqlite does not match its manifest") |
| tmp.replace(local) |
| return local |
|
|
| @property |
| def db(self) -> sqlite3.Connection: |
| conn = getattr(self._db_local, "conn", None) |
| if conn is None: |
| conn = sqlite3.connect(f"file:{self._catalog}?mode=ro", uri=True) |
| self._db_local.conn = conn |
| return conn |
|
|
| |
|
|
| def splits(self) -> list[str]: |
| return [s for s in SPLITS if self.manifest["splits"].get(s)] |
|
|
| def count(self, split: str) -> int: |
| if split not in self.manifest["splits"]: |
| raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}") |
| return int(self.manifest["splits"][split]) |
|
|
| def _task(self, row) -> dict: |
| task_id, payload = row |
| with self._lock: |
| task = self._tasks_cache.get(task_id) |
| if task is not None: |
| self._tasks_cache.move_to_end(task_id) |
| return task |
| task = json.loads(zlib.decompress(payload)) |
| with self._lock: |
| self._tasks_cache[task_id] = task |
| while len(self._tasks_cache) > 4096: |
| self._tasks_cache.popitem(last=False) |
| while len(self._docs) > 1024: |
| self._docs.popitem(last=False) |
| return task |
|
|
| def at(self, split: str, index: int) -> dict: |
| n = self.count(split) |
| if not 0 <= index < n: |
| raise IndexError(f"{split} has {n} tasks; index {index} is out of range") |
| row = self.db.execute("SELECT task_id, payload FROM tasks WHERE split=? AND position=?", |
| (split, index)).fetchone() |
| return self._task(row) |
|
|
| def get(self, task_id: str) -> dict: |
| row = self.db.execute("SELECT task_id, payload FROM tasks WHERE task_id=?", (task_id,)).fetchone() |
| if row is None: |
| raise KeyError(f"unknown task_id {task_id!r}") |
| return self._task(row) |
|
|
| def public(self, task: dict) -> dict: |
| row = self.db.execute("SELECT public FROM tasks WHERE task_id=?", (task["task_id"],)).fetchone() |
| return json.loads(row[0]) |
|
|
| def task_range(self, split: str, start=None, stop=None) -> list[dict]: |
| n = self.count(split) |
| lo, hi, _ = slice(start, stop).indices(n) |
| rows = self.db.execute("SELECT public FROM tasks WHERE split=? AND position>=? AND position<? " |
| "ORDER BY position", (split, lo, hi)).fetchall() |
| return [json.loads(r[0]) for r in rows] |
|
|
| def rows(self, split: str | None = None) -> list[dict]: |
| """Every indexed design's selection features (for building evaluation sets).""" |
| q = "SELECT task_id, split, crello_id, features, score, tier, shape, doc_sha256 FROM tasks" |
| args = () |
| if split: |
| q += " WHERE split=?" |
| args = (split,) |
| return [{"task_id": t, "split": s, "crello_id": c, **json.loads(f), "score": sc, |
| "tier": tier, "shape": sh, "doc_sha256": d} |
| for t, s, c, f, sc, tier, sh, d in self.db.execute(q + " ORDER BY split, position", args)] |
|
|
| |
|
|
| def _location(self, task_id: str): |
| return self.db.execute("SELECT file, row_group, row, crello_id, doc_sha256 FROM tasks " |
| "WHERE task_id=?", (task_id,)).fetchone() |
|
|
| def _group(self, file: str, rg: int) -> Path: |
| spec = self._specs[file] |
| key = f"{(spec.get('xet_hash') or spec.get('sha256'))[:16]}-{rg}" |
|
|
| def fill(dest: Path): |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| self.source.verify(spec) |
| with self.source.open(file) as fh: |
| pf = pq.ParquetFile(fh, pre_buffer=True) |
| |
| |
| cols = [n for n in pf.schema_arrow.names if n != "preview"] |
| table = pf.read_row_group(rg, columns=cols) |
| with pa.OSFile(str(dest / "group.arrow"), "wb") as sink, \ |
| pa.ipc.new_file(sink, table.schema) as writer: |
| writer.write_table(table) |
|
|
| return self.groups.get(key, fill) / "group.arrow" |
|
|
| def materialize(self, task: dict) -> Path: |
| """Make sure this design's assets and preview are on disk; return its directory.""" |
| file, rg, row_i, crello_id, digest = self._location(task["task_id"]) |
|
|
| def fill(dest: Path): |
| import pyarrow as pa |
|
|
| from .crello import convert_row |
|
|
| with pa.memory_map(str(self._group(file, rg))) as src: |
| table = pa.ipc.open_file(src).read_all() |
| row = table.slice(row_i, 1).to_pylist()[0] |
| if str(row["id"]) != crello_id: |
| raise ValueError(f"{task['task_id']}: row {file}:{rg}:{row_i} is not design {crello_id}") |
| if self._features is None: |
| import datasets |
|
|
| self._features = datasets.Features.from_arrow_schema(table.schema) |
| design = convert_row(row, self._features, self.fonts) |
| if doc_digest(design.doc.model_dump(mode="json", exclude_defaults=False)) != digest: |
| raise ValueError(f"{task['task_id']}: converted design differs from the index; " |
| "the converter changed since indexing, rebuild the index") |
| for a in task["assets"]: |
| img = design.assets[a["id"]] |
| if sha256_bytes(img.tobytes() + repr(img.size).encode()) != a["sha256"]: |
| raise ValueError(f"{task['task_id']}: asset {a['id']} differs from the index") |
| staging = dest / "_stage" |
| write_design(staging, design) |
| for p in (staging / "designs" / crello_id).iterdir(): |
| p.replace(dest / p.name) |
| shutil.rmtree(staging) |
|
|
| |
| |
| return self.designs.get(crello_id, fill) |
|
|
| def prefetch(self, tasks: list[dict]) -> None: |
| for task in tasks: |
| self._prefetch.submit(self.materialize, task) |
|
|
| def warm(self, splits=("train", "validation", "test"), workers: int = 4, progress=None) -> dict: |
| """Materialise every design in `splits`, one row group at a time, so each ~50 MB group is |
| downloaded once. Afterwards resets never touch the network: run it before training or a |
| large eval so a burst of cold resets does not queue behind bucket downloads.""" |
| by_group: dict[tuple, list[str]] = {} |
| for split in splits: |
| for task_id, file, rg in self.db.execute( |
| "SELECT task_id, file, row_group FROM tasks WHERE split=? ORDER BY file, row_group, row", (split,)): |
| by_group.setdefault((file, rg), []).append(task_id) |
| done = {"groups": 0, "designs": 0, "failed": []} |
| lock = threading.Lock() |
|
|
| def one(ids): |
| for task_id in ids: |
| try: |
| self.materialize(self.get(task_id)) |
| except Exception as exc: |
| with lock: |
| done["failed"].append(f"{task_id}: {type(exc).__name__}: {exc}"[:300]) |
| with lock: |
| done["groups"] += 1 |
| done["designs"] += len(ids) |
| if progress: |
| progress(done["groups"], len(by_group), done["designs"]) |
|
|
| with ThreadPoolExecutor(workers, thread_name_prefix="design-canvas-warm") as pool: |
| list(pool.map(one, by_group.values())) |
| return {"warmed_row_groups": done["groups"], "warmed_designs": done["designs"], "failed": done["failed"], |
| "cache": self.stats()} |
|
|
| def asset_paths(self, task: dict) -> dict[str, Path]: |
| self.materialize(task) |
| return super().asset_paths(task) |
|
|
| def preview_path(self, task: dict) -> Path: |
| path = super().preview_path(task) |
| if not path.exists(): |
| import pyarrow.parquet as pq |
|
|
| self.materialize(task) |
| file, rg, row_i, crello_id, _ = self._location(task["task_id"]) |
| with self.source.open(file) as fh: |
| col = pq.ParquetFile(fh, pre_buffer=True).read_row_group(rg, columns=["id", "preview"]) |
| row = col.slice(row_i, 1).to_pylist()[0] |
| if str(row["id"]) != crello_id or not row["preview"]: |
| raise FileNotFoundError(f"{task['task_id']} has no Crello preview") |
| from .crello import _to_pil |
|
|
| tmp = path.with_suffix(".part") |
| _to_pil(row["preview"]).convert("RGB").save(tmp, "JPEG", quality=88) |
| tmp.replace(path) |
| return path |
|
|
| def asset_similarity(self, task: dict): |
| self.materialize(task) |
| return super().asset_similarity(task) |
|
|
| def stats(self) -> dict: |
| return {"snapshot_id": self.snapshot_id, "designs": self.designs.stats(), |
| "row_groups": self.groups.stats(), |
| "transport": "mounted-filesystem" if self.source.root else "http"} |
|
|
