feat(migrate): wandb → Pluto historical data migration

[asaiacai]

Summary

Two-phase migration tool for importing legacy wandb data into Pluto with full timestamp fidelity — built for the first customer coming off the wandb shim who wants their historical runs transferred.

• pluto migrate wandb export — stages complete runs from the wandb cloud API on disk (parquet long-schema parts + run.json manifest + media/artifact files): full scan_history metrics, media, tables/histograms, events-stream system metrics, console logs, artifacts (size-cappable). Atomic per-run staging with sentinel-based resume.
• pluto migrate wandb load — replays staged runs through the public client API with original wall-clock timestamps, wandb::{entity}/{project}/{run_id} external-id dedup, crash-healing resume (typed RunExistsError → resume + re-replay), batched metric enqueue, bounded backpressure, and --dry-run.
• pluto migrate wandb all — both phases; loads whatever staged even if some exports failed.

Client groundwork

• op.log(..., timestamp=) — epoch-seconds override threaded through the sync store (wire format already carried per-point time); works in sync, legacy, and perf-queue modes.
• Op._log_console / Op._log_metrics_batch — batched backfill helpers.
• settings.disable_system_metrics — keeps the migration host's hardware/health metrics (sys/*, sys/pluto.*, run systemMetadata) out of imported runs, including in the sync subprocess.
• New migrate extra (pip install 'pluto-ml[migrate]'): wandb + pyarrow, lazily imported so the base CLI/package are unaffected.

Server dependency

Run createdAt/updatedAt backfill needs the companion server PR (Trainy-ai/server-private branch feat/run-createdat-backfill); metric/file/console point timestamps already round-trip against today's prod ingest. Until it deploys, imported runs show import-day creation dates only.

Test plan

• 71 new/updated unit tests (TDD): timestamp threading, monitor suppression, parquet round-trip/rotation, exporter (fake wandb API fixtures), loader (mocked init), CLI wiring.
• tests/test_migrate_staging_e2e.py: live round-trip against the dev channel, gated on PLUTO_STAGING_API_KEY.
• Multi-angle code review with 10 confirmed findings — all fixed in the final commit.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches core logging/sync paths and performs bulk server writes with idempotency edge cases (re-replay may duplicate media); scope is large but guarded by extensive tests and optional install.

Overview
Introduces pluto migrate wandb export|load|all for importing historical Weights & Biases runs into Pluto. Export pulls runs from the wandb API into a local staging layout (run.json, rotated parquet parts, downloaded files) with per-run atomic writes and sentinel-based resume. Load replays staged data through pluto.init using stable external ids (wandb::{entity}/{project}/{run_id}), original per-point timestamps, loaded_runs.json skip/resume, and sync-queue backpressure. The migrate optional extra (wandb, pyarrow) is lazy-loaded so the base CLI stays light.

The public client gains backfill-oriented behavior: op.log(..., timestamp=) for wall-clock replay; Op._log_metrics_batch and _log_console; typed RunExistsError for crash-healing resume; settings.disable_system_metrics so migration hosts do not pollute imported runs (monitor, run-create metadata, sync health uploads); and enqueue_metrics_batch on the sync path for faster replay.

Packaging adds the migrate extra in pyproject.toml / lockfile; coverage includes unit tests across schema, exporter, loader, CLI, and an optional staging e2e gated on PLUTO_STAGING_API_KEY. Run createdAt/updatedAt backfill still depends on a companion server deploy; metric/console timestamps are expected to round-trip today.

