diff --git a/README.md b/README.md index 0ddb5d5..7855299 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,10 @@ HF_TOKEN=... uv run --project scripts scripts/hf-sync.py plan \ Review or edit `plan.json` before applying it. +See [Planning and applying a dataset conversion](docs/plan-and-apply.md) for a +complete walkthrough, plan review checklist, resume behavior, and operational +tuning guidance. + ## Apply a plan ```sh diff --git a/docs/plan-and-apply.md b/docs/plan-and-apply.md new file mode 100644 index 0000000..0670fc6 --- /dev/null +++ b/docs/plan-and-apply.md @@ -0,0 +1,206 @@ +# Planning and applying a dataset conversion + +`hf-sync.py` deliberately separates discovery from execution: + +1. `plan` inspects the source and destination repositories and writes a JSON + action plan. It does not download, convert, or upload dataset files. +2. You review the plan to confirm the immutable revisions, selected files, + output paths, and commit batches. +3. `apply` executes exactly that plan and records resumable progress in a + checkpoint. + +This split makes a large conversion inspectable before it writes anything to +the destination repository. + +## Prerequisites + +Install the Python environment from the repository root: + +```sh +uv sync --project scripts +``` + +Set `HF_TOKEN` to a Hugging Face token that can read the source dataset and +write to the destination dataset: + +```sh +export HF_TOKEN=hf_... +``` + +Both repositories must be Xet-enabled. Vortex output also requires a `vx` +binary built from the Vortex repository with the `unstable_encodings` feature. +The `parquet-zstd6` format only requires the Python environment. + +## Step 1: create a plan + +The following example plans a small FineWeb conversion: + +```sh +uv run --project scripts scripts/hf-sync.py plan \ + --repo HuggingFaceFW/fineweb \ + --revision main \ + --prefix sample/10BT \ + --mode first \ + --limit 10 \ + --formats vortex,vortex-compact \ + --upload-repo vortex-data/fineweb \ + --upload-revision main \ + --plan-file plan.json +``` + +Planning resolves both requested revisions to immutable commit IDs. Later +changes to `main` therefore cannot silently change the source files represented +by the plan. + +### Selecting source files + +Use one of these selection modes: + +- `--mode first --limit N` selects the first `N` matching shards. +- `--mode sample --target-size 10GB --seed 0` selects shards spread across the + ordered dataset until their approximate total size reaches the target. +- `--mode all` selects every matching shard. + +`--prefix` selects the folder to scan. Use `/` for the repository root. +`--include` is matched against paths below that prefix and defaults to +`*.parquet`. Repeat `--filter` to restrict the selection with full repository +path globs: + +```sh +--prefix sample --filter 'sample/10BT/*' --filter 'sample/100BT/*' +``` + +### Selecting output formats + +`--formats` accepts a comma-separated subset of: + +- `parquet-zstd6`: Parquet rewritten with Zstandard level 6. +- `vortex`: Vortex using the `btrblocks` strategy. +- `vortex-compact`: Vortex using the `compact` strategy. + +Destination paths are deterministic: + +```text +// +``` + +Use `--upload-prefix` to place all generated formats below an additional +destination folder. + +## Step 2: review the plan + +Open `plan.json` and verify: + +- `source.revision` and `destination.revision` are the expected immutable + commits. +- `work_chunks` contains the intended source shards. +- Every action has the expected `create` or `skip` decision and destination + path. +- `upload_batches` has acceptable commit boundaries and messages. +- `summary` has plausible file, byte, create, and skip counts. + +The apply command validates that every create action belongs to exactly one +upload batch. If you edit the plan, keep `work_chunks` and `upload_batches` +consistent. For routine selection changes, generating a new plan is safer than +editing the JSON manually. + +## Step 3: apply the plan + +Run the reviewed plan: + +```sh +uv run --project scripts scripts/hf-sync.py apply plan.json \ + --vx /path/to/vx \ + --output-dir data/fineweb-run \ + --delete-after-upload +``` + +During execution, the script: + +1. validates the plan and destination revision; +2. downloads selected Parquet shards with bounded concurrency; +3. converts the requested formats in parallel; +4. preuploads and commits outputs in the planned batches; +5. writes checkpoint, status, and metrics files under `--output-dir`. + +All three stages share one `--workers` pool. A worker can claim any kind of +work. The scheduler drains an upload queue at its high-water mark, otherwise +refills the download buffer, then handles ready uploads, then converts a ready +source shard. Queue claims are atomic, so two workers cannot process the same +file. + +Downloads and uploads are real multi-file operations. Their batch sizes start +at the configured initial value, grow while aggregate throughput improves, and +back off after a significant regression. File and byte limits cap every claim. +Upload preupload batches remain separate from planned commit batches: files may +transfer together, but a commit is created only when all members of its planned +batch are ready. + +`--delete-after-upload` removes a local output only after its upload is +committed successfully. Source downloads are removed after conversion unless +`--keep-downloads` is set. + +### Test without writing to Hugging Face + +Use a local destination to exercise downloading and conversion without remote +uploads: + +```sh +uv run --project scripts scripts/hf-sync.py apply plan.json \ + --vx /path/to/vx \ + --output-dir data/test-run \ + --upload-local-dir data/test-sink +``` + +The local sink preserves the same destination path layout as a Hugging Face +dataset repository. + +## Resume an interrupted run + +Re-run the same apply command with the same `--output-dir`: + +```sh +uv run --project scripts scripts/hf-sync.py apply plan.json \ + --vx /path/to/vx \ + --output-dir data/fineweb-run \ + --delete-after-upload +``` + +The checkpoint reuses completed uploads and valid local encodings. Transfer +tuning such as Xet cache location or concurrency may change between attempts, +but changing the selected files, formats, destination, or encoding-affecting +arguments requires a new output directory or a new plan. + +Before resuming, do not remove local outputs that have finished conversion but +have not yet been committed. They are reusable and avoid repeated conversion. + +## Operational controls + +The most useful apply controls are: + +- `--workers` sets the size of the unified worker pool. +- `--download-initial-concurrency` and `--download-max-concurrency` control the + adaptive number of files in each download batch. +- `--download-buffer-files` and `--download-buffer-size` bound admitted source + downloads. +- `--upload-workers` and `--upload-max-concurrency` control the adaptive number + of files in each upload batch. +- `--upload-buffer-files` and `--upload-buffer-size` bound pending output work. +- `--xet-cache`, `--xet-range-gets`, and `--xet-high-performance` tune Xet. +- `--status-interval` controls how frequently live status is written. + +Start conservatively. Increase concurrency only when CPU, disk, memory, and +network headroom are visible in the live status and host metrics. + +## Generated run files + +The output directory contains the durable state needed to inspect or resume a +run: + +- `checkpoint.json` records per-file conversion and upload state. +- `status.json` is an atomically replaced live pipeline snapshot. +- `metrics/selection.json` records the effective run configuration. +- final metric reports summarize selected files and generated formats. + +Keep the plan and output directory together until the destination has been +verified and no resume is required. diff --git a/pyproject.toml b/pyproject.toml index c73d521..db15947 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,3 +3,6 @@ required-version = ">=0.8.0" [tool.uv.workspace] members = ["scripts"] + +[tool.ruff] +line-length = 120 diff --git a/scripts/hf-sync.py b/scripts/hf-sync.py index 547f0fc..0babb55 100755 --- a/scripts/hf-sync.py +++ b/scripts/hf-sync.py @@ -16,8 +16,8 @@ import hashlib import json import os -import random import queue +import random import subprocess import sys import threading @@ -25,7 +25,6 @@ from collections import deque from pathlib import Path - DEFAULT_TARGET_BYTES = 10_000_000_000 DEFAULT_DATA_DIR = Path(__file__).resolve().parents[1] / "data" / "hf-sync" DEFAULT_XET_CACHE = DEFAULT_DATA_DIR / "xet-cache" @@ -46,7 +45,8 @@ def safe_print(*values, file=None, **kwargs): with LOG_LOCK: print(*values, file=target, **kwargs) except BrokenPipeError: - replacement = open(os.devnull, "w") + # The replacement must stay open because it becomes a process-global stream. + replacement = open(os.devnull, "w") # noqa: SIM115 if target is sys.stdout: sys.stdout = replacement elif target is sys.stderr: @@ -65,6 +65,7 @@ def retryable_hub_error(error): messages.append(str(current).lower()) try: import httpx + if isinstance(current, httpx.TransportError): return True except ImportError: @@ -77,11 +78,20 @@ def retryable_hub_error(error): return True current = current.__cause__ or current.__context__ message = " ".join(messages) - return any(fragment in message for fragment in ( - "timed out", "timeout", "request body", "connection reset", - "connection aborted", "connection closed", "broken pipe", - "temporarily unavailable", "internal error", - )) + return any( + fragment in message + for fragment in ( + "timed out", + "timeout", + "request body", + "connection reset", + "connection aborted", + "connection closed", + "broken pipe", + "temporarily unavailable", + "internal error", + ) + ) def retry_call(operation, attempts, base_delay=1.0, max_delay=60.0): @@ -95,8 +105,7 @@ def retry_call(operation, attempts, base_delay=1.0, max_delay=60.0): delay = min(max_delay, base_delay * 2 ** (attempt - 1)) delay += random.uniform(0, delay * 0.25) safe_print( - f"transient Hub error (attempt {attempt}/{attempts}): {error}; " - f"retrying in {delay:.1f}s", + f"transient Hub error (attempt {attempt}/{attempts}): {error}; retrying in {delay:.1f}s", file=sys.stderr, flush=True, ) @@ -104,14 +113,23 @@ def retry_call(operation, attempts, base_delay=1.0, max_delay=60.0): def parse_size(value): - units = {"b": 1, "kb": 1000, "mb": 1000**2, "gb": 1000**3, "tb": 1000**4, - "kib": 1024, "mib": 1024**2, "gib": 1024**3, "tib": 1024**4} + units = { + "b": 1, + "kb": 1000, + "mb": 1000**2, + "gb": 1000**3, + "tb": 1000**4, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, + "tib": 1024**4, + } value = value.strip().lower() number = value unit = "b" for suffix in sorted(units, key=len, reverse=True): if value.endswith(suffix): - number = value[:-len(suffix)] + number = value[: -len(suffix)] unit = suffix break return int(float(number) * units[unit]) @@ -133,8 +151,8 @@ def resolve_dataset_revision(api, repo, revision, attempts, timeout=None): def require_xet_repository(api, repo, revision, attempts): """Fail before transfer if a source or destination repository is not Xet-backed.""" info = retry_call( - lambda: api.repo_info(repo, repo_type="dataset", revision=revision, - expand=["xetEnabled"]), attempts) + lambda: api.repo_info(repo, repo_type="dataset", revision=revision, expand=["xetEnabled"]), attempts + ) enabled = getattr(info, "xet_enabled", None) if enabled is None: enabled = info.__dict__.get("xetEnabled") @@ -145,8 +163,13 @@ def require_xet_repository(api, repo, revision, attempts): def list_shards(api, repo, revision, prefix, include, filters=(), attempts=3): shards = [] tree = retry_call( - lambda: list(api.list_repo_tree(repo, path_in_repo=prefix or None, recursive=True, - revision=revision, repo_type="dataset")), attempts) + lambda: list( + api.list_repo_tree( + repo, path_in_repo=prefix or None, recursive=True, revision=revision, repo_type="dataset" + ) + ), + attempts, + ) for item in tree: path = getattr(item, "path", "") size = getattr(item, "size", None) @@ -193,33 +216,54 @@ def list_repository_files(api, repo, revision, prefix="", attempts=3): files = {} try: tree = retry_call( - lambda: list(api.list_repo_tree(repo, path_in_repo=prefix or None, recursive=True, - revision=revision, repo_type="dataset")), attempts) + lambda: list( + api.list_repo_tree( + repo, path_in_repo=prefix or None, recursive=True, revision=revision, repo_type="dataset" + ) + ), + attempts, + ) except RemoteEntryNotFoundError: return {} for item in tree: path = getattr(item, "path", None) if path is not None and getattr(item, "size", None) is not None: - files[path] = {"size": item.size, "oid": getattr(item, "blob_id", None), - "lfs": getattr(item, "lfs", None)} + files[path] = {"size": item.size, "oid": getattr(item, "blob_id", None), "lfs": getattr(item, "lfs", None)} return files -def download_shard(repo, revision, shard, download_root, attempts, etag_timeout): - from huggingface_hub import hf_hub_download +def download_shard_batch(repo, revision, shards, download_root, attempts, etag_timeout): + """Download and verify several shards through one concurrent Hub request.""" + from huggingface_hub import snapshot_download started = time.monotonic() - path = Path(retry_call( - lambda: hf_hub_download(repo, shard["path"], repo_type="dataset", revision=revision, - local_dir=download_root, etag_timeout=etag_timeout), attempts)) - elapsed = time.monotonic() - started - actual_size = path.stat().st_size - if actual_size != shard["size"]: - raise RuntimeError(f"downloaded {actual_size} bytes for {path}, expected {shard['size']}") - if shard.get("sha256") and file_sha256(path) != shard["sha256"]: - path.unlink(missing_ok=True) - raise RuntimeError(f"SHA-256 mismatch for {shard['path']}") - return path, elapsed + root = Path( + retry_call( + lambda: snapshot_download( + repo, + repo_type="dataset", + revision=revision, + allow_patterns=[shard["path"] for shard in shards], + local_dir=download_root, + max_workers=len(shards), + etag_timeout=etag_timeout, + ), + attempts, + ) + ) + elapsed = max(time.monotonic() - started, 0.001) + total_bytes = sum(shard["size"] for shard in shards) + results = [] + for shard in shards: + path = root / shard["path"] + actual_size = path.stat().st_size + if actual_size != shard["size"]: + raise RuntimeError(f"downloaded {actual_size} bytes for {path}, expected {shard['size']}") + if shard.get("sha256") and file_sha256(path) != shard["sha256"]: + path.unlink(missing_ok=True) + raise RuntimeError(f"SHA-256 mismatch for {shard['path']}") + results.append((shard, path, elapsed * shard["size"] / max(total_bytes, 1))) + return results def parquet_zstd6(source, destination): @@ -231,8 +275,7 @@ def parquet_zstd6(source, destination): partial = destination.with_suffix(destination.suffix + ".part") parquet = pq.ParquetFile(source) started = time.monotonic() - writer = pq.ParquetWriter(partial, parquet.schema_arrow, - compression="zstd", compression_level=6) + writer = pq.ParquetWriter(partial, parquet.schema_arrow, compression="zstd", compression_level=6) try: for batch in parquet.iter_batches(batch_size=PARQUET_BATCH_ROWS): writer.write_batch(batch, row_group_size=PARQUET_BATCH_ROWS) @@ -286,12 +329,17 @@ def existing_files(self, prefix=""): """Return existing sink files keyed by destination path.""" return {} + def upload_batch(self, items): + """Upload a claimed batch; backends may override this with a native batch call.""" + return [self.upload(**item) for item in items] + class HuggingFaceBatchUploader(Uploader): """Preupload outputs and commit up to a fixed file count per format.""" - def __init__(self, api, repo, revision, batch_files, totals, attempts=5, timeout=30, - batch_bytes=None, planned_batches=None): + def __init__( + self, api, repo, revision, batch_files, totals, attempts=5, timeout=30, batch_bytes=None, planned_batches=None + ): self.api = api self.repo = repo self.revision = revision @@ -300,48 +348,87 @@ def __init__(self, api, repo, revision, batch_files, totals, attempts=5, timeout self.attempts = attempts self.timeout = timeout self.batch_bytes = batch_bytes - self.planned_batches = { - batch["id"]: batch for batch in (planned_batches or []) - } + self.planned_batches = {batch["id"]: batch for batch in (planned_batches or [])} self.batch_for_path = { - path: batch["id"] - for batch in self.planned_batches.values() - for path in batch["destination_paths"] + path: batch["id"] for batch in self.planned_batches.values() for path in batch["destination_paths"] } self.parent_commit = resolve_dataset_revision(api, repo, revision, attempts, timeout) - self.pending = ({batch_id: [] for batch_id in self.planned_batches} - if self.planned_batches else {fmt: [] for fmt in totals}) + self.pending = ( + {batch_id: [] for batch_id in self.planned_batches} if self.planned_batches else {fmt: [] for fmt in totals} + ) self.lock = threading.Lock() def upload(self, local_path, destination_path, *, format_name, ordinal, **metadata): + return self.upload_batch( + [ + { + "local_path": local_path, + "destination_path": destination_path, + "format_name": format_name, + "ordinal": ordinal, + } + ] + )[0] + + def upload_batch(self, items): from huggingface_hub import CommitOperationAdd - operation = CommitOperationAdd(path_in_repo=destination_path, path_or_fileobj=local_path) started = time.monotonic() - retry_call(lambda: self.api.preupload_lfs_files( - self.repo, additions=[operation], repo_type="dataset", - revision=self.revision, num_threads=1), self.attempts) - entry = {"operation": operation, "local_path": str(local_path), - "hub_path": destination_path, "size_bytes": local_path.stat().st_size, - "format": format_name, "ordinal": ordinal} + operations = [ + CommitOperationAdd(path_in_repo=item["destination_path"], path_or_fileobj=item["local_path"]) + for item in items + ] + retry_call( + lambda: self.api.preupload_lfs_files( + self.repo, + additions=operations, + repo_type="dataset", + revision=self.revision, + num_threads=len(operations), + ), + self.attempts, + ) + results = [] with self.lock: - pending_key = self.batch_for_path.get(destination_path, format_name) - if self.planned_batches and pending_key == format_name: - raise RuntimeError(f"destination is not present in the action plan: {destination_path}") - pending = self.pending[pending_key] - pending.append(entry) - if self.planned_batches: - expected = len(self.planned_batches[pending_key]["destination_paths"]) - if len(pending) == expected: - return self._flush_format_locked(pending_key) - return {"status": "preuploaded", "hub_path": destination_path, - "seconds": time.monotonic() - started} - pending_bytes = sum(item["size_bytes"] for item in pending) - if (len(pending) >= self.batch_files - or (self.batch_bytes is not None and pending_bytes >= self.batch_bytes)): - return self._flush_format_locked(format_name) - return {"status": "preuploaded", "hub_path": destination_path, - "seconds": time.monotonic() - started} + for item, operation in zip(items, operations, strict=True): + destination_path = item["destination_path"] + format_name = item["format_name"] + entry = { + "operation": operation, + "local_path": str(item["local_path"]), + "hub_path": destination_path, + "size_bytes": item["local_path"].stat().st_size, + "format": format_name, + "ordinal": item["ordinal"], + } + pending_key = self.batch_for_path.get(destination_path, format_name) + if self.planned_batches and pending_key == format_name: + raise RuntimeError(f"destination is not present in the action plan: {destination_path}") + pending = self.pending[pending_key] + pending.append(entry) + if self.planned_batches: + expected = len(self.planned_batches[pending_key]["destination_paths"]) + if len(pending) == expected: + results.append(self._flush_format_locked(pending_key)) + else: + results.append( + { + "status": "preuploaded", + "hub_path": destination_path, + "seconds": time.monotonic() - started, + } + ) + continue + pending_bytes = sum(entry["size_bytes"] for entry in pending) + if len(pending) >= self.batch_files or ( + self.batch_bytes is not None and pending_bytes >= self.batch_bytes + ): + results.append(self._flush_format_locked(format_name)) + else: + results.append( + {"status": "preuploaded", "hub_path": destination_path, "seconds": time.monotonic() - started} + ) + return results def flush(self): with self.lock: @@ -364,61 +451,81 @@ def _flush_format_locked(self, pending_key): start = planned["start"] if planned else batch[0]["ordinal"] end = planned["end"] if planned else batch[-1]["ordinal"] total_files = planned["total_files"] if planned else self.totals[format_name] - message = (planned["commit_message"] if planned else - f"Upload {format_name} files {start}-{end} of {total_files}") + message = planned["commit_message"] if planned else f"Upload {format_name} files {start}-{end} of {total_files}" operations = [entry["operation"] for entry in batch] result = None for attempt in range(1, self.attempts + 1): try: result = self.api.create_commit( - self.repo, operations=operations, repo_type="dataset", - revision=self.revision, parent_commit=self.parent_commit, - commit_message=message) + self.repo, + operations=operations, + repo_type="dataset", + revision=self.revision, + parent_commit=self.parent_commit, + commit_message=message, + ) break except Exception: - remote = list_repository_files(self.api, self.repo, self.revision, - attempts=1) - committed = all(remote.get(entry["hub_path"], {}).get("size") - == entry["size_bytes"] for entry in batch) + remote = list_repository_files(self.api, self.repo, self.revision, attempts=1) + committed = all(remote.get(entry["hub_path"], {}).get("size") == entry["size_bytes"] for entry in batch) if committed: commit_id = resolve_dataset_revision( - self.api, self.repo, self.revision, attempts=1, timeout=self.timeout) - result = type("CommitResult", (), { - "oid": commit_id, - "commit_url": f"https://huggingface.co/datasets/{self.repo}/commit/{commit_id}", - })() + self.api, self.repo, self.revision, attempts=1, timeout=self.timeout + ) + result = type( + "CommitResult", + (), + { + "oid": commit_id, + "commit_url": f"https://huggingface.co/datasets/{self.repo}/commit/{commit_id}", + }, + )() break if attempt == self.attempts: raise time.sleep(2 ** (attempt - 1)) operations = [] for entry in batch: - operation = CommitOperationAdd(path_in_repo=entry["hub_path"], - path_or_fileobj=entry["local_path"]) - retry_call(lambda op=operation: self.api.preupload_lfs_files( - self.repo, additions=[op], repo_type="dataset", - revision=self.revision, num_threads=1), self.attempts) + operation = CommitOperationAdd(path_in_repo=entry["hub_path"], path_or_fileobj=entry["local_path"]) + retry_call( + lambda op=operation: self.api.preupload_lfs_files( + self.repo, additions=[op], repo_type="dataset", revision=self.revision, num_threads=1 + ), + self.attempts, + ) operations.append(operation) if result is None: raise RuntimeError(f"failed to commit {format_name} files {start}-{end}") - committed = [{key: entry[key] - for key in ("local_path", "hub_path", "size_bytes", "format", "ordinal")} - for entry in batch] + committed = [ + {key: entry[key] for key in ("local_path", "hub_path", "size_bytes", "format", "ordinal")} + for entry in batch + ] total_bytes = sum(entry["size_bytes"] for entry in batch) self.pending[pending_key] = [] self.parent_commit = result.oid - return {"status": "complete", "url": result.commit_url, - "commit_id": result.oid, "committed": committed, "size_bytes": total_bytes, - "format": format_name, "start": start, "end": end, - "total_files": total_files, "commit_message": message} + return { + "status": "complete", + "url": result.commit_url, + "commit_id": result.oid, + "committed": committed, + "size_bytes": total_bytes, + "format": format_name, + "start": start, + "end": end, + "total_files": total_files, + "commit_message": message, + } def config(self): - return {"type": "huggingface-preupload-batch", "repo": self.repo, - "revision": self.revision, "batch_files": self.batch_files} + return { + "type": "huggingface-preupload-batch", + "repo": self.repo, + "revision": self.revision, + "batch_files": self.batch_files, + } def existing_files(self, prefix=""): - return list_repository_files(self.api, self.repo, self.revision, - prefix, self.attempts) + return list_repository_files(self.api, self.repo, self.revision, prefix, self.attempts) class LocalCopyUploader(Uploader): @@ -444,9 +551,14 @@ def upload(self, local_path, destination_path, **metadata): chunks += 1 bytes_copied += len(block) partial.replace(output) - return {"status": "complete", "hub_path": destination_path, - "url": output.as_uri(), "seconds": time.monotonic() - started, - "chunks": chunks, "bytes_uploaded": bytes_copied} + return { + "status": "complete", + "hub_path": destination_path, + "url": output.as_uri(), + "seconds": time.monotonic() - started, + "chunks": chunks, + "bytes_uploaded": bytes_copied, + } def config(self): return {"type": "local-copy", "destination": str(self.destination)} @@ -455,8 +567,11 @@ def existing_files(self, prefix=""): root = self.destination / prefix.strip("/") if not root.exists(): return {} - return {str(path.relative_to(self.destination)): {"size": path.stat().st_size} - for path in root.rglob("*") if path.is_file() and not path.name.endswith(".part")} + return { + str(path.relative_to(self.destination)): {"size": path.stat().st_size} + for path in root.rglob("*") + if path.is_file() and not path.name.endswith(".part") + } def upload_then_maybe_delete(uploader, local_path, destination_path, delete_after_upload): @@ -485,15 +600,30 @@ def fits_download_buffer(inflight_bytes, next_bytes, limit_bytes): return inflight_bytes == 0 or inflight_bytes + next_bytes <= limit_bytes -def completed_futures(futures): - """Yield (key, future) in completion order, removing each from the input mapping.""" - while futures: - done, _ = concurrent.futures.wait(futures.values(), - return_when=concurrent.futures.FIRST_COMPLETED) - for completed in done: - key = next(key for key, future in futures.items() if future is completed) - futures.pop(key) - yield key, completed +def adapt_concurrency(current, maximum, recent_rates, rate): + """Adjust bounded concurrency from a short throughput history.""" + previous = sum(recent_rates) / len(recent_rates) if recent_rates else 0 + recent_rates.append(rate) + if not previous or rate >= previous * 0.95: + return min(maximum, current + 1) + if len(recent_rates) == recent_rates.maxlen and rate < previous * 0.75: + return max(1, current // 2) + return current + + +def choose_work_stage(*, upload_waiting, upload_high, download_available, convert_waiting, active): + """Choose work using disk safety, refill, drain, then conversion priority.""" + if upload_high: + return "upload" + if download_available: + return "download" + if upload_waiting: + return "upload" + if convert_waiting: + return "convert" + if not active: + return "done" + return None def needs_source_download(file_state, outputs): @@ -501,9 +631,8 @@ def needs_source_download(file_state, outputs): for fmt, local_path in outputs.items(): output_state = file_state.get("outputs", {}).get(fmt, {}) upload_complete = output_state.get("upload", {}).get("status") == "complete" - local_complete = ( - local_path.exists() - and local_path.stat().st_size == output_state.get("metrics", {}).get("size_bytes") + local_complete = local_path.exists() and local_path.stat().st_size == output_state.get("metrics", {}).get( + "size_bytes" ) encoding_complete = output_state.get("status") == "complete" and local_complete if not upload_complete and not encoding_complete: @@ -539,8 +668,9 @@ def start(self): def write(self, final=False): value = self.snapshot() - value.update(schema_version=1, elapsed_seconds=time.monotonic() - self.started, - updated_at_unix=time.time(), final=final) + value.update( + schema_version=1, elapsed_seconds=time.monotonic() - self.started, updated_at_unix=time.time(), final=final + ) atomic_json(self.path, value) def _run(self): @@ -571,13 +701,21 @@ def checkpoint_records(checkpoint): if outputs[fmt].get("status") == "complete": source_size = checkpoint["files"][source_name]["size"] metrics = outputs[fmt]["metrics"] - records.append({"source": source_name, "format": fmt, - "source_parquet_bytes": source_size, - "size_bytes": metrics["size_bytes"], - "ratio_to_parquet": metrics["size_bytes"] / source_size, - "bytes_saved_vs_parquet": source_size - metrics["size_bytes"], - **{key: value for key, value in metrics.items() - if key not in ("size_bytes", "ratio_to_download")}}) + records.append( + { + "source": source_name, + "format": fmt, + "source_parquet_bytes": source_size, + "size_bytes": metrics["size_bytes"], + "ratio_to_parquet": metrics["size_bytes"] / source_size, + "bytes_saved_vs_parquet": source_size - metrics["size_bytes"], + **{ + key: value + for key, value in metrics.items() + if key not in ("size_bytes", "ratio_to_download") + }, + } + ) return records @@ -588,8 +726,17 @@ def write_reports(checkpoint, output_dir, selection): with (metrics_dir / "files.jsonl").open("w") as output: for record in records: output.write(json.dumps(record, sort_keys=True) + "\n") - fields = ["source", "format", "source_parquet_bytes", "size_bytes", "ratio_to_parquet", - "bytes_saved_vs_parquet", "seconds", "mb_per_second", "sha256"] + fields = [ + "source", + "format", + "source_parquet_bytes", + "size_bytes", + "ratio_to_parquet", + "bytes_saved_vs_parquet", + "seconds", + "mb_per_second", + "sha256", + ] with (metrics_dir / "files.csv").open("w", newline="") as output: writer = csv.DictWriter(output, fieldnames=fields) writer.writeheader() @@ -598,55 +745,75 @@ def write_reports(checkpoint, output_dir, selection): source_totals = {} for record in records: totals[record["format"]] = totals.get(record["format"], 0) + record["size_bytes"] - source_totals[record["format"]] = (source_totals.get(record["format"], 0) - + record["source_parquet_bytes"]) + source_totals[record["format"]] = source_totals.get(record["format"], 0) + record["source_parquet_bytes"] comparisons = { - fmt: {"source_parquet_bytes": source_totals[fmt], "encoded_bytes": size, - "ratio_to_parquet": size / source_totals[fmt], - "bytes_saved_vs_parquet": source_totals[fmt] - size} + fmt: { + "source_parquet_bytes": source_totals[fmt], + "encoded_bytes": size, + "ratio_to_parquet": size / source_totals[fmt], + "bytes_saved_vs_parquet": source_totals[fmt] - size, + } for fmt, size in totals.items() } - summary = {"selected_source_bytes": sum(item["size"] for item in selection), - "selected_files": len(selection), "format_size_bytes": totals, - "format_comparison": comparisons} + summary = { + "selected_source_bytes": sum(item["size"] for item in selection), + "selected_files": len(selection), + "format_size_bytes": totals, + "format_comparison": comparisons, + } atomic_json(metrics_dir / "summary.json", summary) safe_print(json.dumps(summary, indent=2, sort_keys=True)) def metric(source_name, fmt, path, input_size, seconds): size = path.stat().st_size - return {"source": source_name, "format": fmt, "size_bytes": size, - "ratio_to_download": size / input_size, "seconds": seconds, - "mb_per_second": (input_size / 1_000_000 / seconds) if seconds else None, - "sha256": file_sha256(path)} + return { + "source": source_name, + "format": fmt, + "size_bytes": size, + "ratio_to_download": size / input_size, + "seconds": seconds, + "mb_per_second": (input_size / 1_000_000 / seconds) if seconds else None, + "sha256": file_sha256(path), + } def add_selection_arguments(parser): - parser.add_argument("--mode", choices=("sample", "first", "all"), default="sample", - help="sample evenly, take the first --limit shards, or stream every shard") - parser.add_argument("--limit", type=int, default=10, - help="number of ordered shards selected by --mode first") + parser.add_argument( + "--mode", + choices=("sample", "first", "all"), + default="sample", + help="sample evenly, take the first --limit shards, or stream every shard", + ) + parser.add_argument("--limit", type=int, default=10, help="number of ordered shards selected by --mode first") parser.add_argument("--target-size", type=parse_size, default=DEFAULT_TARGET_BYTES) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--repo", required=True, help="source Hugging Face dataset repository") - parser.add_argument("--revision", required=True, - help="source branch, tag, or commit (resolved before planning)") - parser.add_argument("--prefix", required=True, - help="source repository folder; use / to scan the repository root") - parser.add_argument("--include", default="*.parquet", - help="glob matched against paths below --prefix (default: *.parquet)") - parser.add_argument("--filter", action="append", default=[], - help="repeatable full repository-path glob, e.g. 'sample/10BT/*'") - parser.add_argument("--formats", default="parquet-zstd6,vortex,vortex-compact", - help="comma-separated outputs: parquet-zstd6,vortex,vortex-compact") - parser.add_argument("--upload-repo", required=True, - help="destination Hugging Face dataset repository") + parser.add_argument("--revision", required=True, help="source branch, tag, or commit (resolved before planning)") + parser.add_argument("--prefix", required=True, help="source repository folder; use / to scan the repository root") + parser.add_argument( + "--include", default="*.parquet", help="glob matched against paths below --prefix (default: *.parquet)" + ) + parser.add_argument( + "--filter", action="append", default=[], help="repeatable full repository-path glob, e.g. 'sample/10BT/*'" + ) + parser.add_argument( + "--formats", + default="parquet-zstd6,vortex,vortex-compact", + help="comma-separated outputs: parquet-zstd6,vortex,vortex-compact", + ) + parser.add_argument("--upload-repo", required=True, help="destination Hugging Face dataset repository") parser.add_argument("--upload-revision", default="main") parser.add_argument("--upload-prefix", default="") - parser.add_argument("--upload-batch-files", type=int, default=100, - help="maximum actions in each planned Hugging Face commit") - parser.add_argument("--upload-batch-size", type=parse_size, default=100_000_000_000, - help="maximum planned source bytes represented by one target commit") + parser.add_argument( + "--upload-batch-files", type=int, default=100, help="maximum actions in each planned Hugging Face commit" + ) + parser.add_argument( + "--upload-batch-size", + type=parse_size, + default=100_000_000_000, + help="maximum planned source bytes represented by one target commit", + ) parser.add_argument("--hub-attempts", type=int, default=8) parser.add_argument("--hub-timeout", type=int, default=30) @@ -654,24 +821,24 @@ def add_selection_arguments(parser): def add_operational_arguments(parser): parser.add_argument("--output-dir", type=Path, default=DEFAULT_DATA_DIR) parser.add_argument("--vx", type=Path, default=Path("target/release/vx")) - parser.add_argument("--format-workers", type=int, default=3) - parser.add_argument("--transcode-workers", type=int, default=max(1, os.cpu_count() or 1), - help="maximum single-core transcodes active across downloaded shards") - parser.add_argument("--shard-workers", type=int, default=2, - help="initial downloads (deprecated alias)") - parser.add_argument("--download-initial-concurrency", type=int) - parser.add_argument("--download-max-concurrency", type=int, default=8) + parser.add_argument("--workers", type=int, help="workers in the unified download/upload/convert pool") + parser.add_argument("--format-workers", type=int, default=3, help="deprecated; formats run sequentially per job") + parser.add_argument( + "--transcode-workers", + type=int, + help="deprecated alias for --workers", + ) + parser.add_argument("--shard-workers", type=int, default=2, help="deprecated download batch-size alias") + parser.add_argument("--download-initial-concurrency", type=int, help="initial adaptive download batch size") + parser.add_argument("--download-max-concurrency", type=int, default=8, help="maximum download batch size") parser.add_argument("--download-buffer-files", type=int, default=100) - parser.add_argument("--download-buffer-size", type=parse_size, - default=DEFAULT_DOWNLOAD_BUFFER_BYTES) - parser.add_argument("--upload-workers", type=int, default=2) - parser.add_argument("--upload-max-concurrency", type=int, default=8) - parser.add_argument("--upload-local-dir", type=Path, - help="copy outputs to a local sink instead of Hugging Face") + parser.add_argument("--download-buffer-size", type=parse_size, default=DEFAULT_DOWNLOAD_BUFFER_BYTES) + parser.add_argument("--upload-workers", type=int, default=2, help="initial adaptive upload batch size") + parser.add_argument("--upload-max-concurrency", type=int, default=8, help="maximum upload batch size") + parser.add_argument("--upload-local-dir", type=Path, help="copy outputs to a local sink instead of Hugging Face") parser.add_argument("--upload-buffer-files", type=int, default=100) parser.add_argument("--upload-batch-files", type=int, default=100) - parser.add_argument("--upload-buffer-size", type=parse_size, - default=DEFAULT_UPLOAD_BUFFER_BYTES) + parser.add_argument("--upload-buffer-size", type=parse_size, default=DEFAULT_UPLOAD_BUFFER_BYTES) parser.add_argument("--hub-attempts", type=int, default=8) parser.add_argument("--hub-timeout", type=int, default=30) parser.add_argument("--xet-range-gets", type=int, default=4) @@ -717,8 +884,7 @@ def plan_remote_metadata(remote): lfs = remote.get("lfs") or {} if not isinstance(lfs, dict): lfs = getattr(lfs, "__dict__", {}) - return {"size": remote.get("size"), "oid": remote.get("oid"), - "sha256": lfs.get("sha256")} + return {"size": remote.get("size"), "oid": remote.get("oid"), "sha256": lfs.get("sha256")} def create_action_plan(api, args): @@ -731,19 +897,18 @@ def create_action_plan(api, args): prefix = args.prefix.strip("/") require_xet_repository(api, args.repo, args.revision, args.hub_attempts) require_xet_repository(api, args.upload_repo, args.upload_revision, args.hub_attempts) - source_commit = resolve_dataset_revision( - api, args.repo, args.revision, args.hub_attempts, args.hub_timeout) + source_commit = resolve_dataset_revision(api, args.repo, args.revision, args.hub_attempts, args.hub_timeout) destination_commit = resolve_dataset_revision( - api, args.upload_repo, args.upload_revision, args.hub_attempts, args.hub_timeout) - shards = list_shards(api, args.repo, source_commit, prefix, args.include, - args.filter, args.hub_attempts) + api, args.upload_repo, args.upload_revision, args.hub_attempts, args.hub_timeout + ) + shards = list_shards(api, args.repo, source_commit, prefix, args.include, args.filter, args.hub_attempts) selection = select_shards(shards, args.mode, args.limit, args.target_size, args.seed) existing = {} for fmt in formats: - format_prefix = "/".join( - part for part in (args.upload_prefix.strip("/"), fmt) if part) - existing.update(list_repository_files( - api, args.upload_repo, destination_commit, format_prefix, args.hub_attempts)) + format_prefix = "/".join(part for part in (args.upload_prefix.strip("/"), fmt) if part) + existing.update( + list_repository_files(api, args.upload_repo, destination_commit, format_prefix, args.hub_attempts) + ) chunks = [] counts = {"create": 0, "skip": 0} for ordinal, shard in enumerate(selection, 1): @@ -751,8 +916,7 @@ def create_action_plan(api, args): for fmt in formats: sink_path = destination_path(shard["path"], fmt, args.upload_prefix) remote = existing.get(sink_path) - action = {"action": "skip" if remote else "create", "format": fmt, - "destination_path": sink_path} + action = {"action": "skip" if remote else "create", "format": fmt, "destination_path": sink_path} if remote: action["existing"] = plan_remote_metadata(remote) counts[action["action"]] += 1 @@ -762,7 +926,8 @@ def create_action_plan(api, args): for fmt in formats: format_actions = [ (chunk, action) - for chunk in chunks for action in chunk["actions"] + for chunk in chunks + for action in chunk["actions"] if action["format"] == fmt and action["action"] == "create" ] total = len(format_actions) @@ -771,8 +936,7 @@ def create_action_plan(api, args): group_bytes = 0 for member in format_actions: source_bytes = member[0]["source"]["size"] - if group and (len(group) >= args.upload_batch_files - or group_bytes + source_bytes > upload_batch_size): + if group and (len(group) >= args.upload_batch_files or group_bytes + source_bytes > upload_batch_size): groups.append(group) group = [] group_bytes = 0 @@ -789,50 +953,63 @@ def create_action_plan(api, args): for _, action in members: action["upload_batch"] = batch_id action["commit_message"] = message - upload_batches.append({ - "id": batch_id, - "format": fmt, - "start": start, - "end": end, - "total_files": total, - "commit_message": message, - "source_ordinals": [chunk["ordinal"] for chunk, _ in members], - "destination_paths": [action["destination_path"] for _, action in members], - "planned_source_bytes": sum(chunk["source"]["size"] for chunk, _ in members), - }) + upload_batches.append( + { + "id": batch_id, + "format": fmt, + "start": start, + "end": end, + "total_files": total, + "commit_message": message, + "source_ordinals": [chunk["ordinal"] for chunk, _ in members], + "destination_paths": [action["destination_path"] for _, action in members], + "planned_source_bytes": sum(chunk["source"]["size"] for chunk, _ in members), + } + ) offset = end return { "version": PLAN_VERSION, "kind": "hf-sync-plan", "created_at_unix_seconds": time.time(), - "source": {"repo": args.repo, "requested_revision": args.revision, - "revision": source_commit, "prefix": prefix, "include": args.include, - "filters": args.filter}, - "destination": {"repo": args.upload_repo, - "requested_revision": args.upload_revision, - "revision": destination_commit, - "prefix": args.upload_prefix.strip("/")}, - "selection": {"mode": args.mode, "limit": args.limit, - "target_bytes": args.target_size, "seed": args.seed}, + "source": { + "repo": args.repo, + "requested_revision": args.revision, + "revision": source_commit, + "prefix": prefix, + "include": args.include, + "filters": args.filter, + }, + "destination": { + "repo": args.upload_repo, + "requested_revision": args.upload_revision, + "revision": destination_commit, + "prefix": args.upload_prefix.strip("/"), + }, + "selection": {"mode": args.mode, "limit": args.limit, "target_bytes": args.target_size, "seed": args.seed}, "formats": list(formats), "work_chunks": chunks, "upload_batches": upload_batches, - "summary": {"source_files": len(selection), - "source_bytes": sum(shard["size"] for shard in selection), - "create_actions": counts["create"], "skip_actions": counts["skip"]}, + "summary": { + "source_files": len(selection), + "source_bytes": sum(shard["size"] for shard in selection), + "create_actions": counts["create"], + "skip_actions": counts["skip"], + }, } def load_action_plan(path): with path.open() as source: plan = json.load(source) - if (plan.get("version") != PLAN_VERSION - or plan.get("kind") != "hf-sync-plan" - or not isinstance(plan.get("work_chunks"), list)): + if ( + plan.get("version") != PLAN_VERSION + or plan.get("kind") != "hf-sync-plan" + or not isinstance(plan.get("work_chunks"), list) + ): raise RuntimeError(f"unsupported or invalid action plan: {path}") batches = plan.get("upload_batches") if not isinstance(batches, list): - raise RuntimeError(f"unsupported or invalid action plan: {path}") + raise RuntimeError(f"unsupported or invalid action plan: {path}") # noqa: TRY004 batch_paths = {} for batch in batches: batch_id = batch.get("id") @@ -845,7 +1022,7 @@ def load_action_plan(path): create_paths = set() for chunk in plan["work_chunks"]: if not isinstance(chunk.get("source"), dict) or not isinstance(chunk.get("actions"), list): - raise RuntimeError(f"unsupported or invalid action plan: {path}") + raise RuntimeError(f"unsupported or invalid action plan: {path}") # noqa: TRY004 for action in chunk["actions"]: if action.get("action") not in ("create", "skip"): raise RuntimeError(f"unsupported action in {path}: {action.get('action')}") @@ -892,8 +1069,12 @@ def main(): selection = apply_plan_arguments(args, plan) if args.format_workers < 1 or args.format_workers > 3: raise RuntimeError("--format-workers must be between 1 and 3") - if args.transcode_workers < 1: - raise RuntimeError("--transcode-workers must be positive") + if args.workers is None: + args.workers = args.transcode_workers or max(1, os.cpu_count() or 1) + elif args.transcode_workers is not None: + raise RuntimeError("use --workers or deprecated --transcode-workers, not both") + if args.workers < 1: + raise RuntimeError("--workers must be positive") if args.shard_workers < 1: raise RuntimeError("--shard-workers must be positive") if args.download_initial_concurrency is None: @@ -946,7 +1127,8 @@ def main(): raise RuntimeError(f"vx binary not found: {vx}; build vortex-tui with unstable_encodings") available_cpus = sorted(os.sched_getaffinity(0)) VX_CPU_SLOTS = queue.Queue() - for cpu in available_cpus[:min(args.transcode_workers, len(available_cpus))]: + cpu_slots = available_cpus[: min(args.workers, len(available_cpus))] + for cpu in cpu_slots: VX_CPU_SLOTS.put(cpu) api = hub_api() require_xet_repository(api, args.repo, args.revision, args.hub_attempts) @@ -958,64 +1140,85 @@ def main(): flush=True, ) safe_print( - f"Transfer buffers: downloads={args.download_initial_concurrency}-" - f"{args.download_max_concurrency} workers/" - f"{args.download_buffer_size / 1e9:.1f} GB, uploads={args.upload_workers}-" - f"{args.upload_max_concurrency} workers/" + f"Unified pool: workers={args.workers}, download_batch={args.download_initial_concurrency}-" + f"{args.download_max_concurrency} files/{args.download_buffer_size / 1e9:.1f} GB, " + f"upload_batch={args.upload_workers}-{args.upload_max_concurrency} files/" f"{args.upload_buffer_size / 1e9:.1f} GB", flush=True, ) source_commit = plan["source"]["revision"] checkpoint_path = args.output_dir / "checkpoint.json" checkpoint = load_checkpoint(checkpoint_path) - current_destination = (resolve_dataset_revision( - api, args.upload_repo, args.upload_revision, args.hub_attempts, args.hub_timeout) - if args.upload_local_dir is None else plan["destination"]["revision"]) + current_destination = ( + resolve_dataset_revision(api, args.upload_repo, args.upload_revision, args.hub_attempts, args.hub_timeout) + if args.upload_local_dir is None + else plan["destination"]["revision"] + ) resumable_commits = { output.get("upload", {}).get("commit_id") for state in checkpoint.get("files", {}).values() for output in state.get("outputs", {}).values() } resumable_commits.discard(None) - if (current_destination != plan["destination"]["revision"] - and current_destination not in resumable_commits): + if current_destination != plan["destination"]["revision"] and current_destination not in resumable_commits: raise RuntimeError( f"stale plan: destination {args.upload_repo}@{args.upload_revision} changed from " - f"{plan['destination']['revision']} to {current_destination}; create a new plan") + f"{plan['destination']['revision']} to {current_destination}; create a new plan" + ) create_totals = { - fmt: sum(action["action"] == "create" - for chunk in plan["work_chunks"] for action in chunk["actions"] - if action["format"] == fmt) + fmt: sum( + action["action"] == "create" + for chunk in plan["work_chunks"] + for action in chunk["actions"] + if action["format"] == fmt + ) for fmt in formats } if args.upload_local_dir is not None: uploader = LocalCopyUploader(args.upload_local_dir) else: uploader = HuggingFaceBatchUploader( - api, args.upload_repo, args.upload_revision, args.upload_batch_files, - create_totals, args.hub_attempts, args.hub_timeout, + api, + args.upload_repo, + args.upload_revision, + args.upload_batch_files, + create_totals, + args.hub_attempts, + args.hub_timeout, batch_bytes=max(1, args.upload_buffer_size // len(formats)), - planned_batches=plan.get("upload_batches", [])) - run_config = {"repo": args.repo, "requested_revision": args.revision, - "revision": source_commit, "prefix": args.prefix, - "include": args.include, "filters": args.filter, "mode": args.mode, "limit": args.limit, - "target_bytes": args.target_size, "formats": list(formats), - "xet_range_gets": args.xet_range_gets, - "xet_high_performance": args.xet_high_performance, - "xet_cache": str(args.xet_cache), - "seed": args.seed, "uploader": uploader.config() if uploader else None, - "upload_prefix": args.upload_prefix, - "delete_after_upload": args.delete_after_upload, "files": selection} + planned_batches=plan.get("upload_batches", []), + ) + run_config = { + "repo": args.repo, + "requested_revision": args.revision, + "revision": source_commit, + "prefix": args.prefix, + "include": args.include, + "filters": args.filter, + "mode": args.mode, + "limit": args.limit, + "target_bytes": args.target_size, + "formats": list(formats), + "xet_range_gets": args.xet_range_gets, + "xet_high_performance": args.xet_high_performance, + "xet_cache": str(args.xet_cache), + "seed": args.seed, + "uploader": uploader.config() if uploader else None, + "upload_prefix": args.upload_prefix, + "delete_after_upload": args.delete_after_upload, + "files": selection, + } previous_config = checkpoint.get("config") - if (previous_config is not None - and resumability_config(previous_config) != resumability_config(run_config)): + if previous_config is not None and resumability_config(previous_config) != resumability_config(run_config): raise RuntimeError( - f"arguments do not match {checkpoint_path}; resume with the same arguments or use another --output-dir") + f"arguments do not match {checkpoint_path}; resume with the same arguments or use another --output-dir" + ) checkpoint["config"] = run_config checkpoint_lock = threading.Lock() existing_sink_files = { action["destination_path"]: action["existing"] - for chunk in plan["work_chunks"] for action in chunk["actions"] + for chunk in plan["work_chunks"] + for action in chunk["actions"] if action["action"] == "skip" } for shard in selection: @@ -1027,12 +1230,22 @@ def main(): continue output_state = state["outputs"].setdefault(fmt, {}) output_state["status"] = "complete" - output_state["upload"] = {"status": "complete", "hub_path": sink_path, - "size_bytes": remote.get("size"), "discovered": True} - output_state.setdefault("metrics", { + output_state["upload"] = { + "status": "complete", + "hub_path": sink_path, "size_bytes": remote.get("size"), - "ratio_to_download": remote.get("size") / shard["size"] if remote.get("size") else None, - "seconds": None, "mb_per_second": None, "sha256": remote.get("oid")}) + "discovered": True, + } + output_state.setdefault( + "metrics", + { + "size_bytes": remote.get("size"), + "ratio_to_download": remote.get("size") / shard["size"] if remote.get("size") else None, + "seconds": None, + "mb_per_second": None, + "sha256": remote.get("oid"), + }, + ) atomic_json(checkpoint_path, checkpoint) atomic_json(args.output_dir / "metrics" / "selection.json", run_config) @@ -1043,11 +1256,13 @@ def save_checkpoint(): def shard_paths(shard): source_name = shard["path"] short_hash = hashlib.sha256(source_name.encode()).hexdigest()[:10] - stem = source_name[:-len(".parquet")].replace("/", "__") + "__" + short_hash + stem = source_name[: -len(".parquet")].replace("/", "__") + "__" + short_hash raw_path = args.output_dir / "downloads" / source_name - available = {"parquet-zstd6": args.output_dir / "parquet-zstd6" / f"{stem}.parquet", - "vortex": args.output_dir / "vortex" / f"{stem}.vortex", - "vortex-compact": args.output_dir / "vortex-compact" / f"{stem}.vortex"} + available = { + "parquet-zstd6": args.output_dir / "parquet-zstd6" / f"{stem}.parquet", + "vortex": args.output_dir / "vortex" / f"{stem}.vortex", + "vortex-compact": args.output_dir / "vortex-compact" / f"{stem}.vortex", + } return raw_path, {fmt: available[fmt] for fmt in formats} def shard_needs_download(shard): @@ -1055,73 +1270,30 @@ def shard_needs_download(shard): return needs_source_download(checkpoint["files"][shard["path"]], outputs) missing_shards = [shard for shard in selection if shard_needs_download(shard)] - download_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=min(args.download_max_concurrency, max(1, len(missing_shards)))) - upload_executor = concurrent.futures.ThreadPoolExecutor(max_workers=args.upload_max_concurrency) - pending_uploads = [] - upload_control = threading.Condition() - upload_active = 0 - upload_concurrency = args.upload_workers + work_condition = threading.Condition() + download_queue = deque(missing_shards) + conversion_queue = deque() + upload_queue = deque() + active_work = {"download": 0, "convert": 0, "upload": 0} + download_reserved_files = 0 + download_reserved_bytes = 0 + upload_reserved_bytes = 0 + worker_failure = [] + download_batch_size = args.download_initial_concurrency + upload_batch_size = args.upload_workers recent_upload_rates = deque(maxlen=4) - download_futures = {} - download_sizes = {} - missing_index = 0 - download_concurrency = args.download_initial_concurrency recent_download_rates = deque(maxlen=4) pipeline_started = time.monotonic() - transfer_totals = {"download_bytes": 0, "download_seconds": 0.0, - "upload_bytes": 0, "upload_seconds": 0.0, - "download_failures": 0, "upload_failures": 0} + transfer_totals = { + "download_bytes": 0, + "download_seconds": 0.0, + "upload_bytes": 0, + "upload_seconds": 0.0, + "download_failures": 0, + "upload_failures": 0, + } transfer_totals_lock = threading.Lock() - def run_adaptive_upload(operation, size_bytes): - nonlocal upload_active, upload_concurrency - with upload_control: - while upload_active >= upload_concurrency: - upload_control.wait() - upload_active += 1 - started = time.monotonic() - try: - return operation() - finally: - elapsed = max(time.monotonic() - started, 0.001) - rate = size_bytes / elapsed - with transfer_totals_lock: - transfer_totals["upload_bytes"] += size_bytes - transfer_totals["upload_seconds"] += elapsed - with upload_control: - previous = (sum(recent_upload_rates) / len(recent_upload_rates) - if recent_upload_rates else 0) - recent_upload_rates.append(rate) - if not previous or rate >= previous * 0.95: - upload_concurrency = min(args.upload_max_concurrency, upload_concurrency + 1) - elif len(recent_upload_rates) == recent_upload_rates.maxlen and rate < previous * 0.75: - upload_concurrency = max(1, upload_concurrency // 2) - upload_active -= 1 - upload_control.notify_all() - - def schedule_download(shard): - download_futures[shard["path"]] = download_executor.submit( - download_shard, args.repo, source_commit, shard, - args.output_dir / "downloads", args.hub_attempts, args.hub_timeout) - download_sizes[shard["path"]] = shard["size"] - - def fill_download_window(active_bytes=0): - nonlocal missing_index - while len(download_futures) < min(download_concurrency, args.download_max_concurrency, - args.download_buffer_files): - if missing_index >= len(missing_shards): - break - next_shard = missing_shards[missing_index] - inflight_bytes = active_bytes + sum(download_sizes.values()) - if not fits_download_buffer( - inflight_bytes, next_shard["size"], args.download_buffer_size): - break - schedule_download(next_shard) - missing_index += 1 - - fill_download_window() - def apply_committed_batch(result): if not result.get("committed"): return @@ -1135,26 +1307,34 @@ def apply_committed_batch(result): continue if args.delete_after_upload: Path(entry["local_path"]).unlink(missing_ok=True) - output_state["upload"] = {"status": "complete", "hub_path": sink_path, - "url": result["url"], - "commit_id": result.get("commit_id"), - "batch_size_bytes": result["size_bytes"], - "batch_start": result.get("start"), - "batch_end": result.get("end"), - "batch_total_files": result.get("total_files")} + output_state["upload"] = { + "status": "complete", + "hub_path": sink_path, + "url": result["url"], + "commit_id": result.get("commit_id"), + "batch_size_bytes": result["size_bytes"], + "batch_start": result.get("start"), + "batch_end": result.get("end"), + "batch_total_files": result.get("total_files"), + } output_state["local_deleted"] = args.delete_after_upload atomic_json(checkpoint_path, checkpoint) + def process_downloaded_shard(position, shard, source_name, raw, outputs, state): - all_jobs = {"parquet-zstd6": lambda: parquet_zstd6(raw, outputs["parquet-zstd6"]), - "vortex": lambda: run_vx(vx, raw, outputs["vortex"], "btrblocks"), - "vortex-compact": lambda: run_vx(vx, raw, outputs["vortex-compact"], "compact")} + all_jobs = { + "parquet-zstd6": lambda: parquet_zstd6(raw, outputs["parquet-zstd6"]), + "vortex": lambda: run_vx(vx, raw, outputs["vortex"], "btrblocks"), + "vortex-compact": lambda: run_vx(vx, raw, outputs["vortex-compact"], "compact"), + } jobs = tuple((fmt, all_jobs[fmt]) for fmt in formats) + def process_format(fmt, job): destination = outputs[fmt] output_state = state["outputs"].setdefault(fmt, {}) upload_complete = output_state.get("upload", {}).get("status") == "complete" - local_complete = (destination.exists() and - destination.stat().st_size == output_state.get("metrics", {}).get("size_bytes")) + local_complete = destination.exists() and destination.stat().st_size == output_state.get("metrics", {}).get( + "size_bytes" + ) durable_complete = upload_complete if uploader else local_complete complete = output_state.get("status") == "complete" and durable_complete if complete: @@ -1198,58 +1378,23 @@ def process_format(fmt, job): output_state["upload"] = {"status": "queued", "hub_path": sink_path} with checkpoint_lock: atomic_json(checkpoint_path, checkpoint) - - def record_upload_failure(error): - with transfer_totals_lock: - transfer_totals["upload_failures"] += 1 - with checkpoint_lock: - output_state["upload"] = {"status": "failed", "hub_path": sink_path, - "error": str(error)} - state["status"] = "failed" - state["error"] = f"upload {fmt}: {error}" - atomic_json(checkpoint_path, checkpoint) - - def do_upload_task(): - with checkpoint_lock: - output_state["upload"] = {"status": "uploading", "hub_path": sink_path} - atomic_json(checkpoint_path, checkpoint) - if isinstance(uploader, HuggingFaceBatchUploader): - try: - upload = uploader.upload(destination, sink_path, - format_name=fmt, ordinal=position) - except Exception as error: - record_upload_failure(error) - raise - if upload["status"] == "complete": - apply_committed_batch(upload) - safe_print(f" committed batch: {upload['url']}", flush=True) - else: - with checkpoint_lock: - output_state["upload"] = upload - atomic_json(checkpoint_path, checkpoint) - safe_print(f" preuploaded: {sink_path}", flush=True) - else: - try: - upload = upload_then_maybe_delete( - uploader, destination, sink_path, args.delete_after_upload) - except Exception as error: - record_upload_failure(error) - raise - with checkpoint_lock: - output_state["upload"] = upload - output_state["local_deleted"] = upload["local_deleted"] - atomic_json(checkpoint_path, checkpoint) - safe_print(f" uploaded: {upload['url']}", flush=True) - - def upload_task(): - return run_adaptive_upload(do_upload_task, destination.stat().st_size) - - pending_uploads.append(upload_executor.submit(upload_task)) - - with concurrent.futures.ThreadPoolExecutor(max_workers=args.format_workers) as executor: - futures = [executor.submit(process_format, fmt, job) for fmt, job in jobs] - for future in concurrent.futures.as_completed(futures): - future.result() + upload_item = { + "local_path": destination, + "destination_path": sink_path, + "format_name": fmt, + "ordinal": position, + "output_state": output_state, + "file_state": state, + "size_bytes": destination.stat().st_size, + } + nonlocal upload_reserved_bytes + with work_condition: + upload_queue.append(upload_item) + upload_reserved_bytes += upload_item["size_bytes"] + work_condition.notify_all() + + for fmt, job in jobs: + process_format(fmt, job) if not args.keep_downloads: raw.unlink(missing_ok=True) if "download" in state: @@ -1261,22 +1406,19 @@ def upload_task(): state.pop("error", None) save_checkpoint() - transcode_slots = max(1, args.transcode_workers // max(1, len(formats))) - conversion_executor = concurrent.futures.ThreadPoolExecutor(max_workers=transcode_slots) - conversion_futures = [] - - def submit_conversion(position, shard, raw, outputs, state): + def submit_conversion(position, shard, raw, outputs, state, reserved=False): source_name = shard["path"] state["status"] = "converting" state.pop("error", None) save_checkpoint() - conversion_futures.append(conversion_executor.submit( - process_downloaded_shard, position, shard, source_name, raw, outputs, state)) + with work_condition: + conversion_queue.append((position, shard, source_name, raw, outputs, state, reserved)) + work_condition.notify_all() # Admit resumable local sources immediately. Network downloads below are admitted in # completion order, so one slow low-ordinal shard cannot strand ready CPU work. positions = {shard["path"]: position for position, shard in enumerate(selection, 1)} - shards_by_name = {shard["path"]: shard for shard in selection} + missing_names = {shard["path"] for shard in missing_shards} for position, shard in enumerate(selection, 1): source_name = shard["path"] raw, outputs = shard_paths(shard) @@ -1289,8 +1431,9 @@ def submit_conversion(position, shard, raw, outputs, state): for fmt, destination in outputs.items(): output_state = state["outputs"].get(fmt, {}) upload_complete = output_state.get("upload", {}).get("status") == "complete" - local_complete = (destination.exists() and - destination.stat().st_size == output_state.get("metrics", {}).get("size_bytes")) + local_complete = destination.exists() and destination.stat().st_size == output_state.get("metrics", {}).get( + "size_bytes" + ) durable_complete = upload_complete if uploader else local_complete if output_state.get("status") == "complete" and durable_complete: completed_outputs.append(fmt) @@ -1302,7 +1445,7 @@ def submit_conversion(position, shard, raw, outputs, state): flush=True, ) continue - if source_name not in download_futures and not shard_needs_download(shard): + if source_name not in missing_names and not shard_needs_download(shard): submit_conversion(position, shard, raw, outputs, state) def status_snapshot(): @@ -1311,101 +1454,303 @@ def status_snapshot(): outputs = [output for state in files for output in state.get("outputs", {}).values()] source_done = sum(state.get("status") == "complete" for state in files) converted = sum(output.get("status") == "complete" for output in outputs) - uploaded = sum(output.get("upload", {}).get("status") == "complete" - for output in outputs) - preuploaded = sum(output.get("upload", {}).get("status") == "preuploaded" - for output in outputs) + uploaded = sum(output.get("upload", {}).get("status") == "complete" for output in outputs) + preuploaded = sum(output.get("upload", {}).get("status") == "preuploaded" for output in outputs) with transfer_totals_lock: totals = dict(transfer_totals) + with work_condition: + queue_counts = { + "download": len(download_queue), + "convert": len(conversion_queue), + "upload": len(upload_queue), + } + active_counts = dict(active_work) + reserved_download_bytes = download_reserved_bytes + reserved_upload_bytes = upload_reserved_bytes + current_download_batch = download_batch_size + current_upload_batch = upload_batch_size elapsed = max(time.monotonic() - pipeline_started, 0.001) conversion_source_bytes = sum( - state["size"] for state in files for output in state.get("outputs", {}).values() - if output.get("status") == "complete" and output.get("metrics", {}).get("seconds")) + state["size"] + for state in files + for output in state.get("outputs", {}).values() + if output.get("status") == "complete" and output.get("metrics", {}).get("seconds") + ) return { "queues": { - "download": {"waiting_items": len(missing_shards) - missing_index, - "active_items": len(download_futures), - "reserved_bytes": sum(download_sizes.values())}, - "transcode": {"waiting_items": sum(not future.running() and not future.done() - for future in conversion_futures), - "active_items": sum(future.running() for future in conversion_futures), - "succeeded_items": sum(future.done() and future.exception() is None - for future in conversion_futures), - "active_cpu_slots": args.transcode_workers - VX_CPU_SLOTS.qsize()}, - "upload": {"waiting_items": sum(not future.running() and not future.done() - for future in pending_uploads), - "active_items": upload_active, - "succeeded_items": sum(future.done() and future.exception() is None - for future in pending_uploads), - "concurrency": upload_concurrency}, + "download": { + "waiting_items": queue_counts["download"], + "active_items": active_counts["download"], + "reserved_bytes": reserved_download_bytes, + "batch_size": current_download_batch, + }, + "transcode": { + "waiting_items": queue_counts["convert"], + "active_items": active_counts["convert"], + "active_cpu_slots": len(cpu_slots) - VX_CPU_SLOTS.qsize(), + }, + "upload": { + "waiting_items": queue_counts["upload"], + "active_items": active_counts["upload"], + "reserved_bytes": reserved_upload_bytes, + "batch_size": current_upload_batch, + }, + }, + "progress": { + "source_complete": source_done, + "source_total": len(selection), + "outputs_converted": converted, + "outputs_uploaded": uploaded, + "outputs_preuploaded": preuploaded, + }, + "limits": { + "download_files": args.download_buffer_files, + "download_bytes": args.download_buffer_size, + "upload_files": args.upload_buffer_files, + "upload_bytes": args.upload_buffer_size, }, - "progress": {"source_complete": source_done, "source_total": len(selection), - "outputs_converted": converted, - "outputs_uploaded": uploaded, "outputs_preuploaded": preuploaded}, - "limits": {"download_files": args.download_buffer_files, - "download_bytes": args.download_buffer_size, - "upload_files": args.upload_buffer_files, - "upload_bytes": args.upload_buffer_size}, "throughput": { "download_effective_bytes_per_second": totals["download_bytes"] / elapsed, - "download_worker_bytes_per_second": totals["download_bytes"] - / max(totals["download_seconds"], 0.001), + "download_worker_bytes_per_second": totals["download_bytes"] / max(totals["download_seconds"], 0.001), "conversion_source_bytes_per_second": conversion_source_bytes / elapsed, "upload_effective_bytes_per_second": totals["upload_bytes"] / elapsed, - "upload_worker_bytes_per_second": totals["upload_bytes"] - / max(totals["upload_seconds"], 0.001), + "upload_worker_bytes_per_second": totals["upload_bytes"] / max(totals["upload_seconds"], 0.001), }, - "failures": {"download": totals["download_failures"], - "upload": totals["upload_failures"]}, + "failures": {"download": totals["download_failures"], "upload": totals["upload_failures"]}, } status = LiveStatus(args.output_dir / "status.json", args.status_interval, status_snapshot) status.start() - try: - for source_name, completed in completed_futures(download_futures): - download_sizes.pop(source_name) - shard = shards_by_name[source_name] - raw, outputs = shard_paths(shard) - state = checkpoint["files"][source_name] - try: - downloaded_path, download_seconds = completed.result() - if downloaded_path.resolve() != raw.resolve(): - raise RuntimeError(f"Hub download returned unexpected path: {downloaded_path}") - state["download"] = {"status": "complete", "path": str(raw), - "size_bytes": shard["size"], "seconds": download_seconds} - rate = shard["size"] / max(download_seconds, 0.001) - recent_download_rates.append(rate) + + def claim_batch(queue, maximum_files, maximum_bytes, size): + batch = [] + batch_bytes = 0 + while queue and len(batch) < maximum_files: + item = queue[0] + item_bytes = size(item) + if batch and batch_bytes + item_bytes > maximum_bytes: + break + batch.append(queue.popleft()) + batch_bytes += item_bytes + if batch_bytes >= maximum_bytes: + break + return batch, batch_bytes + + def run_upload_batch(batch): + payload = [ + {key: item[key] for key in ("local_path", "destination_path", "format_name", "ordinal")} for item in batch + ] + with checkpoint_lock: + for item in batch: + item["output_state"]["upload"] = { + "status": "uploading", + "hub_path": item["destination_path"], + } + atomic_json(checkpoint_path, checkpoint) + if isinstance(uploader, HuggingFaceBatchUploader): + results = uploader.upload_batch(payload) + else: + results = [ + upload_then_maybe_delete( + uploader, + item["local_path"], + item["destination_path"], + args.delete_after_upload, + ) + for item in batch + ] + for item, result in zip(batch, results, strict=True): + if result["status"] == "complete" and isinstance(uploader, HuggingFaceBatchUploader): + apply_committed_batch(result) + safe_print(f" committed batch: {result['url']}", flush=True) + elif result["status"] == "preuploaded": + with checkpoint_lock: + item["output_state"]["upload"] = result + atomic_json(checkpoint_path, checkpoint) + safe_print(f" preuploaded: {item['destination_path']}", flush=True) + else: + with checkpoint_lock: + item["output_state"]["upload"] = result + item["output_state"]["local_deleted"] = result.get("local_deleted", False) + atomic_json(checkpoint_path, checkpoint) + safe_print(f" uploaded: {result['url']}", flush=True) + + def worker_loop(): + nonlocal download_reserved_files, download_reserved_bytes, upload_reserved_bytes + nonlocal download_batch_size, upload_batch_size + while True: + with work_condition: + while True: + if worker_failure: + return + upload_high = upload_queue and ( + len(upload_queue) >= args.upload_buffer_files + or upload_reserved_bytes >= args.upload_buffer_size + ) + download_has_capacity = download_queue and ( + download_reserved_files == 0 + or ( + download_reserved_files < args.download_buffer_files + and download_reserved_bytes + download_queue[0]["size"] <= args.download_buffer_size + ) + ) + stage = choose_work_stage( + upload_waiting=bool(upload_queue), + upload_high=bool(upload_high), + download_available=bool(download_has_capacity), + convert_waiting=bool(conversion_queue), + active=any(active_work.values()), + ) + if stage == "done": + return + if stage is None: + work_condition.wait() + continue + + if stage == "download": + remaining_files = max(1, args.download_buffer_files - download_reserved_files) + remaining_bytes = max(1, args.download_buffer_size - download_reserved_bytes) + batch, batch_bytes = claim_batch( + download_queue, + min(download_batch_size, remaining_files), + remaining_bytes, + lambda shard: shard["size"], + ) + download_reserved_files += len(batch) + download_reserved_bytes += batch_bytes + elif stage == "upload": + batch, batch_bytes = claim_batch( + upload_queue, + upload_batch_size, + args.upload_buffer_size, + lambda item: item["size_bytes"], + ) + else: + batch = [conversion_queue.popleft()] + batch_bytes = batch[0][1]["size"] + active_work[stage] += 1 + break + + started = time.monotonic() + try: + if stage == "download": + results = download_shard_batch( + args.repo, + source_commit, + batch, + args.output_dir / "downloads", + args.hub_attempts, + args.hub_timeout, + ) + elapsed = max(time.monotonic() - started, 0.001) with transfer_totals_lock: - transfer_totals["download_bytes"] += shard["size"] - transfer_totals["download_seconds"] += download_seconds - submit_conversion(positions[source_name], shard, raw, outputs, state) - safe_print(f"[{positions[source_name]}/{len(selection)}] downloaded " - f"{source_name} in {download_seconds:.1f}s", flush=True) - except Exception as error: + transfer_totals["download_bytes"] += batch_bytes + transfer_totals["download_seconds"] += elapsed + with work_condition: + download_batch_size = adapt_concurrency( + download_batch_size, + args.download_max_concurrency, + recent_download_rates, + batch_bytes / elapsed, + ) + for shard, downloaded_path, download_seconds in results: + source_name = shard["path"] + raw, outputs = shard_paths(shard) + if downloaded_path.resolve() != raw.resolve(): + raise RuntimeError(f"Hub download returned unexpected path: {downloaded_path}") + state = checkpoint["files"][source_name] + state["download"] = { + "status": "complete", + "path": str(raw), + "size_bytes": shard["size"], + "seconds": download_seconds, + } + save_checkpoint() + submit_conversion(positions[source_name], shard, raw, outputs, state, reserved=True) + safe_print( + f"[{positions[source_name]}/{len(selection)}] downloaded {source_name} " + f"in {download_seconds:.1f}s", + flush=True, + ) + elif stage == "convert": + process_downloaded_shard(*batch[0][:-1]) + if batch[0][-1]: + with work_condition: + download_reserved_files -= 1 + download_reserved_bytes -= batch_bytes + else: + run_upload_batch(batch) + elapsed = max(time.monotonic() - started, 0.001) with transfer_totals_lock: - transfer_totals["download_failures"] += 1 - state["status"] = "failed" - state["error"] = f"download: {error}" - save_checkpoint() - raise - fill_download_window() - while len(pending_uploads) >= args.upload_buffer_files: - pending_uploads.pop(0).result() + transfer_totals["upload_bytes"] += batch_bytes + transfer_totals["upload_seconds"] += elapsed + with work_condition: + upload_batch_size = adapt_concurrency( + upload_batch_size, + args.upload_max_concurrency, + recent_upload_rates, + batch_bytes / elapsed, + ) + upload_reserved_bytes -= batch_bytes + except Exception as error: # noqa: BLE001 - worker boundary records every failure + if stage in ("download", "upload"): + with transfer_totals_lock: + transfer_totals[f"{stage}_failures"] += 1 + if stage == "download": + with checkpoint_lock: + for shard in batch: + state = checkpoint["files"][shard["path"]] + state["status"] = "failed" + state["error"] = f"download: {error}" + atomic_json(checkpoint_path, checkpoint) + elif stage == "convert": + with checkpoint_lock: + state = batch[0][5] + state["status"] = "failed" + state.setdefault("error", f"conversion: {error}") + atomic_json(checkpoint_path, checkpoint) + else: + with checkpoint_lock: + for item in batch: + item["output_state"]["upload"] = { + "status": "failed", + "hub_path": item["destination_path"], + "error": str(error), + } + item["file_state"]["status"] = "failed" + item["file_state"]["error"] = f"upload {item['format_name']}: {error}" + atomic_json(checkpoint_path, checkpoint) + with work_condition: + if stage == "download": + download_reserved_files -= len(batch) + download_reserved_bytes -= batch_bytes + elif stage == "convert" and batch[0][-1]: + download_reserved_files -= 1 + download_reserved_bytes -= batch_bytes + elif stage == "upload": + upload_reserved_bytes -= batch_bytes + worker_failure.append(error) + work_condition.notify_all() + finally: + with work_condition: + active_work[stage] -= 1 + work_condition.notify_all() + + try: + worker_count = args.workers + with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as executor: + workers = [executor.submit(worker_loop) for _ in range(worker_count)] + for worker in workers: + worker.result() + if worker_failure: + raise worker_failure[0] + if isinstance(uploader, HuggingFaceBatchUploader): + for final_batch in uploader.flush(): + apply_committed_batch(final_batch) + safe_print(f" committed final batch: {final_batch['url']}", flush=True) + write_reports(checkpoint, args.output_dir, selection) finally: - status.write() - for future in concurrent.futures.as_completed(conversion_futures): - future.result() - conversion_executor.shutdown(wait=True, cancel_futures=True) - download_executor.shutdown(wait=True, cancel_futures=True) - for future in pending_uploads: - future.result() - upload_executor.shutdown(wait=True, cancel_futures=True) - if isinstance(uploader, HuggingFaceBatchUploader): - for final_batch in uploader.flush(): - apply_committed_batch(final_batch) - safe_print(f" committed final batch: {final_batch['url']}", flush=True) - write_reports(checkpoint, args.output_dir, selection) - status.close() + status.close() return 0 @@ -1418,6 +1763,6 @@ def status_snapshot(): file=sys.stderr, ) sys.exit(130) - except Exception as error: + except Exception as error: # noqa: BLE001 safe_print(f"error: {error}", file=sys.stderr) sys.exit(1) diff --git a/scripts/pyproject.toml b/scripts/pyproject.toml index 3479da2..1b1c911 100644 --- a/scripts/pyproject.toml +++ b/scripts/pyproject.toml @@ -12,3 +12,6 @@ dependencies = [ [tool.uv] package = false + +[dependency-groups] +dev = ["ruff>=0.12.0"] diff --git a/scripts/tests/test_hf_sync.py b/scripts/tests/test_hf_sync.py index 7c4cb60..4962398 100644 --- a/scripts/tests/test_hf_sync.py +++ b/scripts/tests/test_hf_sync.py @@ -5,12 +5,13 @@ import importlib.util import json import sys +import tempfile import time import types -import tempfile import unittest +from collections import deque from pathlib import Path - +from unittest import mock SCRIPT = Path(__file__).parents[1] / "hf-sync.py" SPEC = importlib.util.spec_from_file_location("hf_sync", SCRIPT) @@ -19,20 +20,40 @@ class LocalCopyUploaderTest(unittest.TestCase): - def test_completed_downloads_do_not_wait_for_slow_first_item(self): - def finish(name, delay): - time.sleep(delay) - return name - - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: - futures = { - "slow-first": executor.submit(finish, "slow-first", 0.15), - "fast-second": executor.submit(finish, "fast-second", 0.01), - } - completion_order = [name for name, _ in MODULE.completed_futures(futures)] + def test_unified_worker_priority(self): + choose = MODULE.choose_work_stage + common = {"upload_waiting": True, "download_available": True, "convert_waiting": True, "active": True} + + self.assertEqual(choose(upload_high=True, **common), "upload") + self.assertEqual(choose(upload_high=False, **common), "download") + self.assertEqual( + choose(upload_waiting=True, upload_high=False, download_available=False, convert_waiting=True, active=True), + "upload", + ) + self.assertEqual( + choose( + upload_waiting=False, upload_high=False, download_available=False, convert_waiting=True, active=True + ), + "convert", + ) + self.assertEqual( + choose( + upload_waiting=False, upload_high=False, download_available=False, convert_waiting=False, active=False + ), + "done", + ) + + def test_adaptive_concurrency_increases_holds_and_backs_off(self): + rates = deque(maxlen=4) - self.assertEqual(completion_order, ["fast-second", "slow-first"]) - self.assertEqual(futures, {}) + concurrency = MODULE.adapt_concurrency(2, 4, rates, 100) + self.assertEqual(concurrency, 3) + concurrency = MODULE.adapt_concurrency(concurrency, 4, rates, 90) + self.assertEqual(concurrency, 3) + MODULE.adapt_concurrency(concurrency, 4, rates, 100) + MODULE.adapt_concurrency(concurrency, 4, rates, 100) + concurrency = MODULE.adapt_concurrency(concurrency, 4, rates, 50) + self.assertEqual(concurrency, 1) def test_live_status_is_atomic_and_marks_final_snapshot(self): with tempfile.TemporaryDirectory() as directory: @@ -59,70 +80,120 @@ def dataset_info(self, repo, revision=None, timeout=None): def list_repo_tree(self, repo, path_in_repo=None, **kwargs): if repo == "external/source": - return [types.SimpleNamespace( - path="sample/a.parquet", size=10, - lfs={"sha256": "source-sha"})] + return [types.SimpleNamespace(path="sample/a.parquet", size=10, lfs={"sha256": "source-sha"})] if path_in_repo == "vortex": - return [types.SimpleNamespace( - path="vortex/sample/a.vortex", size=7, - blob_id="destination-oid", lfs={"sha256": "destination-sha"})] + return [ + types.SimpleNamespace( + path="vortex/sample/a.vortex", + size=7, + blob_id="destination-oid", + lfs={"sha256": "destination-sha"}, + ) + ] return [] args = types.SimpleNamespace( - repo="external/source", revision="main", prefix="sample", include="*.parquet", - filter=[], formats="vortex,vortex-compact", upload_repo="vortex-data/source", - upload_revision="main", upload_prefix="", hub_attempts=1, hub_timeout=1, - upload_batch_files=100, mode="all", limit=10, target_size=100, seed=0, + repo="external/source", + revision="main", + prefix="sample", + include="*.parquet", + filter=[], + formats="vortex,vortex-compact", + upload_repo="vortex-data/source", + upload_revision="main", + upload_prefix="", + hub_attempts=1, + hub_timeout=1, + upload_batch_files=100, + mode="all", + limit=10, + target_size=100, + seed=0, ) plan = MODULE.create_action_plan(FakeApi(), args) self.assertEqual(plan["source"]["revision"], "immutable-source") self.assertEqual(plan["destination"]["revision"], "immutable-destination") - self.assertEqual(plan["summary"], { - "source_files": 1, "source_bytes": 10, - "create_actions": 1, "skip_actions": 1, - }) + self.assertEqual( + plan["summary"], + { + "source_files": 1, + "source_bytes": 10, + "create_actions": 1, + "skip_actions": 1, + }, + ) actions = plan["work_chunks"][0]["actions"] self.assertEqual(actions[0]["action"], "skip") self.assertEqual(actions[0]["existing"]["sha256"], "destination-sha") self.assertEqual(actions[1]["action"], "create") self.assertEqual(actions[1]["upload_batch"], "vortex-compact-1-1") - self.assertEqual(plan["upload_batches"], [{ - "id": "vortex-compact-1-1", "format": "vortex-compact", - "start": 1, "end": 1, "total_files": 1, - "commit_message": "Upload vortex-compact files 1-1 of 1", - "source_ordinals": [1], - "destination_paths": ["vortex-compact/sample/a.vortex"], - "planned_source_bytes": 10, - }]) + self.assertEqual( + plan["upload_batches"], + [ + { + "id": "vortex-compact-1-1", + "format": "vortex-compact", + "start": 1, + "end": 1, + "total_files": 1, + "commit_message": "Upload vortex-compact files 1-1 of 1", + "source_ordinals": [1], + "destination_paths": ["vortex-compact/sample/a.vortex"], + "planned_source_bytes": 10, + } + ], + ) def test_plan_round_trip_supplies_apply_work_chunks(self): plan = { "version": MODULE.PLAN_VERSION, "kind": "hf-sync-plan", - "source": {"repo": "external/source", "requested_revision": "main", - "revision": "source-sha", "prefix": "sample", "include": "*.parquet", - "filters": []}, - "destination": {"repo": "vortex-data/source", "requested_revision": "main", - "revision": "destination-sha", "prefix": "converted"}, + "source": { + "repo": "external/source", + "requested_revision": "main", + "revision": "source-sha", + "prefix": "sample", + "include": "*.parquet", + "filters": [], + }, + "destination": { + "repo": "vortex-data/source", + "requested_revision": "main", + "revision": "destination-sha", + "prefix": "converted", + }, "selection": {"mode": "first", "limit": 1, "target_bytes": 10, "seed": 0}, "formats": ["vortex"], - "work_chunks": [{"ordinal": 1, "source": {"path": "sample/a.parquet", - "size": 10, "sha256": "abc"}, - "actions": [{"action": "create", "format": "vortex", - "upload_batch": "vortex-1-1", - "commit_message": "Upload vortex files 1-1 of 1", - "destination_path": - "converted/vortex/sample/a.vortex"}]}], - "upload_batches": [{"id": "vortex-1-1", "format": "vortex", - "start": 1, "end": 1, "total_files": 1, - "commit_message": "Upload vortex files 1-1 of 1", - "source_ordinals": [1], - "destination_paths": [ - "converted/vortex/sample/a.vortex"]}], - "summary": {"source_files": 1, "source_bytes": 10, - "create_actions": 1, "skip_actions": 0}, + "work_chunks": [ + { + "ordinal": 1, + "source": {"path": "sample/a.parquet", "size": 10, "sha256": "abc"}, + "actions": [ + { + "action": "create", + "format": "vortex", + "upload_batch": "vortex-1-1", + "commit_message": "Upload vortex files 1-1 of 1", + "destination_path": "converted/vortex/sample/a.vortex", + } + ], + } + ], + "upload_batches": [ + { + "id": "vortex-1-1", + "format": "vortex", + "start": 1, + "end": 1, + "total_files": 1, + "commit_message": "Upload vortex files 1-1 of 1", + "source_ordinals": [1], + "destination_paths": ["converted/vortex/sample/a.vortex"], + } + ], + "summary": {"source_files": 1, "source_bytes": 10, "create_actions": 1, "skip_actions": 0}, } with tempfile.TemporaryDirectory() as temporary: path = Path(temporary) / "plan.json" @@ -132,8 +203,7 @@ def test_plan_round_trip_supplies_apply_work_chunks(self): selection = MODULE.apply_plan_arguments(args, loaded) - self.assertEqual(selection, [{"path": "sample/a.parquet", "size": 10, - "sha256": "abc"}]) + self.assertEqual(selection, [{"path": "sample/a.parquet", "size": 10, "sha256": "abc"}]) self.assertEqual(args.repo, "external/source") self.assertEqual(args.revision, "source-sha") self.assertEqual(args.upload_repo, "vortex-data/source") @@ -216,55 +286,66 @@ def test_missing_destination_prefix_is_an_empty_listing(self): class FakeApi: def list_repo_tree(self, *args, **kwargs): response = httpx.Response( - 404, request=httpx.Request("GET", "https://huggingface.co/api/datasets/x/tree")) + 404, request=httpx.Request("GET", "https://huggingface.co/api/datasets/x/tree") + ) raise RemoteEntryNotFoundError("missing prefix", response=response) - self.assertEqual( - MODULE.list_repository_files(FakeApi(), "owner/repo", "main", "vortex"), {}) + self.assertEqual(MODULE.list_repository_files(FakeApi(), "owner/repo", "main", "vortex"), {}) def test_huggingface_batches_each_format_by_file_range(self): class FakeApi: def __init__(self): self.preuploads = [] + self.preupload_calls = 0 self.commits = [] def dataset_info(self, repo, revision=None, timeout=None): return types.SimpleNamespace(sha="parent") def preupload_lfs_files(self, repo, additions, **kwargs): + self.preupload_calls += 1 self.preuploads.extend(additions) def create_commit(self, repo, operations, **kwargs): self.commits.append((list(operations), kwargs)) number = len(self.commits) - return types.SimpleNamespace(commit_url=f"https://fixture/commit/{number}", - oid=f"commit-{number}") + return types.SimpleNamespace(commit_url=f"https://fixture/commit/{number}", oid=f"commit-{number}") with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) api = FakeApi() - uploader = MODULE.HuggingFaceBatchUploader( - api, "owner/repo", "main", 2, {"vortex": 3, "vortex-compact": 3}) + uploader = MODULE.HuggingFaceBatchUploader(api, "owner/repo", "main", 2, {"vortex": 3, "vortex-compact": 3}) paths = [] for ordinal in range(1, 4): path = root / f"{ordinal}.vortex" path.write_bytes(bytes([ordinal])) paths.append(path) - pending = uploader.upload(paths[0], "vortex/1.vortex", - format_name="vortex", ordinal=1) - committed = uploader.upload(paths[1], "vortex/2.vortex", - format_name="vortex", ordinal=2) - compact = uploader.upload(paths[2], "vortex-compact/3.vortex", - format_name="vortex-compact", ordinal=3) + pending, committed = uploader.upload_batch( + [ + { + "local_path": paths[0], + "destination_path": "vortex/1.vortex", + "format_name": "vortex", + "ordinal": 1, + }, + { + "local_path": paths[1], + "destination_path": "vortex/2.vortex", + "format_name": "vortex", + "ordinal": 2, + }, + ] + ) + compact = uploader.upload(paths[2], "vortex-compact/3.vortex", format_name="vortex-compact", ordinal=3) final = uploader.flush() self.assertEqual(pending["status"], "preuploaded") self.assertEqual(committed["commit_message"], "Upload vortex files 1-2 of 3") self.assertEqual(compact["status"], "preuploaded") - self.assertEqual(final[0]["commit_message"], - "Upload vortex-compact files 3-3 of 3") + self.assertEqual(final[0]["commit_message"], "Upload vortex-compact files 3-3 of 3") self.assertEqual(len(api.preuploads), 3) + self.assertEqual(api.preupload_calls, 2) self.assertEqual(len(api.commits), 2) self.assertEqual(api.commits[0][1]["parent_commit"], "parent") self.assertEqual(api.commits[1][1]["parent_commit"], "commit-1") @@ -282,8 +363,7 @@ def preupload_lfs_files(self, repo, additions, **kwargs): def create_commit(self, repo, operations, **kwargs): self.commits.append(list(operations)) - return types.SimpleNamespace(commit_url="https://fixture/commit/1", - oid="commit-1") + return types.SimpleNamespace(commit_url="https://fixture/commit/1", oid="commit-1") with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) @@ -292,14 +372,12 @@ def create_commit(self, repo, operations, **kwargs): first.write_bytes(b"123") second.write_bytes(b"456") api = FakeApi() - uploader = MODULE.HuggingFaceBatchUploader( - api, "owner/repo", "main", 100, {"vortex": 2}, batch_bytes=5) + uploader = MODULE.HuggingFaceBatchUploader(api, "owner/repo", "main", 100, {"vortex": 2}, batch_bytes=5) - self.assertEqual(uploader.upload( - first, "vortex/first.vortex", format_name="vortex", ordinal=1 - )["status"], "preuploaded") - result = uploader.upload( - second, "vortex/second.vortex", format_name="vortex", ordinal=2) + self.assertEqual( + uploader.upload(first, "vortex/first.vortex", format_name="vortex", ordinal=1)["status"], "preuploaded" + ) + result = uploader.upload(second, "vortex/second.vortex", format_name="vortex", ordinal=2) self.assertEqual(result["status"], "complete") self.assertEqual(result["size_bytes"], 6) @@ -318,13 +396,20 @@ def preupload_lfs_files(self, repo, additions, **kwargs): def create_commit(self, repo, operations, **kwargs): self.commits.append((list(operations), kwargs)) - return types.SimpleNamespace(commit_url="https://fixture/commit/1", - oid="commit-1") - - planned = [{"id": "vortex-1-2", "format": "vortex", "start": 1, "end": 2, - "total_files": 2, "commit_message": "Planned vortex target batch", - "source_ordinals": [1, 2], - "destination_paths": ["vortex/a.vortex", "vortex/b.vortex"]}] + return types.SimpleNamespace(commit_url="https://fixture/commit/1", oid="commit-1") + + planned = [ + { + "id": "vortex-1-2", + "format": "vortex", + "start": 1, + "end": 2, + "total_files": 2, + "commit_message": "Planned vortex target batch", + "source_ordinals": [1, 2], + "destination_paths": ["vortex/a.vortex", "vortex/b.vortex"], + } + ] with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) first = root / "first.vortex" @@ -333,12 +418,11 @@ def create_commit(self, repo, operations, **kwargs): second.write_bytes(b"second") api = FakeApi() uploader = MODULE.HuggingFaceBatchUploader( - api, "owner/repo", "main", 100, {"vortex": 2}, planned_batches=planned) + api, "owner/repo", "main", 100, {"vortex": 2}, planned_batches=planned + ) - pending = uploader.upload(first, "vortex/a.vortex", - format_name="vortex", ordinal=1) - committed = uploader.upload(second, "vortex/b.vortex", - format_name="vortex", ordinal=2) + pending = uploader.upload(first, "vortex/a.vortex", format_name="vortex", ordinal=1) + committed = uploader.upload(second, "vortex/b.vortex", format_name="vortex", ordinal=2) self.assertEqual(pending["status"], "preuploaded") self.assertEqual(committed["commit_message"], "Planned vortex target batch") @@ -350,23 +434,47 @@ def test_download_buffer_allows_one_oversized_shard_only_when_empty(self): self.assertTrue(MODULE.fits_download_buffer(4, 6, 10)) self.assertFalse(MODULE.fits_download_buffer(4, 7, 10)) + def test_download_batch_uses_one_concurrent_snapshot_request(self): + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + shards = [ + {"path": "data/a.parquet", "size": 1, "sha256": None}, + {"path": "data/b.parquet", "size": 2, "sha256": None}, + ] + for shard in shards: + path = root / shard["path"] + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"x" * shard["size"]) + + with mock.patch("huggingface_hub.snapshot_download", return_value=str(root)) as download: + results = MODULE.download_shard_batch( + "owner/source", "commit", shards, root, attempts=1, etag_timeout=30 + ) + + self.assertEqual([result[0] for result in results], shards) + self.assertEqual(download.call_count, 1) + self.assertEqual(download.call_args.kwargs["allow_patterns"], ["data/a.parquet", "data/b.parquet"]) + self.assertEqual(download.call_args.kwargs["max_workers"], 2) + def test_resume_reuses_local_outputs_and_skips_committed_outputs(self): with tempfile.TemporaryDirectory() as temporary: root = Path(temporary) local = root / "pending.vortex" local.write_bytes(b"encoded") outputs = {"vortex": local, "vortex-compact": root / "committed.vortex"} - state = {"outputs": { - "vortex": { - "status": "complete", - "metrics": {"size_bytes": len(b"encoded")}, - "upload": {"status": "failed"}, - }, - "vortex-compact": { - "status": "complete", - "upload": {"status": "complete"}, - }, - }} + state = { + "outputs": { + "vortex": { + "status": "complete", + "metrics": {"size_bytes": len(b"encoded")}, + "upload": {"status": "failed"}, + }, + "vortex-compact": { + "status": "complete", + "upload": {"status": "complete"}, + }, + } + } self.assertFalse(MODULE.needs_source_download(state, outputs)) local.unlink() @@ -399,18 +507,15 @@ def list_repo_tree(self, *args, **kwargs): with tempfile.TemporaryDirectory() as temporary: artifact = Path(temporary) / "artifact.vortex" artifact.write_bytes(b"keep until commit") - uploader = MODULE.HuggingFaceBatchUploader( - FakeApi(), "owner/repo", "main", 1, {"vortex": 1}, attempts=1) + uploader = MODULE.HuggingFaceBatchUploader(FakeApi(), "owner/repo", "main", 1, {"vortex": 1}, attempts=1) with self.assertRaisesRegex(ValueError, "fixture commit failure"): - uploader.upload(artifact, "vortex/artifact.vortex", - format_name="vortex", ordinal=1) + uploader.upload(artifact, "vortex/artifact.vortex", format_name="vortex", ordinal=1) self.assertEqual(artifact.read_bytes(), b"keep until commit") def test_mirrored_destination_path(self): self.assertEqual( - MODULE.destination_path( - "sample/10BT/000_00000.parquet", "vortex-compact"), + MODULE.destination_path("sample/10BT/000_00000.parquet", "vortex-compact"), "vortex-compact/sample/10BT/000_00000.vortex", ) @@ -423,13 +528,13 @@ def test_parquet_destination_keeps_parquet_extension(self): def test_full_path_filter(self): class FakeApi: def list_repo_tree(self, *args, **kwargs): - return [types.SimpleNamespace(path="sample/10BT/a.parquet", size=10, lfs=None), - types.SimpleNamespace(path="sample/100BT/b.parquet", size=20, lfs=None)] + return [ + types.SimpleNamespace(path="sample/10BT/a.parquet", size=10, lfs=None), + types.SimpleNamespace(path="sample/100BT/b.parquet", size=20, lfs=None), + ] - shards = MODULE.list_shards(FakeApi(), "owner/repo", "main", "sample", "*.parquet", - ["sample/10BT/*"]) - self.assertEqual(shards, [{"path": "sample/10BT/a.parquet", "size": 10, - "sha256": None}]) + shards = MODULE.list_shards(FakeApi(), "owner/repo", "main", "sample", "*.parquet", ["sample/10BT/*"]) + self.assertEqual(shards, [{"path": "sample/10BT/a.parquet", "size": 10, "sha256": None}]) def test_copies_to_hub_shaped_fixture(self): with tempfile.TemporaryDirectory() as temporary: @@ -488,9 +593,11 @@ def test_parallel_real_encoders_upload_and_delete(self): source = root / "source.parquet" pq.write_table(pa.table({"text": ["alpha", "beta"] * 70_000}), source) local = root / "local" - outputs = {"parquet-zstd6": local / "source.parquet", - "vortex": local / "source.vortex", - "vortex-compact": local / "source-compact.vortex"} + outputs = { + "parquet-zstd6": local / "source.parquet", + "vortex": local / "source.vortex", + "vortex-compact": local / "source-compact.vortex", + } sink = MODULE.LocalCopyUploader(root / "sink") def encode_upload(fmt): @@ -500,8 +607,7 @@ def encode_upload(fmt): else: strategy = "compact" if fmt == "vortex-compact" else "btrblocks" MODULE.run_vx(vx, source, destination, strategy) - return MODULE.upload_then_maybe_delete( - sink, destination, f"fixture/{fmt}/{destination.name}", True) + return MODULE.upload_then_maybe_delete(sink, destination, f"fixture/{fmt}/{destination.name}", True) with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: results = list(executor.map(encode_upload, outputs)) @@ -516,9 +622,121 @@ def encode_upload(fmt): self.assertGreater(uploaded.stat().st_size, 0) if fmt == "parquet-zstd6": metadata = pq.ParquetFile(uploaded).metadata - self.assertTrue(all(metadata.row_group(index).num_rows - <= MODULE.PARQUET_BATCH_ROWS - for index in range(metadata.num_row_groups))) + self.assertTrue( + all( + metadata.row_group(index).num_rows <= MODULE.PARQUET_BATCH_ROWS + for index in range(metadata.num_row_groups) + ) + ) + + def test_unified_pool_runs_download_convert_upload_pipeline(self): + try: + import pyarrow as pa + import pyarrow.parquet as pq + except ImportError: + self.skipTest("pyarrow is not installed") + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + output_dir = root / "run" + sink = root / "sink" + shard = {"path": "data/a.parquet", "size": 0, "sha256": None} + plan = { + "source": { + "repo": "owner/source", + "requested_revision": "main", + "revision": "source-commit", + "prefix": "data", + "include": "*.parquet", + "filters": [], + }, + "destination": { + "repo": "owner/destination", + "requested_revision": "main", + "revision": "destination-commit", + "prefix": "", + }, + "selection": {"mode": "all", "limit": 1, "target_bytes": 1, "seed": 0}, + "formats": ["parquet-zstd6"], + "work_chunks": [ + { + "ordinal": 1, + "source": shard, + "actions": [ + { + "action": "create", + "format": "parquet-zstd6", + "destination_path": "parquet-zstd6/data/a.parquet", + "upload_batch": "parquet-zstd6-1-1", + } + ], + } + ], + "upload_batches": [ + { + "id": "parquet-zstd6-1-1", + "format": "parquet-zstd6", + "start": 1, + "end": 1, + "total_files": 1, + "commit_message": "fixture", + "source_ordinals": [1], + "destination_paths": ["parquet-zstd6/data/a.parquet"], + } + ], + } + source_fixture = root / "source.parquet" + pq.write_table(pa.table({"value": [1, 2, 3]}), source_fixture) + shard["size"] = source_fixture.stat().st_size + + args = types.SimpleNamespace( + command="apply", + plan_file=root / "plan.json", + output_dir=output_dir, + vx=Path("/bin/true"), + workers=2, + format_workers=1, + transcode_workers=None, + shard_workers=1, + download_initial_concurrency=1, + download_max_concurrency=2, + download_buffer_files=2, + download_buffer_size=10_000_000, + upload_workers=1, + upload_max_concurrency=2, + upload_local_dir=sink, + upload_buffer_files=2, + upload_buffer_size=10_000_000, + upload_batch_files=2, + hub_attempts=1, + hub_timeout=1, + xet_range_gets=1, + xet_cache=root / "xet", + xet_high_performance=False, + keep_downloads=False, + delete_after_upload=False, + status_interval=0.01, + detached=True, + ) + + def fake_download(_repo, _revision, shards, download_root, _attempts, _timeout): + destination = download_root / shards[0]["path"] + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(source_fixture.read_bytes()) + return [(shards[0], destination, 0.01)] + + with ( + mock.patch.object(MODULE, "parse_args", return_value=args), + mock.patch.object(MODULE, "load_action_plan", return_value=plan), + mock.patch.object(MODULE, "require_xet_repository"), + mock.patch.object(MODULE, "download_shard_batch", side_effect=fake_download), + ): + result = MODULE.main() + + self.assertEqual(result, 0) + uploaded = sink / "parquet-zstd6/data/a.parquet" + self.assertTrue(uploaded.is_file()) + self.assertTrue(json.loads((output_dir / "status.json").read_text())["final"]) if __name__ == "__main__": diff --git a/uv.lock b/uv.lock index 8312fda..79c9800 100644 --- a/uv.lock +++ b/uv.lock @@ -262,6 +262,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] +[[package]] +name = "ruff" +version = "0.16.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/85/c8e12473c93018f92d19dd988a294202e1c27426c47ec4de53ffb847b8d8/ruff-0.16.5.tar.gz", hash = "sha256:1b88500f9ffbcab3dedb0082c9f9492e91ec3d618aac1236a3e0189938f7040b", size = 4912003, upload-time = "2026-08-27T16:34:18.258Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/b6/77c90a970fe2dae17a723acbd011043ea97c98d7deacccefdc4ba74ec512/ruff-0.16.5-py3-none-linux_armv6l.whl", hash = "sha256:12e5f673e774c35fbb62f288809c7653b73445f8ecec6b6063fd6ea3521aa14b", size = 10011941, upload-time = "2026-08-27T16:33:41.287Z" }, + { url = "https://files.pythonhosted.org/packages/4b/46/6cf67cf6411885a1d6f7f6d801682f155536a85176d10b605e2ceffed8bd/ruff-0.16.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:eda58a5802de40e7ed5b32b64e0b32539338cc6fcd2c78f61e3ad6a0d79f51c3", size = 10204049, upload-time = "2026-08-27T16:33:44.056Z" }, + { url = "https://files.pythonhosted.org/packages/46/fd/c8720ca7a090abf0c2fef4abe8a5ef6e5127ed15196d8886ff75a2b370e2/ruff-0.16.5-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c5ae9a7b9a8875131f40f8fe967cc86abf899779efd663cb7ce3d572d01da7eb", size = 9809037, upload-time = "2026-08-27T16:33:46.257Z" }, + { url = "https://files.pythonhosted.org/packages/43/45/a684caacdedaca180f52bacccc40bf0789d2c5a7c75f25324853e9eaedb5/ruff-0.16.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b719b0a1f4d59710d283ab2965f621684a108a9e41da622e3b23f0326cd0025", size = 9964129, upload-time = "2026-08-27T16:33:48.352Z" }, + { url = "https://files.pythonhosted.org/packages/9e/f2/5d2bcdaca6b5b93d1b4dfc166cd2aebf7680143a1b38a28759df13a94d31/ruff-0.16.5-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2298f2780ed1be0c5cb1361e32ab7b1467f3cce7dabe101d2210a314f2fe42e9", size = 9821518, upload-time = "2026-08-27T16:33:50.57Z" }, + { url = "https://files.pythonhosted.org/packages/aa/ff/011cce29accf9257d5974145b733fc653a37985ed6825413a3987cefbfe0/ruff-0.16.5-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:258f29035a2dd021e7861e631b227a5b3f14e50c1184c9a6a122c5f4576154d7", size = 10534835, upload-time = "2026-08-27T16:33:52.522Z" }, + { url = "https://files.pythonhosted.org/packages/d7/5a/f0cf109bada9bba0e96c90c21c9f9251803f57225c32d293327a03c710d6/ruff-0.16.5-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b9a4f0432966834019c74d1b7e5c51224305d7713f3d7faf3e7451f1a3be3cde", size = 11252550, upload-time = "2026-08-27T16:33:54.521Z" }, + { url = "https://files.pythonhosted.org/packages/63/4d/1d481aaea2046c6a7ed7c291f9004c669cce3c087b6b376ed5b08271e3fe/ruff-0.16.5-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b5eb3a8c3d0ade9cea42b591fd530368e8798380e30e0a308b85a5cf718f09ea", size = 10777949, upload-time = "2026-08-27T16:33:56.88Z" }, + { url = "https://files.pythonhosted.org/packages/ee/34/ee245ca55f64443233034b3d02b03236b19242004281247c079390b7facd/ruff-0.16.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ef0f69e191a13a3c9816f63163c88790cb12cd157bbbb384e9c44745702ab105", size = 10311656, upload-time = "2026-08-27T16:33:59.12Z" }, + { url = "https://files.pythonhosted.org/packages/a7/4d/c33a333e341c0a2b96c715b52d89a606f5a34cd4ac493cd9b8d0187186b8/ruff-0.16.5-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:0eeab41fbea2c42f98dfb9822cdccda9d24ba38d49f6dc945b5c236d48f0ef29", size = 10532125, upload-time = "2026-08-27T16:34:01.166Z" }, + { url = "https://files.pythonhosted.org/packages/30/e1/a64cef78b40192497bb98a27a8aa8f2c98ee9ee15bc97f7712d94ef32937/ruff-0.16.5-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:f0768e9df4300713fff30733c87575f68b6f1d8de41184e505b7fdd9c0c95eaf", size = 10097648, upload-time = "2026-08-27T16:34:03.16Z" }, + { url = "https://files.pythonhosted.org/packages/cc/4e/4cdc9ed3c3e109d2f71e62572a37457298d7bc7501ec3138babb7ed32bbd/ruff-0.16.5-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:95cc70cdc7aa80c338de356279d2adbeb2de0f520b9ecd8aba75b94e95e02f91", size = 9829344, upload-time = "2026-08-27T16:34:05.134Z" }, + { url = "https://files.pythonhosted.org/packages/39/4a/31ed35ce31729955fc583ee0d176d6e784c1290cb0b0a75cb2134c1ab72a/ruff-0.16.5-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d185c8398ded1bfd91c0c2cb258346307571eccc473a8490af8c3977399c384a", size = 10277117, upload-time = "2026-08-27T16:34:07.425Z" }, + { url = "https://files.pythonhosted.org/packages/a8/a0/60356d86687b4b666d593df213f4dc3041750d024cb7bf2cfa81cfd65c2e/ruff-0.16.5-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:fb8e3a3c4c6a784150a7ced53b015f4b253fc2bf97a610886419ead64b4756ef", size = 10711653, upload-time = "2026-08-27T16:34:09.712Z" }, + { url = "https://files.pythonhosted.org/packages/ed/20/656d67f5b25ca9bda4e02b1de25867b2954e1d19e03648060f167ad0f4cc/ruff-0.16.5-py3-none-win32.whl", hash = "sha256:288b0a5f080492fe5635db849f9e2e84aa3cce7b7f0e955997d416c507c76a26", size = 10034250, upload-time = "2026-08-27T16:34:11.8Z" }, + { url = "https://files.pythonhosted.org/packages/5b/42/ee8e68a207b9127fcde6c3d7e197def432f346cb1af159e1fa14ca0d1cdc/ruff-0.16.5-py3-none-win_amd64.whl", hash = "sha256:ddc6385fb2137f616357ca03d6c74f4be987f80fed4008566b754f6032b8546f", size = 10516714, upload-time = "2026-08-27T16:34:13.963Z" }, + { url = "https://files.pythonhosted.org/packages/73/e3/7df5a396e445b9ba49ce9a9437439a4d80042c61c0ade199abf8d16de1ac/ruff-0.16.5-py3-none-win_arm64.whl", hash = "sha256:a64abe90968719b851bb7cedffaa8753fbdbdadab483089682db623f3edc587e", size = 10391564, upload-time = "2026-08-27T16:34:16.064Z" }, +] + [[package]] name = "tqdm" version = "4.70.0" @@ -294,6 +319,11 @@ dependencies = [ { name = "pyarrow" }, ] +[package.dev-dependencies] +dev = [ + { name = "ruff" }, +] + [package.metadata] requires-dist = [ { name = "hf-xet", specifier = ">=1.1.10" }, @@ -301,3 +331,6 @@ requires-dist = [ { name = "huggingface-hub", specifier = ">=0.34.0" }, { name = "pyarrow", specifier = ">=17.0.0" }, ] + +[package.metadata.requires-dev] +dev = [{ name = "ruff", specifier = ">=0.12.0" }]