^{Reviewed by Cursor Bugbot for commit 641460f. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces a robust two-phase migration pipeline (pluto.migrate) to export historical experiment data from Weights & Biases to local parquet files and load them into Pluto while preserving original timestamps. The changes include new CLI commands, exporter and loader modules, resume bookkeeping, and support for historical timestamps and batch logging in the core Op class. The code review identified critical path traversal vulnerabilities when handling externally-sourced identifiers (such as run IDs, file names, and artifact names) from the wandb API and staged files. Additionally, several bugs and robustness issues were highlighted, including a missing defaultdict import in pluto/op.py, potential crashes from malformed JSON or non-string inputs, uncleaned temporary directories on export failure, and platform-dependent file opening without explicit UTF-8 encoding.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The defaultdict class is used here but it is not imported in this file, which will raise a NameError when _log_metrics_batch is called. Import defaultdict from collections locally or at the top of the file to prevent this crash.

Suggested change

	new_metric_names: List[str] = []
	new_file_meta: Dict[str, List[str]] = defaultdict(list)
	from collections import defaultdict
	new_metric_names: List[str] = []
	new_file_meta: Dict[str, List[str]] = defaultdict(list)

[gemini-code-assist]

The run.id is an externally-sourced identifier from the wandb API. Using it directly to construct run_dir can lead to path traversal vulnerabilities if a malicious run ID contains .. or path separators. Sanitize it using Path(run.id).name and check for unsafe values like . or .. before constructing the path.

Suggested change

	run_dir = runs_root / run.id
	safe_run_id = Path(run.id).name
	if not safe_run_id or safe_run_id in ('.', '..'):
	logger.warning(f'{tag}: unsafe run ID {run.id!r}, skipping')
	continue
	run_dir = runs_root / safe_run_id

References

1. Sanitize externally-sourced identifiers before using them to construct file paths to prevent path traversal vulnerabilities. Use pathlib.Path(identifier).name or os.path.basename(identifier) and check for unsafe values like . or ...

[gemini-code-assist]

The f.name is an externally-sourced file path from the wandb API. To prevent path traversal vulnerabilities, sanitize it using Path(f.name).name and check for unsafe values like . or .. before downloading.

        for f in run.files():
            if not self.include_files and f.name != 'output.log':
                continue
            safe_name = Path(f.name).name
            if not safe_name or safe_name in ('.', '..') or safe_name != f.name:
                logger.warning(f'{tag}: unsafe file name {f.name!r}, skipping')
                continue
            try:
                f.download(root=str(files_dir), exist_ok=True)

References

1. Sanitize externally-sourced identifiers before using them to construct file paths to prevent path traversal vulnerabilities. Use pathlib.Path(identifier).name or os.path.basename(identifier) and check for unsafe values like . or ...

[gemini-code-assist]

The artifact.name is an externally-sourced identifier from the wandb API. To prevent path traversal vulnerabilities, sanitize it using Path(artifact.name).name and check for unsafe values like . or .. before constructing the download destination path.

            safe_name = Path(artifact.name).name
            if not safe_name or safe_name in ('.', '..'):
                logger.warning(f'{tag}: unsafe artifact name {artifact.name!r}, skipping')
                continue
            dest = tmp_dir / 'artifacts' / safe_name

References

1. Sanitize externally-sourced identifiers before using them to construct file paths to prevent path traversal vulnerabilities. Use pathlib.Path(identifier).name or os.path.basename(identifier) and check for unsafe values like . or ...

[gemini-code-assist]

The row['file_value'] is read from the staged parquet files, which could be manipulated or contain unsafe paths. To prevent path traversal vulnerabilities, sanitize it using Path(file_value).name and check for unsafe values like . or .. before checking for its existence or loading it.

Suggested change

	path = run_dir / (row['file_value'] or '')
	if not row['file_value'] or not path.exists():
	file_value = row['file_value'] or ''
	safe_file_value = Path(file_value).name
	if not safe_file_value or safe_file_value in ('.', '..') or safe_file_value != file_value:
	logger.warning(f'{tag}: unsafe file path {file_value!r}, skipping')
	return
	path = run_dir / safe_file_value
	if not path.exists():

References

1. Sanitize externally-sourced identifiers before using them to construct file paths to prevent path traversal vulnerabilities. Use pathlib.Path(identifier).name or os.path.basename(identifier) and check for unsafe values like . or ...

[gemini-code-assist]

The parse_iso_ms function assumes value is always a string. If the wandb API returns a non-string type (such as an integer or a datetime object), calling .replace will raise an AttributeError and crash the export. Add a defensive check to ensure value is a string.

Suggested change

	def parse_iso_ms(value: Optional[str]) -> Optional[int]:
	"""Parse an ISO-8601 string to epoch milliseconds (UTC assumed if naive)."""
	if not value:
	return None
	def parse_iso_ms(value: Optional[str]) -> Optional[int]:
	"""Parse an ISO-8601 string to epoch milliseconds (UTC assumed if naive)."""
	if not isinstance(value, str):
	return None
	try:

[gemini-code-assist]

If an exception occurs during _export_run (e.g., during history scanning or file downloading), the partial .tmp directory is left on disk. Wrap the export logic in a try-except block to clean up tmp_dir on failure so orphaned temporary directories do not accumulate.

    def _export_run(self, run: Any, run_dir: Path) -> None:
        tmp_dir = run_dir.with_name(run_dir.name + '.tmp')
        if tmp_dir.exists():
            shutil.rmtree(tmp_dir)  # leftovers from an interrupted export
        tmp_dir.mkdir(parents=True)

        try:
            created_ms = parse_iso_ms(getattr(run, 'created_at', None))
            with PartWriter(tmp_dir) as writer:
                self._write_run_json(run, tmp_dir, created_ms)
                self._export_history(run, writer)
                if self.include_system:
                    self._export_system_metrics(run, writer)
                files_dir = tmp_dir / 'files'
                if self.include_files or self.include_console:
                    self._download_files(run, files_dir)
                if self.include_console:
                    self._export_console(run, writer, files_dir, created_ms)
                if self.include_artifacts:
                    self._export_artifacts(run, writer, tmp_dir)

            mark_run_exported(tmp_dir, {'rows': writer.rows_written})
            if run_dir.exists():
                shutil.rmtree(run_dir)
            os.rename(tmp_dir, run_dir)
        except Exception:
            if tmp_dir.exists():
                shutil.rmtree(tmp_dir)
            raise

[gemini-code-assist]

Opening files without specifying an explicit encoding is platform-dependent and can lead to UnicodeDecodeError or corrupt characters on platforms where the default encoding is not UTF-8 (such as Windows). Specify encoding='utf-8' explicitly.

Suggested change

	with open(output_log, errors='replace') as f:
	with open(output_log, encoding='utf-8', errors='replace') as f:

[gemini-code-assist]

Opening files without specifying an explicit encoding is platform-dependent and can lead to UnicodeDecodeError or corrupt characters on platforms where the default encoding is not UTF-8 (such as Windows). Specify encoding='utf-8' explicitly.

Suggested change

	with open(tmp, 'w') as f:
	with open(tmp, 'w', encoding='utf-8') as f:

[gemini-code-assist]

Opening files without specifying an explicit encoding is platform-dependent and can lead to UnicodeDecodeError or corrupt characters on platforms where the default encoding is not UTF-8 (such as Windows). Specify encoding='utf-8' explicitly.

Suggested change

	with open(path) as f:
	with open(path, encoding='utf-8') as f:

[asaiacai]

CI status: all code-quality gates are green (mypy, format, API-docs check, CodeQL, prod + staging contract tests, k8s-multinode). The remaining pytest-smoke failures are an environmental prod flake, not this PR: intermittent 401 {"error":"Unauthorized","message":"Invalid API key"} from /api/runs/create, hitting a different pre-existing test on a different Python leg each round (3.10 → 3.12 → 3.11 → 3.12+3.13 across four runs), while the same key serves hundreds of requests in the same job. Main's recent history shows the same occasional failures. Looks like the server-side api-key lookup transiently failing under load and surfacing as 401 instead of 5xx — filed to look at in server-private's key-validation path.

🤖 Generated with Claude Code

[asaiacai]

Correction on the flake diagnosis above: the root cause was client-side, not server-side — and it's now fixed in this PR (46e52ce).

login() posts the token to /api/slug as a best-effort validation. When that single POST hit a transient network error, the except branch overwrote settings._auth with the '_key' sentinel even when the token was provided via PLUTO_API_KEY, so every subsequent request from that run sent Bearer _key → the server's prefix check correctly returned 401 "Invalid API key" for all retries. That explains the random victim test per round, the all-5-retries-fail pattern, and the load correlation (more concurrent suites → more transient blips → more corrupted runs). The bug dates to the January rename commit, matching the occasional failures in main's history. The server needs no change.

Regression-tested in tests/test_auth_transient.py (reproduces the exact corruption, plus pins the interactive-flow sentinel for the no-token case).

🤖 Generated with Claude Code