diff --git a/docs-api/run.mdx b/docs-api/run.mdx index 92b99c4..ccab6c3 100644 --- a/docs-api/run.mdx +++ b/docs-api/run.mdx @@ -13,10 +13,16 @@ log( data: Dict[str, Any], step: Union[int, None] = None, commit: Union[bool, None] = None, + timestamp: Optional[float] = None, ) -> None ``` -Log run data +Log run data. + +`timestamp` is the wall-clock time of the data points in epoch +seconds (`time.time()` style) and defaults to now. The server +stores it as-is, so backfill/migration tooling can preserve +historical times. Invalid values fall back to now with a warning. --- diff --git a/pluto/__main__.py b/pluto/__main__.py index 85f96b4..ab5cd72 100644 --- a/pluto/__main__.py +++ b/pluto/__main__.py @@ -240,6 +240,13 @@ def main(): help='show detailed sync progress', ) + # `pluto migrate wandb export|load|all` — import historical data. + # cli.py keeps wandb/pyarrow imports inside the handlers, so this + # is safe without the 'migrate' extra installed. + from pluto.migrate.cli import add_migrate_parser, cmd_migrate + + add_migrate_parser(subparsers) + args = parser.parse_args() if args.version: @@ -259,6 +266,8 @@ def main(): logout() elif args.command == 'sync': _cmd_sync(args) + elif args.command == 'migrate': + sys.exit(cmd_migrate(args)) else: parser.print_help() sys.exit(1) diff --git a/pluto/auth.py b/pluto/auth.py index eda752b..7349cd6 100644 --- a/pluto/auth.py +++ b/pluto/auth.py @@ -75,7 +75,12 @@ def login(settings=None, retry=False): ) except Exception as e: tlogger.warning(f'{tag}: server not reachable; reason: {e}') - settings._auth = '_key' + # A transient failure of this best-effort validation POST must not + # corrupt an explicitly provided token: overwriting it with the + # '_key' sentinel makes every later request send 'Bearer _key', + # which servers reject as 401 "Invalid API key" for the whole run. + if not auth_was_provided: + settings._auth = '_key' try: r.raise_for_status() body = r.json() diff --git a/pluto/migrate/__init__.py b/pluto/migrate/__init__.py new file mode 100644 index 0000000..6a9c35a --- /dev/null +++ b/pluto/migrate/__init__.py @@ -0,0 +1,29 @@ +""" +pluto.migrate — import historical experiment data into Pluto. + +Two-phase pipeline: export a source platform's runs to on-disk parquet +(``pluto migrate wandb export``), then load the staged data into Pluto +(``pluto migrate wandb load``). Both phases are resumable. + +Requires the ``migrate`` extra: ``pip install 'pluto-ml[migrate]'``. +""" + +from typing import Any + +_INSTALL_HINT = ( + "pluto.migrate requires the 'migrate' extra. " + "Install it with: pip install 'pluto-ml[migrate]'" +) + + +def __getattr__(name: str) -> Any: + # Lazy so `import pluto` never pays for (or requires) wandb/pyarrow. + if name == 'WandbExporter': + from pluto.migrate.wandb_export import WandbExporter + + return WandbExporter + if name == 'PlutoLoader': + from pluto.migrate.loader import PlutoLoader + + return PlutoLoader + raise AttributeError(f'module {__name__!r} has no attribute {name!r}') diff --git a/pluto/migrate/cli.py b/pluto/migrate/cli.py new file mode 100644 index 0000000..c192514 --- /dev/null +++ b/pluto/migrate/cli.py @@ -0,0 +1,213 @@ +""" +CLI for pluto.migrate: `pluto migrate wandb export|load|all`. + +This module keeps its imports light — wandb/pyarrow (the 'migrate' +extra) load inside the command handlers, so the top-level `pluto` CLI +works without them and missing deps produce an install hint instead of +an ImportError traceback. +""" + +from __future__ import annotations + +import argparse +from typing import List, Optional + +from pluto.migrate import _INSTALL_HINT + + +def _add_export_flags(parser: argparse.ArgumentParser) -> None: + parser.add_argument('--entity', required=True, help='wandb entity (team/user)') + parser.add_argument('--project', required=True, help='wandb project name') + parser.add_argument( + '--output', required=True, help='directory to stage exported data in' + ) + parser.add_argument( + '--wandb-api-key', + help='wandb API key (default: WANDB_API_KEY / wandb login)', + ) + parser.add_argument( + '--run-id', + action='append', + dest='run_ids', + help='only these wandb run ids (repeatable)', + ) + parser.add_argument('--after', help='only runs created after this ISO date') + parser.add_argument('--before', help='only runs created before this ISO date') + parser.add_argument( + '--no-artifacts', action='store_true', help='skip logged artifacts' + ) + parser.add_argument( + '--artifact-max-size-mb', + type=int, + help='skip artifacts larger than this many MB', + ) + parser.add_argument( + '--no-console', action='store_true', help='skip console output.log' + ) + parser.add_argument( + '--no-system-metrics', + action='store_true', + help='skip GPU/CPU system metrics', + ) + parser.add_argument( + '--no-files', action='store_true', help='skip media/file downloads' + ) + + +def _add_load_flags(parser: argparse.ArgumentParser, with_input: bool = True) -> None: + if with_input: + parser.add_argument( + '--input', required=True, help='export directory to load from' + ) + parser.add_argument( + '--run-id', + action='append', + dest='run_ids', + help='only these wandb run ids (repeatable)', + ) + parser.add_argument( + '--dest-project', + help='Pluto project to load into (default: the wandb project name)', + ) + parser.add_argument( + '--dry-run', + action='store_true', + help='print what would be loaded without creating runs', + ) + parser.add_argument( + '--force-resume', + action='store_true', + help='re-load runs already marked loaded (resumes via external id; ' + 'may duplicate media files)', + ) + parser.add_argument( + '--flush-every', + type=int, + default=500, + help='steps between sync-queue backpressure checks (default: 500)', + ) + parser.add_argument( + '--max-pending', + type=int, + default=5000, + help='max queued records before the loader throttles (default: 5000)', + ) + + +def add_migrate_parser(subparsers: argparse._SubParsersAction) -> None: + """Attach the `migrate` subcommand to the top-level pluto CLI.""" + p_migrate = subparsers.add_parser( + 'migrate', help='import historical data from another platform' + ) + sources = p_migrate.add_subparsers(dest='source', required=True) + + p_wandb = sources.add_parser('wandb', help='migrate from Weights & Biases') + actions = p_wandb.add_subparsers(dest='action', required=True) + + p_export = actions.add_parser( + 'export', help='download wandb runs to a local staging directory' + ) + _add_export_flags(p_export) + + p_load = actions.add_parser( + 'load', help='load a staged export directory into Pluto' + ) + _add_load_flags(p_load) + + p_all = actions.add_parser('all', help='export then load in one go') + _add_export_flags(p_all) + _add_load_flags(p_all, with_input=False) + + +def _run_export(args: argparse.Namespace) -> int: + try: + from pluto.migrate.wandb_export import WandbExporter + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + + exporter = WandbExporter( + entity=args.entity, + project=args.project, + output_dir=args.output, + api_key=args.wandb_api_key, + run_ids=args.run_ids, + after=args.after, + before=args.before, + include_artifacts=not args.no_artifacts, + artifact_max_bytes=( + args.artifact_max_size_mb * 1024 * 1024 + if args.artifact_max_size_mb is not None + else None + ), + include_console=not args.no_console, + include_system=not args.no_system_metrics, + include_files=not args.no_files, + ) + summary = exporter.export() + print( + f'export: {summary["exported"]} exported, {summary["skipped"]} skipped, ' + f'{len(summary["failed"])} failed' + ) + for failure in summary['failed']: + print(f' failed {failure["run_id"]}: {failure["error"]}') + return 1 if summary['failed'] else 0 + + +def _run_load(args: argparse.Namespace, input_dir: Optional[str] = None) -> int: + try: + from pluto.migrate.loader import PlutoLoader + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + + loader = PlutoLoader( + input_dir=input_dir if input_dir is not None else args.input, + dest_project=args.dest_project, + flush_every=args.flush_every, + max_pending=args.max_pending, + dry_run=args.dry_run, + run_ids=getattr(args, 'run_ids', None), + force_resume=args.force_resume, + ) + summary = loader.load() + print( + f'load: {summary["loaded"]} loaded, {summary["skipped"]} skipped, ' + f'{len(summary["failed"])} failed' + ) + for failure in summary['failed']: + print(f' failed {failure["run_id"]}: {failure["error"]}') + return 1 if summary['failed'] else 0 + + +def cmd_migrate(args: argparse.Namespace) -> int: + if args.action == 'export': + return _run_export(args) + if args.action == 'load': + return _run_load(args) + if args.action == 'all': + if args.dry_run: + print( + 'error: --dry-run is not supported with `all` (it would still ' + 'download everything). Run `export` first, then preview with ' + '`load --dry-run`.' + ) + return 2 + export_code = _run_export(args) + if export_code == 2: # missing deps — nothing was staged + return export_code + # Per-run export failures must not block loading the runs that DID + # stage successfully; both phases are independently resumable. + load_code = _run_load(args, input_dir=args.output) + return max(export_code, load_code) + raise AssertionError(f'unknown action {args.action!r}') + + +def run_migrate(argv: List[str]) -> int: + """Standalone entry (also used by tests): argv excludes 'migrate'.""" + parser = argparse.ArgumentParser(prog='pluto migrate') + subparsers = parser.add_subparsers(dest='command', required=True) + # Reuse the same tree shape as the top-level CLI: migrate -> wandb -> action + add_migrate_parser(subparsers) + args = parser.parse_args(['migrate', *argv]) + return cmd_migrate(args) diff --git a/pluto/migrate/loader.py b/pluto/migrate/loader.py new file mode 100644 index 0000000..010c477 --- /dev/null +++ b/pluto/migrate/loader.py @@ -0,0 +1,354 @@ +""" +Load staged export directories into Pluto through the public client API. + +Replays each exported run — run.json manifest plus parquet parts — as a +Pluto run with the ORIGINAL wall-clock timestamps (``op.log(timestamp=)``) +and creation time (``settings.compat`` createdAt/updatedAt). Idempotency +is run-level: a ``run_id`` external id (``wandb::{entity}/{project}/{id}``) +makes re-creation collide server-side (the loader then resumes and +re-replays, so an interrupted load heals on the next run), and +``loaded_runs.json`` records finished loads so re-runs skip them. +Replayed metric points carry identical staged timestamps, so the +backend's replace-by-time dedup makes re-replays safe for metrics; +re-replayed media files may duplicate. +""" + +from __future__ import annotations + +import json +import logging +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import pluto +from pluto.migrate.schema import iter_part_tables, part_files +from pluto.migrate.state import ( + LOADED_CACHE_FILENAME, + LoadedCache, + is_run_exported, + read_json, +) +from pluto.op import RunExistsError + +logger = logging.getLogger(f'{__name__.split(".")[0]}') +tag = 'migrate' + +CONSOLE_BATCH_SIZE = 1000 + +# All three take (data, caption=...); table-file and inline histograms are +# handled separately in _replay_media. +_MEDIA_LOADERS = { + 'image-file': pluto.Image, + 'audio-file': pluto.Audio, + 'video-file': pluto.Video, +} + + +class PlutoLoader: + """Replay a pluto.migrate export directory into Pluto.""" + + def __init__( + self, + input_dir: Union[str, Path], + dest_project: Optional[str] = None, + flush_every: int = 500, + max_pending: int = 5000, + dry_run: bool = False, + run_ids: Optional[List[str]] = None, + force_resume: bool = False, + stall_timeout: float = 600.0, + ) -> None: + self.input_dir = Path(input_dir) + self.dest_project = dest_project + self.flush_every = flush_every + self.max_pending = max_pending + self.dry_run = dry_run + self.run_ids = set(run_ids) if run_ids else None + self.force_resume = force_resume + self.stall_timeout = stall_timeout + + def load(self) -> Dict[str, Any]: + """Load all staged runs. Returns {'loaded', 'skipped', 'failed'}.""" + loaded, skipped = 0, 0 + failed: List[Dict[str, str]] = [] + cache = LoadedCache(self.input_dir / LOADED_CACHE_FILENAME) + + for run_dir in self._discover_runs(): + manifest = read_json(run_dir / 'run.json') + run_id = manifest['run_id'] + if self.run_ids is not None and run_id not in self.run_ids: + continue + external_id = f'wandb::{manifest["entity"]}/{manifest["project"]}/{run_id}' + if cache.is_loaded(external_id) and not self.force_resume: + logger.info(f'{tag}: {external_id} already loaded, skipping') + skipped += 1 + continue + if self.dry_run: + self._print_dry_run(run_dir, manifest, external_id) + continue + + try: + op = self._init_run(manifest, external_id) + except RunExistsError: + # The run exists server-side but isn't in loaded_runs.json — + # a previous load was interrupted mid-replay. Resume and + # re-replay (metrics dedup by identical timestamps; media + # from the interrupted portion may duplicate). + logger.warning( + f'{tag}: {external_id} already exists on server ' + '(previous load interrupted?); resuming to re-replay' + ) + op = self._init_run(manifest, external_id, resume=True) + + try: + self._replay_run(run_dir, op) + op.finish(code=0 if manifest.get('state') == 'finished' else 1) + cache.mark_loaded(external_id, {'pluto_run_id': op.settings._op_id}) + loaded += 1 + logger.info(f'{tag}: loaded {external_id}') + except Exception as e: + logger.error(f'{tag}: load failed for {external_id}: {e}') + failed.append({'run_id': run_id, 'error': f'{type(e).__name__}: {e}'}) + try: + op.finish(code=1) + except Exception: + pass + + return {'loaded': loaded, 'skipped': skipped, 'failed': failed} + + def _discover_runs(self) -> List[Path]: + return sorted( + d + for d in self.input_dir.glob('*/*/runs/*') + if d.is_dir() and is_run_exported(d) and (d / 'run.json').exists() + ) + + def _init_run( + self, + manifest: Dict[str, Any], + external_id: str, + resume: Optional[bool] = None, + ) -> Any: + tags = list(manifest.get('tags') or []) + if 'import:wandb' not in tags: + tags.append('import:wandb') + settings: Dict[str, Any] = { + 'compat': { + 'createdAt': manifest.get('createdAt'), + 'updatedAt': manifest.get('updatedAt'), + }, + # Never attribute the migration host's console/hardware to the + # imported run. + 'disable_console': True, + 'disable_system_metrics': True, + # The historical-timestamp path only exists in the sync store; + # force it on even if the user's defaults disable it. + 'sync_process_enabled': True, + } + op = pluto.init( + project=self.dest_project or manifest['project'], + name=manifest.get('name'), + config=manifest.get('config') or None, + tags=tags, + run_id=external_id, + resume=self.force_resume if resume is None else resume, + settings=settings, + ) + wandb_block = { + k: v + for k, v in { + 'notes': manifest.get('notes'), + 'url': manifest.get('url'), + 'state': manifest.get('state'), + 'summary': manifest.get('summary'), + }.items() + if v + } + if wandb_block: + op.update_config({'wandb': wandb_block}) + return op + + @staticmethod + def _sys_metric_name(name: str) -> str: + """Map a source-native system metric name into Pluto's sys/ namespace.""" + if name.startswith('system.'): + return 'sys/' + name[len('system.') :] + if name.startswith('sys/'): + return name + return f'sys/{name}' + + def _replay_run(self, run_dir: Path, op: Any) -> None: + # (attribute_type, step, timestamp_ms) of the group being buffered; + # rows are staged in write order so same-step metrics are contiguous. + group_key: Optional[Tuple[str, int, int]] = None + group_metrics: Dict[str, float] = {} + # Closed groups accumulate here and flush through one SQLite + # transaction per flush_every groups (op._log_metrics_batch). + pending_groups: List[Tuple[Dict[str, float], int, float]] = [] + console_lines: List[Tuple[str, str, float, int]] = [] + + def close_group() -> None: + nonlocal group_key, group_metrics + if group_key is not None and group_metrics: + _, step, timestamp_ms = group_key + pending_groups.append((group_metrics, step, timestamp_ms / 1000)) + group_key, group_metrics = None, {} + + def flush_pending(force: bool = False) -> None: + if pending_groups and (force or len(pending_groups) >= self.flush_every): + op._log_metrics_batch(list(pending_groups)) + pending_groups.clear() + self._wait_for_backpressure(op) + + for table in iter_part_tables(run_dir): + for row in table.to_pylist(): + attr_type = row['attribute_type'] + if attr_type in ('metric', 'system_metric'): + key = (attr_type, row['step'], row['timestamp_ms']) + if key != group_key: + close_group() + flush_pending() + group_key = key + name = row['attribute_path'] + if attr_type == 'system_metric': + name = self._sys_metric_name(name) + group_metrics[name] = row['float_value'] + continue + close_group() + if attr_type == 'media': + self._replay_media(run_dir, op, row) + elif attr_type == 'console': + console_lines.append( + ( + row['string_value'] or '', + 'INFO', + row['timestamp_ms'] / 1000, + row['step'], + ) + ) + if len(console_lines) >= CONSOLE_BATCH_SIZE: + op._log_console(console_lines) + console_lines = [] + elif attr_type == 'artifact': + self._replay_artifact(run_dir, op, row) + + close_group() + flush_pending(force=True) + if console_lines: + op._log_console(console_lines) + + def _replay_media(self, run_dir: Path, op: Any, row: Dict[str, Any]) -> None: + name = row['attribute_path'] + step = row['step'] + timestamp = row['timestamp_ms'] / 1000 + media_type = row['string_value'] or '' + + if media_type.startswith('{'): # inline JSON (histogram) + payload = json.loads(media_type) + if payload.get('_type') == 'histogram': + op.log( + { + name: pluto.Histogram( + [payload.get('values'), payload.get('bins')], + bins=None, + ) + }, + step=step, + timestamp=timestamp, + ) + return + + path = run_dir / (row['file_value'] or '') + if not row['file_value'] or not path.exists(): + logger.warning( + f'{tag}: media file missing for {name!r} ' + f'({row["file_value"]}), skipping' + ) + return + + caption = row['caption'] + if media_type == 'table-file': + try: + table_json = read_json(path) + value: Any = pluto.Table( + data=table_json.get('data'), + columns=table_json.get('columns', []), + ) + except Exception as e: + logger.warning(f'{tag}: bad table file {path}: {e}') + return + else: + make = _MEDIA_LOADERS.get(media_type) + if make is not None: + value = make(str(path), caption=caption) + else: # plotly/html/object3D/unknown -> raw artifact + value = pluto.Artifact(str(path), caption=caption) + op.log({name: value}, step=step, timestamp=timestamp) + + def _replay_artifact(self, run_dir: Path, op: Any, row: Dict[str, Any]) -> None: + path = run_dir / (row['file_value'] or '') + if not row['file_value'] or not path.exists(): + logger.warning( + f'{tag}: artifact file missing ({row["file_value"]}), skipping' + ) + return + metadata = None + if row['string_value']: + try: + metadata = json.loads(row['string_value']) + except ValueError: + metadata = None + op.log( + { + row['attribute_path']: pluto.Artifact( + str(path), caption=path.name, metadata=metadata + ) + }, + step=row['step'], + timestamp=row['timestamp_ms'] / 1000, + ) + + def _wait_for_backpressure(self, op: Any) -> None: + """Bound the sync queue so huge runs don't balloon SQLite/memory. + + Bounded by ``stall_timeout``: if the queue never drains (dead sync + process, unreachable server) the loader logs and moves on rather + than hanging — data stays in the sync store either way. + """ + deadline = time.time() + self.stall_timeout + sleep_s = 0.5 + throttled = False + while True: + try: + pending = op._sync_manager.get_pending_count() + except Exception: + return + if pending <= self.max_pending: + return + if time.time() >= deadline: + logger.warning( + f'{tag}: sync queue still has {pending} pending records ' + f'after {self.stall_timeout:.0f}s; continuing (uploads ' + 'proceed in the background)' + ) + return + if not throttled: + logger.info( + f'{tag}: {pending} records pending upload, throttling loader' + ) + throttled = True + time.sleep(sleep_s) + sleep_s = min(sleep_s * 2, 5.0) + + def _print_dry_run( + self, run_dir: Path, manifest: Dict[str, Any], external_id: str + ) -> None: + parts = part_files(run_dir) + size = sum(p.stat().st_size for p in parts) + print( + f'[dry-run] {external_id} -> project ' + f'{self.dest_project or manifest["project"]!r} ' + f'name={manifest.get("name")!r} ' + f'parts={len(parts)} ({size / 1e6:.1f} MB)' + ) diff --git a/pluto/migrate/schema.py b/pluto/migrate/schema.py new file mode 100644 index 0000000..86f87e8 --- /dev/null +++ b/pluto/migrate/schema.py @@ -0,0 +1,162 @@ +""" +Parquet staging format for pluto.migrate. + +One long/tall schema holds every per-point record of a run (metrics, +system metrics, media references, console lines, artifact files); +``attribute_type`` discriminates. Run-scalar data (name, config, tags, +summary, timestamps) lives in the sibling ``run.json`` manifest instead. + +Rows are written through :class:`PartWriter`, which rotates output files +(``part-00000.parquet``, ``part-00001.parquet``, ...) once the current +part exceeds ``max_part_bytes`` on disk, so arbitrarily long runs stage +in bounded memory and load back part by part. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Iterator, List, Optional, Union + +import pyarrow as pa +import pyarrow.parquet as pq + +ATTRIBUTE_TYPES = { + 'metric', # scalar metric history point (float_value) + 'system_metric', # host stats history point (float_value) + 'media', # image/audio/video/table/histogram (string_value=_type, + # file_value=relative path or inline JSON in string_value) + 'console', # one console line (string_value, step=line number) + 'artifact', # one file inside a logged artifact (file_value) +} + +SCHEMA = pa.schema( + [ + pa.field('project_id', pa.string()), + pa.field('run_id', pa.string()), + pa.field('attribute_path', pa.string()), + pa.field('attribute_type', pa.string()), + pa.field('step', pa.int64()), + pa.field('timestamp_ms', pa.int64()), + pa.field('float_value', pa.float64()), + pa.field('string_value', pa.string()), + pa.field('file_value', pa.string()), + pa.field('caption', pa.string()), + ] +) + +_COLUMNS = [f.name for f in SCHEMA] + +DEFAULT_MAX_PART_BYTES = 50 * 1024 * 1024 +DEFAULT_ROWS_PER_FLUSH = 20_000 + +PART_PREFIX = 'part-' +PART_SUFFIX = '.parquet' + + +class PartWriter: + """Buffered writer producing rotated parquet parts for one run. + + Rows accumulate in memory and flush to the open part every + ``rows_per_flush`` rows; after each flush the part rotates if it + outgrew ``max_part_bytes``. Use as a context manager so the trailing + buffer is flushed and the last part closed. + """ + + def __init__( + self, + run_dir: Union[str, Path], + max_part_bytes: int = DEFAULT_MAX_PART_BYTES, + rows_per_flush: int = DEFAULT_ROWS_PER_FLUSH, + ) -> None: + self._run_dir = Path(run_dir) + self._run_dir.mkdir(parents=True, exist_ok=True) + self._max_part_bytes = max_part_bytes + self._rows_per_flush = rows_per_flush + self._buffer: List[dict] = [] + self._part_index = 0 + self._writer: Optional[pq.ParquetWriter] = None + self._part_path: Optional[Path] = None + self.rows_written = 0 + + def __enter__(self) -> 'PartWriter': + return self + + def __exit__(self, *exc_info: Any) -> None: + self.close() + + def write_row( + self, + *, + project_id: str, + run_id: str, + attribute_path: str, + attribute_type: str, + timestamp_ms: int, + step: Optional[int] = None, + float_value: Optional[float] = None, + string_value: Optional[str] = None, + file_value: Optional[str] = None, + caption: Optional[str] = None, + ) -> None: + if attribute_type not in ATTRIBUTE_TYPES: + raise ValueError( + f'unknown attribute_type {attribute_type!r}, ' + f'expected one of {sorted(ATTRIBUTE_TYPES)}' + ) + self._buffer.append( + { + 'project_id': project_id, + 'run_id': run_id, + 'attribute_path': attribute_path, + 'attribute_type': attribute_type, + 'step': step, + 'timestamp_ms': timestamp_ms, + 'float_value': float_value, + 'string_value': string_value, + 'file_value': file_value, + 'caption': caption, + } + ) + if len(self._buffer) >= self._rows_per_flush: + self._flush() + + def close(self) -> None: + self._flush() + self._close_part() + + def _flush(self) -> None: + if not self._buffer: + return + table = pa.Table.from_pylist(self._buffer, schema=SCHEMA) + self._buffer = [] + if self._writer is None: + self._part_path = self._run_dir / ( + f'{PART_PREFIX}{self._part_index:05d}{PART_SUFFIX}' + ) + self._writer = pq.ParquetWriter( + self._part_path, SCHEMA, compression='snappy' + ) + self._part_index += 1 + self._writer.write_table(table) + self.rows_written += table.num_rows + assert self._part_path is not None # set alongside _writer above + if os.path.getsize(self._part_path) >= self._max_part_bytes: + self._close_part() + + def _close_part(self) -> None: + if self._writer is not None: + self._writer.close() + self._writer = None + self._part_path = None + + +def part_files(run_dir: Union[str, Path]) -> List[Path]: + """Return a run's parquet parts in write order.""" + return sorted(Path(run_dir).glob(f'{PART_PREFIX}*{PART_SUFFIX}')) + + +def iter_part_tables(run_dir: Union[str, Path]) -> Iterator[pa.Table]: + """Yield each parquet part as a table, in write order (bounded memory).""" + for path in part_files(run_dir): + yield pq.read_table(path, schema=SCHEMA) diff --git a/pluto/migrate/state.py b/pluto/migrate/state.py new file mode 100644 index 0000000..a225507 --- /dev/null +++ b/pluto/migrate/state.py @@ -0,0 +1,61 @@ +""" +Resume bookkeeping for pluto.migrate. + +Export marks each fully-staged run with a sentinel file written last, so +an interrupted export re-does at most one run. Load records finished +runs in a single ``loaded_runs.json`` next to the export, written only +after ``finish()`` drains — so a re-run skips completed runs and retries +partial ones. All writes are atomic (tmp file + fsync + rename). +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any, Dict, Optional, Union + +EXPORT_SENTINEL = '_export_complete.json' +LOADED_CACHE_FILENAME = 'loaded_runs.json' + + +def write_json_atomic(path: Union[str, Path], obj: Any) -> None: + path = Path(path) + tmp = path.with_name(path.name + '.tmp') + with open(tmp, 'w') as f: + json.dump(obj, f, indent=2, default=str) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def read_json(path: Union[str, Path]) -> Any: + with open(path) as f: + return json.load(f) + + +def is_run_exported(run_dir: Union[str, Path]) -> bool: + return (Path(run_dir) / EXPORT_SENTINEL).exists() + + +def mark_run_exported( + run_dir: Union[str, Path], summary: Optional[Dict[str, Any]] = None +) -> None: + write_json_atomic(Path(run_dir) / EXPORT_SENTINEL, summary or {}) + + +class LoadedCache: + """Load-phase resume cache: which external ids finished loading.""" + + def __init__(self, path: Union[str, Path]) -> None: + self._path = Path(path) + self._entries: Dict[str, Any] = {} + if self._path.exists(): + self._entries = read_json(self._path) + + def is_loaded(self, external_id: str) -> bool: + return external_id in self._entries + + def mark_loaded(self, external_id: str, info: Dict[str, Any]) -> None: + self._entries[external_id] = info + write_json_atomic(self._path, self._entries) diff --git a/pluto/migrate/wandb_export.py b/pluto/migrate/wandb_export.py new file mode 100644 index 0000000..b747eca --- /dev/null +++ b/pluto/migrate/wandb_export.py @@ -0,0 +1,414 @@ +""" +Export wandb cloud runs to the on-disk pluto.migrate staging format. + +Reads complete run data through the wandb public API (``wandb.Api()``) +and stages each run as ``run.json`` + parquet parts + downloaded files +under ``output_dir/{entity}/{project}/runs/{run_id}/``. Runs are staged +in a ``.tmp`` directory renamed into place only after the export +sentinel is written, so an interrupted export never leaves a directory +that looks complete; re-running skips finished runs. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, TypedDict, Union + +from pluto.migrate.schema import PartWriter +from pluto.migrate.state import is_run_exported, mark_run_exported, write_json_atomic + +logger = logging.getLogger(f'{__name__.split(".")[0]}') +tag = 'migrate' + +MANIFEST_FILENAME = 'manifest.json' + + +class _RowBase(TypedDict): + """Identity columns shared by every staged row (see schema.write_row).""" + + project_id: str + run_id: str + + +# scan_history dict values whose media file lives under the run's files/ +_FILE_MEDIA_TYPES = { + 'image-file', + 'audio-file', + 'video-file', + 'table-file', + 'plotly-file', + 'object3D-file', + 'html-file', +} + +# Leading " " console lines (wandb writes these +# when x_show_timestamps is enabled). +_CONSOLE_TS_RE = re.compile( + r'^(\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:[.,]\d+)?' + r'(?:Z|[+-]\d{2}:?\d{2})?)\s+(.*)$' +) + + +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 + try: + normalized = value.replace(',', '.').replace('Z', '+00:00') + dt = datetime.fromisoformat(normalized) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return int(dt.timestamp() * 1000) + except ValueError: + return None + + +class WandbExporter: + """Stage one wandb project's runs on disk for later loading into Pluto.""" + + def __init__( + self, + entity: str, + project: str, + output_dir: Union[str, Path], + api: Optional[Any] = None, + api_key: Optional[str] = None, + run_ids: Optional[List[str]] = None, + after: Optional[str] = None, + before: Optional[str] = None, + include_artifacts: bool = True, + artifact_max_bytes: Optional[int] = None, + include_console: bool = True, + include_system: bool = True, + include_files: bool = True, + history_page_size: int = 1000, + ) -> None: + self.entity = entity + self.project = project + self.output_dir = Path(output_dir) + self._api = api + self._api_key = api_key + self.run_ids = set(run_ids) if run_ids else None + self.after_ms = parse_iso_ms(after) + self.before_ms = parse_iso_ms(before) + self.include_artifacts = include_artifacts + self.artifact_max_bytes = artifact_max_bytes + self.include_console = include_console + self.include_system = include_system + self.include_files = include_files + self.history_page_size = history_page_size + + @property + def api(self) -> Any: + if self._api is None: + import wandb + + self._api = wandb.Api(api_key=self._api_key) + return self._api + + @property + def project_path(self) -> str: + return f'{self.entity}/{self.project}' + + def export(self) -> Dict[str, Any]: + """Export all matching runs. Returns {'exported', 'skipped', 'failed'}.""" + exported, skipped = 0, 0 + failed: List[Dict[str, str]] = [] + + runs_root = self.output_dir / self.entity / self.project / 'runs' + for run in self.api.runs(self.project_path): + if self.run_ids is not None and run.id not in self.run_ids: + continue + created_ms = parse_iso_ms(getattr(run, 'created_at', None)) + if created_ms is not None: + if self.after_ms is not None and created_ms < self.after_ms: + continue + if self.before_ms is not None and created_ms > self.before_ms: + continue + + run_dir = runs_root / run.id + if is_run_exported(run_dir): + logger.info(f'{tag}: {run.id} already exported, skipping') + skipped += 1 + continue + + try: + self._export_run(run, run_dir) + exported += 1 + logger.info(f'{tag}: exported {run.id} ({run.name})') + except Exception as e: # keep going: one bad run must not stop all + logger.error(f'{tag}: export failed for {run.id}: {e}') + failed.append({'run_id': run.id, 'error': f'{type(e).__name__}: {e}'}) + + summary = {'exported': exported, 'skipped': skipped, 'failed': failed} + write_json_atomic( + self.output_dir / MANIFEST_FILENAME, + { + 'source': 'wandb', + 'project': self.project_path, + **summary, + }, + ) + return summary + + 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) + + 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) + + def _write_run_json( + self, run: Any, tmp_dir: Path, created_ms: Optional[int] + ) -> None: + summary_dict = getattr(getattr(run, 'summary', None), '_json_dict', None) or {} + summary_dict = {k: v for k, v in summary_dict.items() if not k.startswith('_')} + updated_ms = parse_iso_ms(getattr(run, 'heartbeat_at', None)) or created_ms + write_json_atomic( + tmp_dir / 'run.json', + { + 'entity': self.entity, + 'project': self.project, + 'run_id': run.id, + 'name': run.name, + 'notes': getattr(run, 'notes', None), + 'tags': list(getattr(run, 'tags', []) or []), + 'state': getattr(run, 'state', None), + 'config': dict(getattr(run, 'config', {}) or {}), + 'summary': summary_dict, + 'createdAt': created_ms, + 'updatedAt': updated_ms, + 'url': getattr(run, 'url', None), + 'metadata': getattr(run, 'metadata', None), + }, + ) + + def _row_base(self, run: Any) -> '_RowBase': + return {'project_id': self.project_path, 'run_id': run.id} + + def _export_history(self, run: Any, writer: PartWriter) -> None: + for row in run.scan_history(page_size=self.history_page_size): + step = row.get('_step') + ts = row.get('_timestamp') + if step is None or ts is None: + continue + timestamp_ms = int(float(ts) * 1000) + for key, value in row.items(): + if key.startswith('_'): + continue + self._export_history_value( + run, writer, key, value, int(step), timestamp_ms + ) + + def _export_history_value( + self, + run: Any, + writer: PartWriter, + key: str, + value: Any, + step: int, + timestamp_ms: int, + ) -> None: + base = self._row_base(run) + if isinstance(value, bool) or value is None: + return + if isinstance(value, (int, float)): + writer.write_row( + **base, + attribute_path=key, + attribute_type='metric', + step=step, + timestamp_ms=timestamp_ms, + float_value=float(value), + ) + return + if isinstance(value, dict): + media_type = value.get('_type') + if media_type in _FILE_MEDIA_TYPES and value.get('path'): + writer.write_row( + **base, + attribute_path=key, + attribute_type='media', + step=step, + timestamp_ms=timestamp_ms, + string_value=media_type, + file_value=f'files/{value["path"]}', + caption=value.get('caption'), + ) + elif media_type == 'histogram': + writer.write_row( + **base, + attribute_path=key, + attribute_type='media', + step=step, + timestamp_ms=timestamp_ms, + string_value=json.dumps( + { + '_type': 'histogram', + 'values': value.get('values'), + 'bins': value.get('bins'), + } + ), + ) + elif media_type == 'images/separated' and value.get('filenames'): + captions = value.get('captions') or [] + for i, filename in enumerate(value['filenames']): + writer.write_row( + **base, + attribute_path=key, + attribute_type='media', + step=step, + timestamp_ms=timestamp_ms, + string_value='image-file', + file_value=f'files/{filename}', + caption=captions[i] if i < len(captions) else None, + ) + else: + logger.debug( + f'{tag}: skipping unsupported history value ' + f'{key!r} (_type={media_type!r})' + ) + + def _export_system_metrics(self, run: Any, writer: PartWriter) -> None: + base = self._row_base(run) + try: + events = run.history(stream='events', pandas=False) + except Exception as e: + logger.warning(f'{tag}: system metrics unavailable for {run.id}: {e}') + return + for index, row in enumerate(events): + ts = row.get('_timestamp') + if ts is None: + continue + timestamp_ms = int(float(ts) * 1000) + for key, value in row.items(): + if not key.startswith('system.'): + continue + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + # Source-native name; the loader translates to Pluto's sys/ + # namespace, so staged exports stay platform-agnostic. + writer.write_row( + **base, + attribute_path=key, + attribute_type='system_metric', + step=index, + timestamp_ms=timestamp_ms, + float_value=float(value), + ) + + def _download_files(self, run: Any, files_dir: Path) -> None: + files_dir.mkdir(parents=True, exist_ok=True) + for f in run.files(): + if not self.include_files and f.name != 'output.log': + continue + try: + f.download(root=str(files_dir), exist_ok=True) + except Exception as e: + logger.warning(f'{tag}: failed to download {f.name}: {e}') + + def _export_console( + self, + run: Any, + writer: PartWriter, + files_dir: Path, + created_ms: Optional[int], + ) -> None: + output_log = files_dir / 'output.log' + if not output_log.exists(): + return + base = self._row_base(run) + fallback_ms = created_ms or 0 + with open(output_log, errors='replace') as f: + for line_number, raw_line in enumerate(f, start=1): + message = raw_line.rstrip('\n') + if not message.strip(): + continue + timestamp_ms = self._parse_console_line_time(message, fallback_ms) + writer.write_row( + **base, + attribute_path='console', + attribute_type='console', + step=line_number, + timestamp_ms=timestamp_ms, + string_value=message, + ) + + @staticmethod + def _parse_console_line_time(message: str, fallback_ms: int) -> int: + """Best-effort per-line timestamp; the message itself is never altered + (a leading ISO prefix may be the user's own logging format).""" + match = _CONSOLE_TS_RE.match(message) + if match: + parsed = parse_iso_ms(match.group(1)) + if parsed is not None: + return parsed + return fallback_ms + + def _export_artifacts(self, run: Any, writer: PartWriter, tmp_dir: Path) -> None: + base = self._row_base(run) + try: + artifacts: Iterable[Any] = run.logged_artifacts() + except Exception as e: + logger.warning(f'{tag}: artifacts unavailable for {run.id}: {e}') + return + for artifact in artifacts: + size = getattr(artifact, 'size', None) + if ( + self.artifact_max_bytes is not None + and size is not None + and size > self.artifact_max_bytes + ): + logger.info( + f'{tag}: skipping artifact {artifact.name} ' + f'({size} bytes > cap {self.artifact_max_bytes})' + ) + continue + dest = tmp_dir / 'artifacts' / artifact.name + try: + artifact.download(root=str(dest)) + except Exception as e: + logger.warning(f'{tag}: failed to download {artifact.name}: {e}') + continue + timestamp_ms = parse_iso_ms(getattr(artifact, 'created_at', None)) or 0 + meta = json.dumps( + { + 'name': artifact.name, + 'type': getattr(artifact, 'type', None), + 'size': size, + } + ) + for path in sorted(p for p in dest.rglob('*') if p.is_file()): + writer.write_row( + **base, + attribute_path=artifact.name, + attribute_type='artifact', + step=0, + timestamp_ms=timestamp_ms, + string_value=meta, + file_value=str(path.relative_to(tmp_dir)), + ) diff --git a/pluto/op.py b/pluto/op.py index b37870c..f76e701 100644 --- a/pluto/op.py +++ b/pluto/op.py @@ -1,6 +1,7 @@ import atexit import builtins import logging +import math import os import queue import signal @@ -210,7 +211,16 @@ def _unregister_excepthook(): LoggedNumbers = Dict[str, Any] LoggedData = Dict[str, List[Data]] LoggedFiles = Dict[str, List[File]] -QueueItem = Tuple[Dict[str, Any], Optional[int]] +QueueItem = Tuple[Dict[str, Any], Optional[int], Optional[float]] + + +class RunExistsError(RuntimeError): + """A run with the given externalId already exists (init without resume). + + Typed so callers that intentionally reuse external ids (e.g. + pluto.migrate re-loading after a crash) can catch the collision and + retry with ``resume=True`` instead of string-matching the message. + """ class OpMonitor: @@ -257,16 +267,19 @@ def stop(self, code: Union[int, None] = None) -> None: def _worker_monitor(self, stop): while not stop(): try: - # Collect system metrics - sys_metrics = make_compat_monitor_v1(self.op.settings._sys.monitor()) - timestamp_ms = int(time.time() * 1000) - - # Send system metrics via sync process if enabled - if self.op._sync_manager is not None: - self.op._sync_manager.enqueue_system_metrics( - metrics=sys_metrics, - timestamp_ms=timestamp_ms, + if not self.op.settings.disable_system_metrics: + # Collect system metrics + sys_metrics = make_compat_monitor_v1( + self.op.settings._sys.monitor() ) + timestamp_ms = int(time.time() * 1000) + + # Send system metrics via sync process if enabled + if self.op._sync_manager is not None: + self.op._sync_manager.enqueue_system_metrics( + metrics=sys_metrics, + timestamp_ms=timestamp_ms, + ) # Send heartbeat/trigger to server # Use short timeout and no retries: if it fails, the next @@ -338,7 +351,7 @@ def __init__(self, config, settings, tags=None, resume=False) -> None: make_compat_start_v1( self.config, self.settings, - self.settings._sys.get_info(), + self._start_info(), self.tags, ), client=tmp_iface.client_api, @@ -370,7 +383,7 @@ def __init__(self, config, settings, tags=None, resume=False) -> None: ) else: external_id = self.settings._external_id - raise RuntimeError( + raise RunExistsError( f"Run with externalId '{external_id}' already exists. " f'This often happens when random.seed() or ' f'L.seed_everything() makes run IDs deterministic. ' @@ -427,6 +440,16 @@ def __init__(self, config, settings, tags=None, resume=False) -> None: self._dropped_item_warned: set = set() atexit.register(self.finish) + def _start_info(self) -> Dict[str, Any]: + """System info for the run-create payload. + + Suppressed under disable_system_metrics so a backfill host's + hardware isn't recorded as the imported run's systemMetadata. + """ + if self.settings.disable_system_metrics: + return {} + return self.settings._sys.get_info() + def _init_sync_manager(self) -> None: """Initialize the sync process manager.""" # Generate run_id for sync process @@ -459,6 +482,7 @@ def _init_sync_manager(self) -> None: 'sync_process_retry_backoff': self.settings.sync_process_retry_backoff, 'sync_process_batch_size': self.settings.sync_process_batch_size, 'sync_process_file_batch_size': self.settings.sync_process_file_batch_size, + 'disable_system_metrics': self.settings.disable_system_metrics, } self._sync_manager = SyncProcessManager( @@ -532,7 +556,7 @@ def start(self) -> None: self._monitor.start() # Register system metric names with server (required for dashboard display) - if self._iface: + if self._iface and not self.settings.disable_system_metrics: sys_metric_names = list( make_compat_monitor_v1(self.settings._sys.monitor()).keys() ) @@ -558,12 +582,20 @@ def log( data: Dict[str, Any], step: Union[int, None] = None, commit: Union[bool, None] = None, + timestamp: Optional[float] = None, ) -> None: - """Log run data""" + """Log run data. + + ``timestamp`` is the wall-clock time of the data points in epoch + seconds (``time.time()`` style) and defaults to now. The server + stores it as-is, so backfill/migration tooling can preserve + historical times. Invalid values fall back to now with a warning. + """ + timestamp = self._validate_timestamp(timestamp) # Use sync process if enabled (default: uploads data to server) if self._sync_manager is not None: try: - self._log_via_sync(data=data, step=step) + self._log_via_sync(data=data, step=step, timestamp=timestamp) except sqlite3.OperationalError as e: # Never let a transient SQLite error crash the user's training. # The data for this step is lost, but training continues. @@ -571,23 +603,40 @@ def log( '%s: dropping log data due to database error: %s', tag, e ) elif self.settings.mode == 'perf': - self._queue.put((data, step), block=False) + self._queue.put((data, step, timestamp), block=False) else: # Legacy offline mode (sync_process_enabled=False) # Data stored locally in SQLite only, not uploaded to server - self._log(data=data, step=step) + self._log(data=data, step=step, t=timestamp) + + def _validate_timestamp(self, timestamp: Optional[float]) -> Optional[float]: + """Return a usable explicit timestamp or None (meaning "now").""" + if timestamp is None: + return None + if ( + not isinstance(timestamp, (int, float)) + or isinstance(timestamp, bool) + or not math.isfinite(timestamp) + or timestamp <= 0 + ): + logger.warning( + f'{tag}: ignoring invalid timestamp {timestamp!r}; using current time' + ) + return None + return float(timestamp) def _log_via_sync( self, data: Dict[str, Any], step: Optional[int] = None, + timestamp: Optional[float] = None, ) -> None: """Log data via sync process (writes to SQLite, picked up by sync).""" if self._sync_manager is None: return self._step = self._step + 1 if step is None else step - timestamp_ms = int(time.time() * 1000) + timestamp_ms = int((timestamp if timestamp is not None else time.time()) * 1000) metrics: Dict[str, Any] = {} new_metric_names: List[str] = [] @@ -619,6 +668,51 @@ def _log_via_sync( if (new_metric_names or new_file_meta) and self._iface: self._iface._update_meta(num=new_metric_names, df=dict(new_file_meta)) + def _log_console(self, lines: List[Tuple[str, str, float, int]]) -> None: + """Enqueue console lines with explicit timestamps (backfill path). + + ``lines`` are ``(message, log_type, timestamp_seconds, line_number)`` + tuples with ``log_type`` in ``{'INFO', 'ERROR'}``. Used by + ``pluto.migrate`` loaders to replay another platform's console + output with the original wall-clock times; live console capture + goes through ``pluto.log.ConsoleHandler`` instead. + """ + if self._sync_manager is None: + return + self._sync_manager.enqueue_console_batch( + [ + (message, log_type, int(ts * 1000), line_number) + for message, log_type, ts, line_number in lines + ] + ) + + def _log_metrics_batch( + self, groups: List[Tuple[Dict[str, Any], int, float]] + ) -> None: + """Enqueue many numeric metric groups in one transaction (backfill path). + + ``groups`` are ``(metrics, step, timestamp_seconds)`` tuples with + numeric values only. Used by ``pluto.migrate`` loaders, where one + SQLite transaction per step would dominate replay time; live + training goes through :meth:`log`. + """ + if self._sync_manager is None or not groups: + return + new_metric_names: List[str] = [] + new_file_meta: Dict[str, List[str]] = defaultdict(list) + items: List[Tuple[Dict[str, Any], int, int]] = [] + for metrics, step, timestamp in groups: + clean: Dict[str, Any] = {} + for key, value in metrics.items(): + key = get_char(key) + self._register_meta_sync(key, value, new_metric_names, new_file_meta) + clean[key] = value + self._step = step + items.append((clean, int(timestamp * 1000), step)) + self._sync_manager.enqueue_metrics_batch(items) + if new_metric_names and self._iface: + self._iface._update_meta(num=new_metric_names) + def _warn_dropped_item(self, key: str, value: Any, exc: Exception) -> None: """Report a log item dropped due to an unexpected error. diff --git a/pluto/sets.py b/pluto/sets.py index 916fd81..1b39e13 100644 --- a/pluto/sets.py +++ b/pluto/sets.py @@ -41,6 +41,10 @@ class Settings: disable_iface: bool = False disable_progress: bool = True disable_console: bool = False # disable file-based logging + # Skip sampling/uploading this host's CPU/GPU stats. Used by backfill + # tooling (pluto.migrate) so the importing machine's hardware doesn't + # show up in migrated runs. + disable_system_metrics: bool = False sanitize_logs: bool = True # redact secrets before uploading console logs _op_name: Optional[str] = None diff --git a/pluto/sync/process.py b/pluto/sync/process.py index dfa8c4b..b1263b8 100644 --- a/pluto/sync/process.py +++ b/pluto/sync/process.py @@ -22,7 +22,7 @@ import time from enum import Enum from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Tuple, Union try: from filelock import FileLock @@ -226,6 +226,26 @@ def enqueue_metrics( # Update heartbeat to show we're alive self.store.heartbeat(self.run_id) + def enqueue_metrics_batch( + self, + items: List[Tuple[Dict[str, Any], int, int]], + ) -> None: + """Enqueue many metric groups in one SQLite transaction. + + ``items`` are ``(metrics, timestamp_ms, step)`` tuples. Used by + backfill tooling (pluto.migrate) where per-group ``enqueue_metrics`` + transactions would dominate replay time. + """ + if not items: + return + self.store.enqueue_batch( + [ + (self.run_id, RecordType.METRIC, metrics, timestamp_ms, step) + for metrics, timestamp_ms, step in items + ] + ) + self.store.heartbeat(self.run_id) + def enqueue_config(self, config: Dict[str, Any], timestamp_ms: int) -> None: """Enqueue config update for upload.""" self.store.enqueue( @@ -1149,6 +1169,10 @@ def upload_health_stats(self, stats: Dict[str, Any]) -> None: """ if not self.url_num: return + if self.settings.get('disable_system_metrics'): + # Backfill runs (pluto.migrate) must not receive current-time + # sync-health datapoints from the migration host. + return data = {f'sys/pluto.{k}': v for k, v in stats.items()} timestamp_ms = int(time.time() * 1000) diff --git a/poetry.lock b/poetry.lock index b32647f..a52b3c4 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,17 @@ -# This file is automatically @generated by Poetry 2.3.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. + +[[package]] +name = "annotated-types" +version = "0.7.0" +description = "Reusable constraint types to use with typing.Annotated" +optional = true +python-versions = ">=3.8" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53"}, + {file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"}, +] [[package]] name = "antlr4-python3-runtime" @@ -163,7 +176,7 @@ version = "3.4.4" description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." optional = false python-versions = ">=3.7" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "charset_normalizer-3.4.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e824f1492727fa856dd6eda4f7cee25f8518a12f3c4a56a74e8095695089cf6d"}, {file = "charset_normalizer-3.4.4-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4bd5d4137d500351a30687c2d3971758aac9a19208fc110ccb9d7188fbe709e8"}, @@ -279,6 +292,23 @@ files = [ {file = "charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f"}, {file = "charset_normalizer-3.4.4.tar.gz", hash = "sha256:94537985111c35f28720e43603b8e7b43a6ecfb2ce1d3058bbe955b73404e21a"}, ] +markers = {main = "extra == \"migrate\""} + +[[package]] +name = "click" +version = "8.4.2" +description = "Composable command line interface toolkit" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76"}, + {file = "click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} [[package]] name = "colorama" @@ -286,11 +316,12 @@ version = "0.4.6" description = "Cross-platform colored terminal text." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] +markers = {main = "extra == \"migrate\" and platform_system == \"Windows\""} [[package]] name = "contourpy" @@ -1941,11 +1972,12 @@ version = "25.0" description = "Core utilities for Python packages" optional = false python-versions = ">=3.8" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484"}, {file = "packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f"}, ] +markers = {main = "extra == \"migrate\""} [[package]] name = "pathspec" @@ -2068,6 +2100,19 @@ test-arrow = ["arro3-compute", "arro3-core", "nanoarrow", "pyarrow"] tests = ["check-manifest", "coverage (>=7.4.2)", "defusedxml", "markdown2", "olefile", "packaging", "pyroma (>=5)", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "trove-classifiers (>=2024.10.12)"] xmp = ["defusedxml"] +[[package]] +name = "platformdirs" +version = "4.10.0" +description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`." +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a"}, + {file = "platformdirs-4.10.0.tar.gz", hash = "sha256:31e761a6a0ca04faf7353ea759bdba55652be214725111e5aac52dfa29d4bef7"}, +] + [[package]] name = "pluggy" version = "1.6.0" @@ -2099,6 +2144,25 @@ files = [ [package.dependencies] tqdm = "*" +[[package]] +name = "protobuf" +version = "7.35.1" +description = "" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6"}, + {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799"}, + {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4"}, + {file = "protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4"}, + {file = "protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30"}, + {file = "protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87"}, + {file = "protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9"}, + {file = "protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a"}, +] + [[package]] name = "psutil" version = "7.1.3" @@ -2132,6 +2196,67 @@ files = [ dev = ["abi3audit", "black", "check-manifest", "colorama ; os_name == \"nt\"", "coverage", "packaging", "pylint", "pyperf", "pypinfo", "pyreadline ; os_name == \"nt\"", "pytest", "pytest-cov", "pytest-instafail", "pytest-subtests", "pytest-xdist", "pywin32 ; os_name == \"nt\" and platform_python_implementation != \"PyPy\"", "requests", "rstcheck", "ruff", "setuptools", "sphinx", "sphinx_rtd_theme", "toml-sort", "twine", "validate-pyproject[all]", "virtualenv", "vulture", "wheel", "wheel ; os_name == \"nt\" and platform_python_implementation != \"PyPy\"", "wmi ; os_name == \"nt\" and platform_python_implementation != \"PyPy\""] test = ["pytest", "pytest-instafail", "pytest-subtests", "pytest-xdist", "pywin32 ; os_name == \"nt\" and platform_python_implementation != \"PyPy\"", "setuptools", "wheel ; os_name == \"nt\" and platform_python_implementation != \"PyPy\"", "wmi ; os_name == \"nt\" and platform_python_implementation != \"PyPy\""] +[[package]] +name = "pyarrow" +version = "24.0.0" +description = "Python library for Apache Arrow" +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +files = [ + {file = "pyarrow-24.0.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:7c2b98645d576a0b9616892ead22b64a83a5f043c5e2ca15ebcefcb5b70c80cb"}, + {file = "pyarrow-24.0.0-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:644a246325b8c69c595ad1dd4b463eba4b0cdb731370e4a86137d433208d6147"}, + {file = "pyarrow-24.0.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:3a577bd840ca83f646f0a625dbc571dba7044c43c2d1503afc378b570954345c"}, + {file = "pyarrow-24.0.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:e3268e43984d0b1a185c89b4cfff282a7ead12fc93f56cfd7088bdbcbe727041"}, + {file = "pyarrow-24.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:2392d954fcb920f42d230284b677605e4e2fbb11f2821e823e642abd67fbb491"}, + {file = "pyarrow-24.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:bec9373df11544592b0ba7ec2af0e35059e5f0e7647c6183a854dedd193298f1"}, + {file = "pyarrow-24.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:c42ab9439498270139cc63e18847a02afe5c8b3ed9c931266533cfe378bd3591"}, + {file = "pyarrow-24.0.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:b0e131f880cda8d04e076cee175a46fc0e8bc8b65c99c6c09dff6669335fde74"}, + {file = "pyarrow-24.0.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:1b2fe7f9a5566401a0ef2571f197eb92358925c1f0c8dba305d6e43ea0871bb3"}, + {file = "pyarrow-24.0.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:0b3537c00fb8d384f15ac1e79b6eb6db04a16514c8c1d22e59a9b95c8ba42868"}, + {file = "pyarrow-24.0.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:14e31a3c9e35f1ab6356c6378f6f72830e6d2d5f1791df3774a7b097d18a6a1e"}, + {file = "pyarrow-24.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b7d9a514e73bc42711e6a35aaccf3587c520024fe0a25d830a1a8a27c15f4f57"}, + {file = "pyarrow-24.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:b196eb3f931862af3fa84c2a253514d859c08e0d8fe020e07be12e75a5a9780c"}, + {file = "pyarrow-24.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:35405aecb474e683fb36af650618fd5340ee5471fc65a21b36076a18bbc6c981"}, + {file = "pyarrow-24.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6233c9ed9ab9d1db47de57d9753256d9dcffbf42db341576099f0fd9f6bf4810"}, + {file = "pyarrow-24.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:f7616236ec1bc2b15bfdec22a71ab38851c86f8f05ff64f379e1278cf20c634a"}, + {file = "pyarrow-24.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1617043b99bd33e5318ae18eb2919af09c71322ef1ca46566cdafc6e6712fb66"}, + {file = "pyarrow-24.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:6165461f55ef6314f026de6638d661188e3455d3ec49834556a0ebbdbace18bb"}, + {file = "pyarrow-24.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b13dedfe76a0ad2d1d859b0811b53827a4e9d93a0bcb05cf59333ab4980cc7e"}, + {file = "pyarrow-24.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:25ea65d868eb04015cd18e6df2fbe98f07e5bda2abefabcb88fce39a947716f6"}, + {file = "pyarrow-24.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:295f0a7f2e242dabd513737cf076007dc5b2d59237e3eca37b05c0c6446f3826"}, + {file = "pyarrow-24.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:02b001b3ed4723caa44f6cd1af2d5c86aa2cf9971dacc2ffa55b21237713dfba"}, + {file = "pyarrow-24.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:04920d6a71aabd08a0417709efce97d45ea8e6fb733d9ca9ecffb13c67839f68"}, + {file = "pyarrow-24.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a964266397740257f16f7bb2e4f08a0c81454004beab8ff59dd531b73610e9f2"}, + {file = "pyarrow-24.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6f066b179d68c413374294bc1735f68475457c933258df594443bb9d88ddc2a0"}, + {file = "pyarrow-24.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1183baeb14c5f587b1ec52831e665718ce632caab84b7cd6b85fd44f96114495"}, + {file = "pyarrow-24.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:806f24b4085453c197a5078218d1ee08783ebbba271badd153d1ae22a3ee804f"}, + {file = "pyarrow-24.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:e4505fc6583f7b05ab854934896bcac8253b04ac1171a77dfb73efef92076d91"}, + {file = "pyarrow-24.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:1a4e45017efbf115032e4475ee876d525e0e36c742214fbe405332480ecd6275"}, + {file = "pyarrow-24.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:7986f1fa71cee060ad00758bcc79d3a93bab8559bf978fab9e53472a2e25a17b"}, + {file = "pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d3e0b61e8efb24ed38898e5cdc5fffa9124be480008d401a1f8071500494ae42"}, + {file = "pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:55a3bc1e3df3b5567b7d27ef551b2283f0c68a5e86f1cd56abc569da4f31335b"}, + {file = "pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:641f795b361874ac9da5294f8f443dfdbee355cf2bd9e3b8d97aaac2306b9b37"}, + {file = "pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8adc8e6ce5fccf5dc707046ae4914fd537def529709cc0d285d37a7f9cd442ca"}, + {file = "pyarrow-24.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:9b18371ad2f44044b81a8d23bc2d8a9b6a6226dca775e8e16cfee640473d6c5d"}, + {file = "pyarrow-24.0.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:1cc9057f0319e26333b357e17f3c2c022f1a83739b48a88b25bfd5fa2dc18838"}, + {file = "pyarrow-24.0.0-cp314-cp314-macosx_12_0_x86_64.whl", hash = "sha256:e6f1278ee4785b6db21229374a1c9e54ec7c549de5d1efc9630b6207de7e170b"}, + {file = "pyarrow-24.0.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:adbbedc55506cbdabb830890444fb856bfb0060c46c6f8026c6c2f2cf86ae795"}, + {file = "pyarrow-24.0.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:ae8a1145af31d903fa9bb166824d7abe9b4681a000b0159c9fb99c11bc11ad26"}, + {file = "pyarrow-24.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:d7027eba1df3b2069e2e8d80f644fa0918b68c46432af3d088ddd390d063ecde"}, + {file = "pyarrow-24.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e56a1ffe9bf7b727432b89104cc0849c21582949dd7bdcb34f17b2001a351a76"}, + {file = "pyarrow-24.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:38be1808cdd068605b787e6ca9119b27eb275a0234e50212c3492331680c3b1e"}, + {file = "pyarrow-24.0.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:418e48ce50a45a6a6c73c454677203a9c75c966cb1e92ca3370959185f197a05"}, + {file = "pyarrow-24.0.0-cp314-cp314t-macosx_12_0_x86_64.whl", hash = "sha256:2f16197705a230a78270cdd4ea8a1d57e86b2fdcbc34a1f6aebc72e65c986f9a"}, + {file = "pyarrow-24.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:fb24ac194bfc5e86839d7dcd52092ee31e5fe6733fe11f5e3b06ef0812b20072"}, + {file = "pyarrow-24.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:9700ebd9a51f5895ce75ff4ac4b3c47a7d4b42bc618be8e713e5d56bacf5f931"}, + {file = "pyarrow-24.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:d8ddd2768da81d3ee08cfea9b597f4abb4e8e1dc8ae7e204b608d23a0d3ab699"}, + {file = "pyarrow-24.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:61a3d7eaa97a14768b542f3d284dc6400dd2470d9f080708b13cd46b6ae18136"}, + {file = "pyarrow-24.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c91d00057f23b8d353039520dc3a6c09d8608164c692e9f59a175a42b2ae0c19"}, + {file = "pyarrow-24.0.0.tar.gz", hash = "sha256:85fe721a14dd823aca09127acbb06c3ca723efbd436c004f16bca601b04dcc83"}, +] +markers = {main = "extra == \"migrate\""} + [[package]] name = "pycparser" version = "2.23" @@ -2145,6 +2270,163 @@ files = [ {file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"}, ] +[[package]] +name = "pydantic" +version = "2.13.4" +description = "Data validation using Python type hints" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba"}, + {file = "pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6"}, +] + +[package.dependencies] +annotated-types = ">=0.6.0" +pydantic-core = "2.46.4" +typing-extensions = ">=4.14.1" +typing-inspection = ">=0.4.2" + +[package.extras] +email = ["email-validator (>=2.0.0)"] +timezone = ["tzdata ; python_version >= \"3.9\" and platform_system == \"Windows\""] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +description = "Core functionality for Pydantic validation and serialization" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "pydantic_core-2.46.4-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4"}, + {file = "pydantic_core-2.46.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_31_riscv64.whl", hash = "sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01"}, + {file = "pydantic_core-2.46.4-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d"}, + {file = "pydantic_core-2.46.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4"}, + {file = "pydantic_core-2.46.4-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f"}, + {file = "pydantic_core-2.46.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39"}, + {file = "pydantic_core-2.46.4-cp310-cp310-win32.whl", hash = "sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d"}, + {file = "pydantic_core-2.46.4-cp310-cp310-win_amd64.whl", hash = "sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf"}, + {file = "pydantic_core-2.46.4-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594"}, + {file = "pydantic_core-2.46.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398"}, + {file = "pydantic_core-2.46.4-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3"}, + {file = "pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848"}, + {file = "pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3"}, + {file = "pydantic_core-2.46.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109"}, + {file = "pydantic_core-2.46.4-cp311-cp311-win32.whl", hash = "sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda"}, + {file = "pydantic_core-2.46.4-cp311-cp311-win_amd64.whl", hash = "sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33"}, + {file = "pydantic_core-2.46.4-cp311-cp311-win_arm64.whl", hash = "sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d"}, + {file = "pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2"}, + {file = "pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987"}, + {file = "pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b"}, + {file = "pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458"}, + {file = "pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b"}, + {file = "pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c"}, + {file = "pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894"}, + {file = "pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89"}, + {file = "pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a"}, + {file = "pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008"}, + {file = "pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262"}, + {file = "pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e"}, + {file = "pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd"}, + {file = "pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be"}, + {file = "pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d"}, + {file = "pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb"}, + {file = "pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292"}, + {file = "pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d"}, + {file = "pydantic_core-2.46.4-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb"}, + {file = "pydantic_core-2.46.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c"}, + {file = "pydantic_core-2.46.4-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb"}, + {file = "pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898"}, + {file = "pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e"}, + {file = "pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519"}, + {file = "pydantic_core-2.46.4-cp314-cp314-win32.whl", hash = "sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4"}, + {file = "pydantic_core-2.46.4-cp314-cp314-win_amd64.whl", hash = "sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac"}, + {file = "pydantic_core-2.46.4-cp314-cp314-win_arm64.whl", hash = "sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-win32.whl", hash = "sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-win_amd64.whl", hash = "sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5"}, + {file = "pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596"}, + {file = "pydantic_core-2.46.4-cp39-cp39-macosx_10_12_x86_64.whl", hash = "sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae"}, + {file = "pydantic_core-2.46.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_31_riscv64.whl", hash = "sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066"}, + {file = "pydantic_core-2.46.4-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6"}, + {file = "pydantic_core-2.46.4-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9"}, + {file = "pydantic_core-2.46.4-cp39-cp39-musllinux_1_1_armv7l.whl", hash = "sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29"}, + {file = "pydantic_core-2.46.4-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9"}, + {file = "pydantic_core-2.46.4-cp39-cp39-win32.whl", hash = "sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1"}, + {file = "pydantic_core-2.46.4-cp39-cp39-win_amd64.whl", hash = "sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac"}, + {file = "pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c"}, + {file = "pydantic_core-2.46.4-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b"}, + {file = "pydantic_core-2.46.4-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b"}, + {file = "pydantic_core-2.46.4-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea"}, + {file = "pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7"}, + {file = "pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df"}, + {file = "pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526"}, + {file = "pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc"}, + {file = "pydantic_core-2.46.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983"}, + {file = "pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1"}, +] + +[package.dependencies] +typing-extensions = ">=4.14.1" + [[package]] name = "pygments" version = "2.19.2" @@ -2254,7 +2536,7 @@ version = "6.0.3" description = "YAML parser and emitter for Python" optional = false python-versions = ">=3.8" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "PyYAML-6.0.3-cp38-cp38-macosx_10_13_x86_64.whl", hash = "sha256:c2514fceb77bc5e7a2f7adfaa1feb2fb311607c9cb518dbc378688ec73d8292f"}, {file = "PyYAML-6.0.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c57bb8c96f6d1808c030b1687b9b5fb476abaa47f0db9c0101f5e9f394e97f4"}, @@ -2330,6 +2612,7 @@ files = [ {file = "pyyaml-6.0.3-cp39-cp39-win_amd64.whl", hash = "sha256:2e71d11abed7344e42a8849600193d15b6def118602c4c176f748e4583246007"}, {file = "pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f"}, ] +markers = {main = "extra == \"migrate\""} [[package]] name = "requests" @@ -2337,11 +2620,12 @@ version = "2.32.5" description = "Python HTTP for Humans." optional = false python-versions = ">=3.9" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6"}, {file = "requests-2.32.5.tar.gz", hash = "sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf"}, ] +markers = {main = "extra == \"migrate\""} [package.dependencies] certifi = ">=2017.4.17" @@ -2833,7 +3117,23 @@ files = [ {file = "typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548"}, {file = "typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466"}, ] -markers = {main = "python_version < \"3.13\""} +markers = {main = "python_version < \"3.13\" or extra == \"migrate\""} + +[[package]] +name = "typing-inspection" +version = "0.4.2" +description = "Runtime typing introspection tools" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7"}, + {file = "typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464"}, +] + +[package.dependencies] +typing-extensions = ">=4.12.0" [[package]] name = "urllib3" @@ -2853,6 +3153,53 @@ h2 = ["h2 (>=4,<5)"] socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] zstd = ["backports-zstd (>=1.0.0) ; python_version < \"3.14\""] +[[package]] +name = "wandb" +version = "0.28.0" +description = "A CLI and library for interacting with the Weights & Biases API." +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"migrate\"" +files = [ + {file = "wandb-0.28.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:c3dab1205a5aca4abbad1eca08902cdba86add0edfa83d8d61b4429d0e79fa87"}, + {file = "wandb-0.28.0-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:ae255da18726ee8e731ef82cbc85035b901a28ae14cf91604c361b44b8d44ce0"}, + {file = "wandb-0.28.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:6dbcba12ab168aa37561f2f32dcdef8713495fc25fa7d30fdc9bfb37989694dd"}, + {file = "wandb-0.28.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:325b2d0bd88be6eda5db10542499bad3710927f2569c81a84dc5eeaffc76825c"}, + {file = "wandb-0.28.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:8954bc1c62ae43914dce2bebfd1d9957f72350f8fbb78e5cdfe2ca9b6be8a7b8"}, + {file = "wandb-0.28.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:9fec6c908554c2dad33110c1312bc3028cc2e430f0679f16b84f82c8ea801e3b"}, + {file = "wandb-0.28.0-py3-none-win32.whl", hash = "sha256:8834ef3a7c8c43b701654162783caa7ad37af48a0ff06fc35d0d65a411f76ccd"}, + {file = "wandb-0.28.0-py3-none-win_amd64.whl", hash = "sha256:ac1f82292e2da4f98297b78c3a46726b3a6c5734ecb75fc39b8db2c8a4989159"}, + {file = "wandb-0.28.0-py3-none-win_arm64.whl", hash = "sha256:c5b0faf1b84cf79ebabed77538c1940a4c6053e815f767a4004e877a1354bed1"}, + {file = "wandb-0.28.0.tar.gz", hash = "sha256:b20e5af0fe80e2e2a466b0466a1d60cedcc578dce0f036eca04f4a0adcad95b6"}, +] + +[package.dependencies] +click = ">=8.2.0" +gitpython = ">=1.0.0,<3.1.29 || >3.1.29" +packaging = "*" +platformdirs = "*" +protobuf = ">4.21.0,<5.28.0 || >5.28.0,<5.29.0 || >5.29.0,<8" +pydantic = ">=2.6,<3" +pyyaml = "*" +requests = ">=2.0.0,<3" +sentry-sdk = ">=2.0.0" +typing-extensions = ">=4.8,<5" + +[package.extras] +aws = ["boto3", "botocore (>=1.5.76)"] +azure = ["azure-identity", "azure-storage-blob"] +eval-table = ["weave (>=0.52.41)"] +eval-table-video-support = ["weave[video-support] (>=0.52.41)"] +gcp = ["google-cloud-storage"] +kubeflow = ["google-cloud-storage", "kubernetes", "minio", "sh"] +launch = ["awscli", "azure-containerregistry", "azure-identity", "azure-storage-blob", "boto3", "botocore (>=1.5.76)", "chardet", "google-auth", "google-cloud-aiplatform", "google-cloud-artifact-registry", "google-cloud-compute", "google-cloud-storage", "iso8601", "jsonschema", "kubernetes", "kubernetes-asyncio", "nbconvert", "nbformat", "optuna", "pydantic", "pyyaml (>=6.0.0)", "tomli", "tornado (>=6.5.0)", "typing-extensions"] +media = ["bokeh", "imageio (>=2.28.1)", "moviepy (>=1.0.0)", "numpy", "pillow", "plotly (>=5.18.0)", "rdkit", "soundfile"] +models = ["cloudpickle"] +sandbox = ["cwsandbox[cli] (>=0.20.0)"] +sweeps = ["sweeps (>=0.2.0)"] +workspaces = ["wandb-workspaces"] + [[package]] name = "zipp" version = "3.23.0" @@ -2876,8 +3223,9 @@ type = ["pytest-mypy"] [extras] full = ["nvidia-ml-py"] +migrate = ["pyarrow", "wandb"] [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "6668aa67f414632c07ec8050bd4eee21670c916759d41644f43e12cdf36680f6" +content-hash = "d4792c7ce74fd4778321d7613fe3a56e12af0b52889468caf57d773e321ca81d" diff --git a/pyproject.toml b/pyproject.toml index e2ae061..5437539 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,9 @@ nvidia-ml-py = ">=11.515.0" rich = ">=13.0" sentry-sdk = ">=2.0" soundfile = ">=0.12" +# migrate extra: wandb -> pluto historical data migration (pluto migrate wandb) +wandb = { version = ">=0.16", optional = true } +pyarrow = { version = ">=15", optional = true } [tool.poetry.group.dev.dependencies] mypy = "^1.10.0" @@ -50,12 +53,17 @@ torchvision = "*" imageio = "^2.37.0" moviepy = "^1.0.3" omegaconf = "^2.3" +pyarrow = ">=15" # pluto.migrate tests [tool.poetry.extras] full = [ "nvidia-ml-py", ] +migrate = [ + "wandb", + "pyarrow", +] [tool.poetry.scripts] pluto = 'pluto.__main__:main' diff --git a/tests/test_auth_transient.py b/tests/test_auth_transient.py new file mode 100644 index 0000000..6cf75f8 --- /dev/null +++ b/tests/test_auth_transient.py @@ -0,0 +1,49 @@ +""" +Regression tests for transient login-validation failures. + +login() posts the token to /api/slug as a best-effort validation. A +TRANSIENT network failure of that single request (timeout, reset — CI +under load) must not corrupt a token that was explicitly provided via +PLUTO_API_KEY/keyring: overwriting it with the '_key' sentinel made +every subsequent request send 'Bearer _key', which fails the server's +prefix check as 401 "Invalid API key" for the run's entire lifetime. +""" + +from __future__ import annotations + +from unittest import mock + +import httpx + +from pluto import auth +from pluto.sets import Settings + +PROVIDED_TOKEN = 'mlpi_env_provided_token' + + +def _settings_with_token(token) -> Settings: + settings = Settings() + settings._auth = token + settings.update_host() + return settings + + +def test_transient_login_failure_keeps_provided_token(): + settings = _settings_with_token(PROVIDED_TOKEN) + with mock.patch.object(auth.httpx, 'Client') as client_cls: + client_cls.return_value.post.side_effect = httpx.ConnectTimeout('boom') + auth.login(settings=settings) + assert settings._auth == PROVIDED_TOKEN + + +def test_unreachable_server_without_token_still_marks_unauthenticated(): + # Preserve the interactive-flow sentinel when no token was provided. + settings = _settings_with_token(None) + with ( + mock.patch.object(auth.httpx, 'Client') as client_cls, + mock.patch.object(auth.keyring, 'get_password', return_value=None), + mock.patch.object(auth.sys, 'stdin', None), # no interactive prompt + ): + client_cls.return_value.post.side_effect = httpx.ConnectTimeout('boom') + auth.login(settings=settings) + assert settings._auth == '_key' diff --git a/tests/test_disable_system_metrics.py b/tests/test_disable_system_metrics.py new file mode 100644 index 0000000..1bb6946 --- /dev/null +++ b/tests/test_disable_system_metrics.py @@ -0,0 +1,115 @@ +""" +Unit tests for settings.disable_system_metrics and Op._log_console. + +Migration tooling (pluto.migrate) replays runs recorded elsewhere: the +migration host's own GPU/CPU stats must not leak into the imported run, +and historical console lines must be enqueueable with their original +timestamps. These tests pin both behaviors. +""" + +from __future__ import annotations + +import os +from unittest import mock + +from pluto.op import Op +from pluto.sets import Settings + +TS = 1600000000.5 +TS_MS = 1600000000500 + + +def _make_op(tmp_path) -> Op: + settings = Settings() + settings.mode = 'noop' + settings.dir = str(tmp_path) + settings.meta = [] # shadow the class-level shared list (test isolation) + os.makedirs(os.path.join(settings.get_dir(), 'files'), exist_ok=True) + op = Op(config={}, settings=settings) + op._sync_manager = mock.MagicMock() + return op + + +def _run_monitor_once(op: Op) -> None: + """Drive exactly one _worker_monitor iteration without sleeping.""" + op._monitor._stop_event.set() # make the end-of-loop wait() return instantly + seq = iter([False, True]) + op._monitor._worker_monitor(lambda: next(seq)) + + +class TestDisableSystemMetrics: + def test_monitor_sends_system_metrics_by_default(self, tmp_path): + op = _make_op(tmp_path) + op.settings._sys = mock.MagicMock() + op.settings._sys.monitor.return_value = {'cpu': 1.0} + _run_monitor_once(op) + op._sync_manager.enqueue_system_metrics.assert_called_once() + + def test_monitor_skips_system_metrics_when_disabled(self, tmp_path): + op = _make_op(tmp_path) + op.settings.disable_system_metrics = True + op.settings._sys = mock.MagicMock() + _run_monitor_once(op) + op._sync_manager.enqueue_system_metrics.assert_not_called() + op.settings._sys.monitor.assert_not_called() + + def test_start_skips_sys_name_registration_when_disabled(self, tmp_path): + op = _make_op(tmp_path) + op.settings.disable_system_metrics = True + op.settings._sys = mock.MagicMock() + op._sync_manager = None + op._monitor = mock.MagicMock() + op._iface = mock.MagicMock() + op.start() + op._iface._update_meta.assert_not_called() + op.settings._sys.monitor.assert_not_called() + + +class TestDisableSystemMetricsPropagation: + def test_flag_forwarded_to_sync_process_settings(self, tmp_path): + op = _make_op(tmp_path) + op.settings.update_host() # populate url_* like setup() does + op.settings.disable_system_metrics = True + op.settings._op_id = 1 + with mock.patch('pluto.op.SyncProcessManager') as spm: + op._init_sync_manager() + settings_dict = spm.call_args.kwargs['settings_dict'] + assert settings_dict['disable_system_metrics'] is True + + def test_health_stats_not_uploaded_when_disabled(self): + from pluto.sync.process import _SyncUploader + + uploader = _SyncUploader( + {'disable_system_metrics': True, 'url_num': 'http://x/ingest/metrics'}, + mock.MagicMock(), + ) + with mock.patch.object(uploader, '_post_with_retry') as post: + uploader.upload_health_stats({'pending': 1}) + post.assert_not_called() + + def test_start_info_suppressed_when_disabled(self, tmp_path): + op = _make_op(tmp_path) + op.settings._sys = mock.MagicMock() + op.settings.disable_system_metrics = True + assert op._start_info() == {} + op.settings._sys.get_info.assert_not_called() + + def test_start_info_default(self, tmp_path): + op = _make_op(tmp_path) + op.settings._sys = mock.MagicMock() + op.settings._sys.get_info.return_value = {'host': 'gpu-box'} + assert op._start_info() == {'host': 'gpu-box'} + + +class TestLogConsole: + def test_log_console_converts_seconds_to_ms(self, tmp_path): + op = _make_op(tmp_path) + op._log_console([('hello', 'INFO', TS, 1), ('oops', 'ERROR', TS + 1, 2)]) + op._sync_manager.enqueue_console_batch.assert_called_once_with( + [('hello', 'INFO', TS_MS, 1), ('oops', 'ERROR', TS_MS + 1000, 2)] + ) + + def test_log_console_is_noop_without_sync_manager(self, tmp_path): + op = _make_op(tmp_path) + op._sync_manager = None + op._log_console([('hello', 'INFO', TS, 1)]) # must not raise diff --git a/tests/test_log_timestamp.py b/tests/test_log_timestamp.py new file mode 100644 index 0000000..f943b61 --- /dev/null +++ b/tests/test_log_timestamp.py @@ -0,0 +1,133 @@ +""" +Unit tests for historical-timestamp support on Op.log (backfill path). + +Migration/backfill tooling (pluto.migrate) must be able to preserve the +original wall-clock time of each data point. These tests pin that an +explicit ``timestamp`` (epoch seconds) passed to ``Op.log`` reaches the +sync layer as ``timestamp_ms`` for metrics, structured data, and files — +and that a missing or invalid timestamp falls back to "now" without +crashing the caller. +""" + +from __future__ import annotations + +import math +import os +import time +from unittest import mock + +import numpy as np +import pytest + +import pluto +from pluto.op import Op +from pluto.sets import Settings + +TS = 1600000000.5 +TS_MS = 1600000000500 + + +def _make_op(tmp_path) -> Op: + settings = Settings() + settings.mode = 'noop' + settings.dir = str(tmp_path) + settings.meta = [] # shadow the class-level shared list (test isolation) + # Op.__init__ only prepares the staging dir outside noop mode. + os.makedirs(os.path.join(settings.get_dir(), 'files'), exist_ok=True) + op = Op(config={}, settings=settings) + op._sync_manager = mock.MagicMock() + return op + + +class TestLogTimestampMetrics: + def test_explicit_timestamp_reaches_enqueue_metrics(self, tmp_path): + op = _make_op(tmp_path) + op.log({'loss': 0.5}, step=5, timestamp=TS) + op._sync_manager.enqueue_metrics.assert_called_once_with( + {'loss': 0.5}, TS_MS, 5 + ) + + def test_default_timestamp_is_now(self, tmp_path): + op = _make_op(tmp_path) + before_ms = int(time.time() * 1000) + op.log({'loss': 0.5}, step=1) + after_ms = int(time.time() * 1000) + (_, timestamp_ms, _), _ = ( + op._sync_manager.enqueue_metrics.call_args.args, + op._sync_manager.enqueue_metrics.call_args.kwargs, + ) + assert before_ms <= timestamp_ms <= after_ms + + @pytest.mark.parametrize('bad', [0, -5, float('nan'), float('inf')]) + def test_invalid_timestamp_falls_back_to_now(self, tmp_path, bad): + op = _make_op(tmp_path) + before_ms = int(time.time() * 1000) + op.log({'loss': 0.5}, step=1, timestamp=bad) + (_, timestamp_ms, _), _ = ( + op._sync_manager.enqueue_metrics.call_args.args, + op._sync_manager.enqueue_metrics.call_args.kwargs, + ) + assert timestamp_ms >= before_ms + assert math.isfinite(timestamp_ms) + + +class TestLogTimestampDataAndFiles: + def test_timestamp_reaches_enqueue_data_for_histogram(self, tmp_path): + op = _make_op(tmp_path) + op.log({'dist': pluto.Histogram(np.arange(10))}, step=3, timestamp=TS) + kwargs = op._sync_manager.enqueue_data.call_args.kwargs + assert kwargs['timestamp_ms'] == TS_MS + assert kwargs['step'] == 3 + + def test_timestamp_reaches_enqueue_file_for_image(self, tmp_path): + op = _make_op(tmp_path) + img = pluto.Image(np.zeros((4, 4, 3), dtype=np.uint8)) + op.log({'sample': img}, step=7, timestamp=TS) + kwargs = op._sync_manager.enqueue_file.call_args.kwargs + assert kwargs['timestamp_ms'] == TS_MS + assert kwargs['step'] == 7 + + +class TestLogMetricsBatch: + def test_batch_enqueues_groups_in_one_call(self, tmp_path): + op = _make_op(tmp_path) + op._log_metrics_batch( + [ + ({'loss': 1.0, 'acc': 0.1}, 0, TS), + ({'loss': 0.5}, 1, TS + 1), + ] + ) + op._sync_manager.enqueue_metrics_batch.assert_called_once_with( + [ + ({'loss': 1.0, 'acc': 0.1}, TS_MS, 0), + ({'loss': 0.5}, TS_MS + 1000, 1), + ] + ) + + def test_batch_registers_new_metric_names(self, tmp_path): + op = _make_op(tmp_path) + op._iface = mock.MagicMock() + op._log_metrics_batch([({'loss': 1.0}, 0, TS)]) + op._iface._update_meta.assert_called_once_with(num=['loss']) + + def test_batch_noop_without_sync_manager(self, tmp_path): + op = _make_op(tmp_path) + op._sync_manager = None + op._log_metrics_batch([({'loss': 1.0}, 0, TS)]) # must not raise + + +class TestLogTimestampLegacyPath: + def test_timestamp_forwarded_to_legacy_log(self, tmp_path): + op = _make_op(tmp_path) + op._sync_manager = None # force legacy offline path + op.settings.mode = 'debug' # not the perf queue + with mock.patch.object(op, '_log') as legacy_log: + op.log({'loss': 0.5}, step=2, timestamp=TS) + legacy_log.assert_called_once_with(data={'loss': 0.5}, step=2, t=TS) + + def test_timestamp_forwarded_in_perf_mode_queue(self, tmp_path): + op = _make_op(tmp_path) + op._sync_manager = None + op.settings.mode = 'perf' + op.log({'loss': 0.5}, step=2, timestamp=TS) + assert op._queue.get_nowait() == ({'loss': 0.5}, 2, TS) diff --git a/tests/test_migrate_cli.py b/tests/test_migrate_cli.py new file mode 100644 index 0000000..58f51ea --- /dev/null +++ b/tests/test_migrate_cli.py @@ -0,0 +1,210 @@ +""" +Unit tests for the `pluto migrate wandb` CLI (pluto.migrate.cli). + +The CLI is thin arg-parsing over WandbExporter/PlutoLoader; both are +mocked here. Heavy deps (wandb/pyarrow) must only be imported inside +command handlers so `pluto --help` stays light — pinned by the +subprocess help test. +""" + +from __future__ import annotations + +import subprocess +import sys +from unittest import mock + +from pluto.migrate.cli import run_migrate + + +def _mock_exporter(summary=None): + exporter = mock.MagicMock() + exporter.export.return_value = summary or { + 'exported': 1, + 'skipped': 0, + 'failed': [], + } + return exporter + + +def _mock_loader(summary=None): + loader = mock.MagicMock() + loader.load.return_value = summary or {'loaded': 1, 'skipped': 0, 'failed': []} + return loader + + +class TestMigrateCli: + def test_export_wires_flags_to_exporter(self, tmp_path): + exporter = _mock_exporter() + with mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ) as cls: + code = run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + '--run-id', + 'r1', + '--run-id', + 'r2', + '--after', + '2025-01-01', + '--no-artifacts', + '--artifact-max-size-mb', + '512', + ] + ) + assert code == 0 + kwargs = cls.call_args.kwargs + assert kwargs['entity'] == 'acme' + assert kwargs['project'] == 'vision' + assert kwargs['run_ids'] == ['r1', 'r2'] + assert kwargs['after'] == '2025-01-01' + assert kwargs['include_artifacts'] is False + assert kwargs['artifact_max_bytes'] == 512 * 1024 * 1024 + exporter.export.assert_called_once() + + def test_load_wires_flags_to_loader(self, tmp_path): + loader = _mock_loader() + with mock.patch('pluto.migrate.loader.PlutoLoader', return_value=loader) as cls: + code = run_migrate( + [ + 'wandb', + 'load', + '--input', + str(tmp_path), + '--dest-project', + 'legacy', + '--dry-run', + ] + ) + assert code == 0 + kwargs = cls.call_args.kwargs + assert kwargs['dest_project'] == 'legacy' + assert kwargs['dry_run'] is True + loader.load.assert_called_once() + + def test_all_exports_then_loads(self, tmp_path): + exporter, loader = _mock_exporter(), _mock_loader() + with ( + mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ), + mock.patch('pluto.migrate.loader.PlutoLoader', return_value=loader), + ): + code = run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + ] + ) + assert code == 0 + exporter.export.assert_called_once() + loader.load.assert_called_once() + + def test_artifact_max_size_zero_means_zero_cap(self, tmp_path): + exporter = _mock_exporter() + with mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ) as cls: + run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + '--artifact-max-size-mb', + '0', + ] + ) + # 0 is an explicit cap (skip everything), not "unlimited" + assert cls.call_args.kwargs['artifact_max_bytes'] == 0 + + def test_all_still_loads_when_some_exports_failed(self, tmp_path): + exporter = _mock_exporter( + {'exported': 499, 'skipped': 0, 'failed': [{'run_id': 'x', 'error': 'e'}]} + ) + loader = _mock_loader() + with ( + mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ), + mock.patch('pluto.migrate.loader.PlutoLoader', return_value=loader), + ): + code = run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + ] + ) + loader.load.assert_called_once() # staged runs still load + assert code == 1 # but the failure is reported + + def test_all_rejects_dry_run(self, tmp_path): + with mock.patch('pluto.migrate.wandb_export.WandbExporter') as cls: + code = run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + '--dry-run', + ] + ) + assert code == 2 + cls.assert_not_called() # must not silently do a full export + + def test_failures_produce_nonzero_exit(self, tmp_path): + exporter = _mock_exporter( + {'exported': 0, 'skipped': 0, 'failed': [{'run_id': 'x', 'error': 'e'}]} + ) + with mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ): + code = run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + ] + ) + assert code == 1 + + def test_top_level_cli_help_does_not_need_migrate_extras(self): + result = subprocess.run( + [sys.executable, '-m', 'pluto', 'migrate', '--help'], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert 'wandb' in result.stdout diff --git a/tests/test_migrate_loader.py b/tests/test_migrate_loader.py new file mode 100644 index 0000000..08d86d9 --- /dev/null +++ b/tests/test_migrate_loader.py @@ -0,0 +1,316 @@ +""" +Unit tests for pluto.migrate.loader.PlutoLoader. + +The loader replays staged export dirs into Pluto through the public +client API. pluto.init is mocked; these tests pin the init kwargs +(external id, compat createdAt, import tag, host-pollution guards), the +per-step metric replay with original timestamps, media/histogram +conversion, console/artifact replay, finish-code mapping, dedup, and +resume caching. +""" + +from __future__ import annotations + +import json +from unittest import mock + +import pytest + +pytest.importorskip('pyarrow') + +import pluto # noqa: E402 +from pluto.migrate.loader import PlutoLoader # noqa: E402 +from pluto.migrate.schema import PartWriter # noqa: E402 +from pluto.migrate.state import ( # noqa: E402 + LoadedCache, + mark_run_exported, + write_json_atomic, +) + +CREATED_AT_MS = 1746093600000 +UPDATED_AT_MS = CREATED_AT_MS + 7200000 +T0_MS = 1746093601000 +EXTERNAL_ID = 'wandb::acme/vision/abc123' + + +def _stage_run(tmp_path, run_id='abc123', state='finished'): + run_dir = tmp_path / 'acme' / 'vision' / 'runs' / run_id + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': 'vision', + 'run_id': run_id, + 'name': 'sunny-lion-1', + 'notes': 'baseline run', + 'tags': ['baseline'], + 'state': state, + 'config': {'lr': 0.1}, + 'summary': {'loss': 0.05}, + 'createdAt': CREATED_AT_MS, + 'updatedAt': UPDATED_AT_MS, + 'url': f'https://wandb.ai/acme/vision/runs/{run_id}', + }, + ) + media_file = run_dir / 'files' / 'media' / 'images' / 'sample_3.png' + media_file.parent.mkdir(parents=True) + media_file.write_bytes(b'PNG') + artifact_file = run_dir / 'artifacts' / 'model-weights:v2' / 'model.pt' + artifact_file.parent.mkdir(parents=True) + artifact_file.write_bytes(b'weights') + + base = dict(project_id='acme/vision', run_id=run_id) + with PartWriter(run_dir) as w: + w.write_row( + **base, + attribute_path='loss', + attribute_type='metric', + step=0, + timestamp_ms=T0_MS, + float_value=1.0, + ) + w.write_row( + **base, + attribute_path='acc', + attribute_type='metric', + step=0, + timestamp_ms=T0_MS, + float_value=0.1, + ) + w.write_row( + **base, + attribute_path='loss', + attribute_type='metric', + step=1, + timestamp_ms=T0_MS + 1000, + float_value=0.5, + ) + w.write_row( + **base, + attribute_path='sample', + attribute_type='media', + step=3, + timestamp_ms=T0_MS + 3000, + string_value='image-file', + file_value='files/media/images/sample_3.png', + caption='a dog', + ) + w.write_row( + **base, + attribute_path='weights', + attribute_type='media', + step=3, + timestamp_ms=T0_MS + 3000, + string_value=json.dumps( + {'_type': 'histogram', 'values': [1, 2, 1], 'bins': [0, 1, 2, 3]} + ), + ) + w.write_row( + **base, + attribute_path='system.gpu.0.gpu', + attribute_type='system_metric', + step=0, + timestamp_ms=T0_MS, + float_value=55.0, + ) + w.write_row( + **base, + attribute_path='console', + attribute_type='console', + step=1, + timestamp_ms=T0_MS, + string_value='starting up', + ) + w.write_row( + **base, + attribute_path='model-weights:v2', + attribute_type='artifact', + step=0, + timestamp_ms=CREATED_AT_MS, + string_value=json.dumps({'name': 'model-weights:v2', 'type': 'model'}), + file_value='artifacts/model-weights:v2/model.pt', + ) + mark_run_exported(run_dir, {'rows': 8}) + return run_dir + + +@pytest.fixture +def mock_init(): + with mock.patch('pluto.init') as init: + op = mock.MagicMock() + op.settings._op_id = 42 + op._sync_manager.get_pending_count.return_value = 0 + init.return_value = op + yield init, op + + +def _log_calls_with(op, predicate): + return [c for c in op.log.call_args_list if predicate(c.args[0])] + + +class TestPlutoLoader: + def test_init_kwargs(self, tmp_path, mock_init): + init, op = mock_init + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + kwargs = init.call_args.kwargs + assert kwargs['project'] == 'vision' + assert kwargs['name'] == 'sunny-lion-1' + assert kwargs['config'] == {'lr': 0.1} + assert kwargs['tags'] == ['baseline', 'import:wandb'] + assert kwargs['run_id'] == EXTERNAL_ID + settings = kwargs['settings'] + assert settings['compat'] == { + 'createdAt': CREATED_AT_MS, + 'updatedAt': UPDATED_AT_MS, + } + assert settings['disable_console'] is True + assert settings['disable_system_metrics'] is True + # The historical-timestamp path only exists in the sync store. + assert settings['sync_process_enabled'] is True + + def test_wandb_scalars_pushed_via_update_config(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + wandb_block = op.update_config.call_args.args[0]['wandb'] + assert wandb_block['notes'] == 'baseline run' + assert wandb_block['state'] == 'finished' + assert wandb_block['summary'] == {'loss': 0.05} + + def test_metrics_batched_per_step_with_timestamps(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + op._log_metrics_batch.assert_called_once() + groups = op._log_metrics_batch.call_args.args[0] + assert groups[0] == ({'loss': 1.0, 'acc': 0.1}, 0, T0_MS / 1000) + assert groups[1] == ({'loss': 0.5}, 1, (T0_MS + 1000) / 1000) + + def test_media_converted_to_pluto_types(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + image_calls = _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Image) for v in d.values()) + ) + assert len(image_calls) == 1 + call = image_calls[0] + assert list(call.args[0]) == ['sample'] + assert call.args[0]['sample']._caption == 'a dog' + assert call.kwargs == {'step': 3, 'timestamp': (T0_MS + 3000) / 1000} + + hist_calls = _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Histogram) for v in d.values()) + ) + assert len(hist_calls) == 1 + hist = hist_calls[0].args[0]['weights'] + assert hist._freq == [1, 2, 1] + assert hist._bins == [0, 1, 2, 3] + + def test_system_metrics_translated_to_sys_names(self, tmp_path, mock_init): + # Staged rows keep wandb's source-native 'system.*' names; the + # loader owns the translation to Pluto's 'sys/' namespace. + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + groups = op._log_metrics_batch.call_args.args[0] + assert ({'sys/gpu.0.gpu': 55.0}, 0, T0_MS / 1000) in groups + + def test_console_replayed_with_timestamps(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + op._log_console.assert_called_once_with( + [('starting up', 'INFO', T0_MS / 1000, 1)] + ) + + def test_artifacts_replayed(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + art_calls = _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Artifact) for v in d.values()) + ) + assert len(art_calls) == 1 + assert art_calls[0].kwargs['step'] == 0 + + def test_finish_code_mapping(self, tmp_path, mock_init): + _, op = mock_init + _stage_run(tmp_path, run_id='crashed1', state='crashed') + PlutoLoader(tmp_path).load() + op.finish.assert_called_once_with(code=1) + + def test_loaded_cache_written_and_resume_skips(self, tmp_path, mock_init): + init, op = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path).load() + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(EXTERNAL_ID) + + init.reset_mock() + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 0, 'skipped': 1, 'failed': []} + init.assert_not_called() + + def test_external_id_collision_resumes_and_replays(self, tmp_path, mock_init): + # A collision means a previous load created the run but never made + # it into loaded_runs.json (e.g. crashed mid-replay) — the loader + # must resume and re-replay, not skip and lose the remaining data. + from pluto.op import RunExistsError + + init, op = mock_init + init.side_effect = [ + RunExistsError( + "Run with externalId 'wandb::acme/vision/abc123' already exists." + ), + op, + ] + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + assert init.call_count == 2 + assert init.call_args_list[0].kwargs['resume'] is False + assert init.call_args_list[1].kwargs['resume'] is True + op._log_metrics_batch.assert_called_once() # actually re-replayed + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(EXTERNAL_ID) + + def test_backpressure_throttles_then_gives_up(self, tmp_path, mock_init): + _, op = mock_init + op._sync_manager.get_pending_count.side_effect = [10, 10, 3] + _stage_run(tmp_path) + with mock.patch('pluto.migrate.loader.time.sleep') as sleep: + PlutoLoader(tmp_path, max_pending=5).load() + assert sleep.called # throttled while pending > max_pending + + op._sync_manager.get_pending_count.side_effect = None + op._sync_manager.get_pending_count.return_value = 10 + (tmp_path / 'loaded_runs.json').unlink() + with mock.patch('pluto.migrate.loader.time.sleep'): + summary = PlutoLoader(tmp_path, max_pending=5, stall_timeout=0).load() + assert summary['loaded'] == 1 # bounded: gives up waiting, keeps going + + def test_dry_run_makes_no_runs(self, tmp_path, mock_init): + init, _ = mock_init + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path, dry_run=True).load() + assert summary['loaded'] == 0 + init.assert_not_called() + assert not (tmp_path / 'loaded_runs.json').exists() + + def test_dest_project_override(self, tmp_path, mock_init): + init, _ = mock_init + _stage_run(tmp_path) + PlutoLoader(tmp_path, dest_project='legacy-wandb').load() + assert init.call_args.kwargs['project'] == 'legacy-wandb' + + def test_missing_media_file_skipped_not_fatal(self, tmp_path, mock_init): + _, op = mock_init + run_dir = _stage_run(tmp_path) + (run_dir / 'files' / 'media' / 'images' / 'sample_3.png').unlink() + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 1 + assert not _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Image) for v in d.values()) + ) diff --git a/tests/test_migrate_schema.py b/tests/test_migrate_schema.py new file mode 100644 index 0000000..905c357 --- /dev/null +++ b/tests/test_migrate_schema.py @@ -0,0 +1,132 @@ +""" +Unit tests for pluto.migrate.schema (parquet part writing/reading) and +pluto.migrate.state (resume bookkeeping). + +The migration pipeline stages source-platform data on disk as parquet +parts in a long/tall schema (one row per data point) plus JSON state +files that make both phases resumable. These tests pin the round-trip, +part rotation, and resume semantics. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip('pyarrow') + +from pluto.migrate.schema import ( # noqa: E402 + ATTRIBUTE_TYPES, + PartWriter, + iter_part_tables, + part_files, +) +from pluto.migrate.state import ( # noqa: E402 + LoadedCache, + is_run_exported, + mark_run_exported, + read_json, + write_json_atomic, +) + + +def _metric_row(i: int) -> dict: + return dict( + project_id='acme/vision', + run_id='abc123', + attribute_path=f'loss/train_{i % 3}', + attribute_type='metric', + step=i, + timestamp_ms=1600000000000 + i, + float_value=float(i) / 7, + ) + + +class TestPartWriter: + def test_round_trip_preserves_rows_and_order(self, tmp_path): + with PartWriter(tmp_path) as w: + for i in range(10): + w.write_row(**_metric_row(i)) + w.write_row( + project_id='acme/vision', + run_id='abc123', + attribute_path='sample', + attribute_type='media', + step=3, + timestamp_ms=1600000000003, + string_value='image-file', + file_value='files/media/images/sample_3.png', + caption='a caption', + ) + + rows = [] + for table in iter_part_tables(tmp_path): + rows.extend(table.to_pylist()) + + assert len(rows) == 11 + assert [r['step'] for r in rows[:10]] == list(range(10)) + assert rows[0]['float_value'] == pytest.approx(0.0) + assert rows[0]['string_value'] is None + media = rows[10] + assert media['attribute_type'] == 'media' + assert media['file_value'] == 'files/media/images/sample_3.png' + assert media['caption'] == 'a caption' + assert media['float_value'] is None + + def test_rotation_creates_multiple_ordered_parts(self, tmp_path): + with PartWriter(tmp_path, max_part_bytes=1, rows_per_flush=10) as w: + for i in range(35): + w.write_row(**_metric_row(i)) + + files = part_files(tmp_path) + assert len(files) > 1 + assert files == sorted(files) + + steps = [] + for table in iter_part_tables(tmp_path): + steps.extend(table.column('step').to_pylist()) + assert steps == list(range(35)) + + def test_invalid_attribute_type_raises(self, tmp_path): + with PartWriter(tmp_path) as w: + with pytest.raises(ValueError, match='attribute_type'): + w.write_row(**{**_metric_row(0), 'attribute_type': 'bogus'}) + + def test_no_rows_writes_no_parts(self, tmp_path): + with PartWriter(tmp_path): + pass + assert part_files(tmp_path) == [] + + def test_attribute_types_cover_migration_scope(self): + assert ATTRIBUTE_TYPES == { + 'metric', + 'system_metric', + 'media', + 'console', + 'artifact', + } + + +class TestState: + def test_write_json_atomic_round_trip_and_overwrite(self, tmp_path): + path = tmp_path / 'x.json' + write_json_atomic(path, {'a': 1}) + assert read_json(path) == {'a': 1} + write_json_atomic(path, {'a': 2}) + assert read_json(path) == {'a': 2} + + def test_export_sentinel(self, tmp_path): + run_dir = tmp_path / 'run1' + run_dir.mkdir() + assert not is_run_exported(run_dir) + mark_run_exported(run_dir, {'rows': 42}) + assert is_run_exported(run_dir) + + def test_loaded_cache_persists_across_instances(self, tmp_path): + path = tmp_path / 'loaded_runs.json' + cache = LoadedCache(path) + assert not cache.is_loaded('wandb::acme/vision/abc123') + cache.mark_loaded('wandb::acme/vision/abc123', {'pluto_run_id': 7}) + + reopened = LoadedCache(path) + assert reopened.is_loaded('wandb::acme/vision/abc123') + assert not reopened.is_loaded('wandb::acme/vision/other') diff --git a/tests/test_migrate_staging_e2e.py b/tests/test_migrate_staging_e2e.py new file mode 100644 index 0000000..7912cdc --- /dev/null +++ b/tests/test_migrate_staging_e2e.py @@ -0,0 +1,188 @@ +""" +Client<->server integration test for the migration pipeline, run against +the STAGING (dev) environment. + +Loads a hand-staged export (no wandb involved) into a real Pluto server +via PlutoLoader and verifies through the query API that the historical +data actually round-tripped: metric values/steps, original per-point +wall-clock timestamps, tags, and — once the server-side fix is deployed +— the run's backfilled createdAt. + +Gated on PLUTO_STAGING_API_KEY so ordinary CI (which only has a +production PLUTO_API_KEY) skips it: + + PLUTO_STAGING_API_KEY=... poetry run pytest tests/test_migrate_staging_e2e.py + +URL overrides (defaults point at the dev channel): + PLUTO_STAGING_URL_APP https://pluto-dev.trainy.ai + PLUTO_STAGING_URL_API https://pluto-api-dev.trainy.ai + PLUTO_STAGING_URL_INGEST https://pluto-ingest-dev.trainy.ai +""" + +from __future__ import annotations + +import os +import time +import uuid +from datetime import datetime, timezone + +import pytest + +pytest.importorskip('pyarrow') + +from pluto.migrate.loader import PlutoLoader # noqa: E402 +from pluto.migrate.schema import PartWriter # noqa: E402 +from pluto.migrate.state import mark_run_exported, write_json_atomic # noqa: E402 + +STAGING_API_KEY = os.environ.get('PLUTO_STAGING_API_KEY') +URL_APP = os.environ.get('PLUTO_STAGING_URL_APP', 'https://pluto-dev.trainy.ai') +URL_API = os.environ.get('PLUTO_STAGING_URL_API', 'https://pluto-api-dev.trainy.ai') +URL_INGEST = os.environ.get( + 'PLUTO_STAGING_URL_INGEST', 'https://pluto-ingest-dev.trainy.ai' +) +URL_PY = os.environ.get('PLUTO_STAGING_URL_PY', 'https://pluto-py-dev.trainy.ai') + +pytestmark = pytest.mark.skipif( + not STAGING_API_KEY, + reason='PLUTO_STAGING_API_KEY not set; staging integration test skipped', +) + +PROJECT = 'migrate-staging-e2e' +CREATED_AT_MS = 1600000000000 # 2020-09-13T12:26:40Z +T0_MS = CREATED_AT_MS + 60_000 +METRIC_POINTS = [ # (step, timestamp_ms, loss) + (0, T0_MS, 1.0), + (1, T0_MS + 1000, 0.5), + (2, T0_MS + 2000, 0.25), +] + + +def _stage_fixture_run(root, run_id: str): + run_dir = root / 'acme' / 'vision' / 'runs' / run_id + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': PROJECT, + 'run_id': run_id, + 'name': f'staging-e2e-{run_id}', + 'notes': 'staging integration fixture', + 'tags': ['fixture'], + 'state': 'finished', + 'config': {'lr': 0.1, 'optimizer': 'adamw'}, + 'summary': {'loss': 0.25}, + 'createdAt': CREATED_AT_MS, + 'updatedAt': CREATED_AT_MS + 3_600_000, + 'url': f'https://wandb.ai/acme/vision/runs/{run_id}', + }, + ) + with PartWriter(run_dir) as w: + for step, ts_ms, loss in METRIC_POINTS: + w.write_row( + project_id='acme/vision', + run_id=run_id, + attribute_path='loss', + attribute_type='metric', + step=step, + timestamp_ms=ts_ms, + float_value=loss, + ) + w.write_row( + project_id='acme/vision', + run_id=run_id, + attribute_path='console', + attribute_type='console', + step=1, + timestamp_ms=T0_MS, + string_value='hello from 2020', + ) + mark_run_exported(run_dir, {'rows': len(METRIC_POINTS) + 1}) + return run_dir + + +def _parse_point_time_ms(value) -> int: + if isinstance(value, (int, float)): + return int(value if value > 1e11 else value * 1000) + dt = datetime.fromisoformat(str(value).replace('Z', '+00:00')) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return int(dt.timestamp() * 1000) + + +@pytest.fixture +def staging_env(monkeypatch, tmp_path): + monkeypatch.setenv('PLUTO_API_KEY', STAGING_API_KEY) + monkeypatch.setenv('PLUTO_URL_APP', URL_APP) + monkeypatch.setenv('PLUTO_URL_API', URL_API) + monkeypatch.setenv('PLUTO_URL_INGEST', URL_INGEST) + monkeypatch.setenv('PLUTO_URL_PY', URL_PY) + monkeypatch.setenv('PLUTO_DIR', str(tmp_path / 'staging')) + + +def test_migration_round_trip_against_staging(tmp_path, staging_env): + from pluto import query + + run_id = uuid.uuid4().hex[:12] + _stage_fixture_run(tmp_path, run_id) + + summary = PlutoLoader(tmp_path).load() + assert summary['failed'] == [] + assert summary['loaded'] == 1 + + client = query.Client(api_token=STAGING_API_KEY, host=URL_API) + runs = client.list_runs(PROJECT, search=f'staging-e2e-{run_id}') + match = [r for r in runs if r['name'] == f'staging-e2e-{run_id}'] + assert match, f'imported run staging-e2e-{run_id} not found on staging' + run = client.get_run(PROJECT, match[0]['id']) + + assert 'import:wandb' in run['tags'] + assert 'fixture' in run['tags'] + + # Metrics may take a moment to land in ClickHouse; poll get_metrics + # (NOT get_metric_names, which lags minutes behind ingest). + deadline = time.time() + 120 + rows = [] + while time.time() < deadline: + data = client.get_metrics(PROJECT, match[0]['id'], metric_names=['loss']) + rows = data.to_dict('records') if hasattr(data, 'to_dict') else list(data) + if len(rows) >= len(METRIC_POINTS): + break + time.sleep(5) + assert len(rows) == len(METRIC_POINTS), f'expected 3 points, got {rows}' + + by_step = {r['step']: r for r in rows} + for step, ts_ms, loss in METRIC_POINTS: + assert by_step[step]['value'] == pytest.approx(loss) + assert _parse_point_time_ms(by_step[step]['time']) == ts_ms, ( + f'historical timestamp not preserved for step {step}: ' + f'{by_step[step]["time"]!r}' + ) + + # Run createdAt backfill needs the server-side fix; xfail until deployed. + run_created_ms = _parse_point_time_ms(run['createdAt']) + if abs(run_created_ms - CREATED_AT_MS) > 60_000: + pytest.xfail( + 'server does not yet honor createdAt on the run row ' + '(fix pending deploy); run created at ' + f'{run["createdAt"]!r} instead of 2020-09-13' + ) + assert run_created_ms == CREATED_AT_MS + + +def test_reload_is_idempotent_against_staging(tmp_path, staging_env): + run_id = uuid.uuid4().hex[:12] + _stage_fixture_run(tmp_path, run_id) + + first = PlutoLoader(tmp_path).load() + assert first['loaded'] == 1 + + # Local cache skip + second = PlutoLoader(tmp_path).load() + assert second == {'loaded': 0, 'skipped': 1, 'failed': []} + + # Server-side external-id dedup (fresh cache simulates another machine) + (tmp_path / 'loaded_runs.json').unlink() + third = PlutoLoader(tmp_path).load() + assert third['loaded'] == 0 + assert third['skipped'] == 1 diff --git a/tests/test_migrate_wandb_export.py b/tests/test_migrate_wandb_export.py new file mode 100644 index 0000000..f50fd80 --- /dev/null +++ b/tests/test_migrate_wandb_export.py @@ -0,0 +1,295 @@ +""" +Unit tests for pluto.migrate.wandb_export.WandbExporter. + +The exporter reads runs from the wandb cloud API and stages them on disk +(parquet parts + run.json + downloaded files). These tests drive it with +fake wandb API objects — no network, no real wandb — and pin the staged +layout: metric/media/system/console/artifact rows, original timestamps, +resume-by-sentinel, and the artifact size cap. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +pytest.importorskip('pyarrow') + +from pluto.migrate.schema import iter_part_tables # noqa: E402 +from pluto.migrate.state import is_run_exported, read_json # noqa: E402 +from pluto.migrate.wandb_export import WandbExporter # noqa: E402 + +CREATED_AT = '2025-05-01T10:00:00Z' +CREATED_AT_MS = 1746093600000 +T0 = 1746093601.0 # first history point, epoch seconds + + +class FakeSummary: + def __init__(self, d): + self._json_dict = d + + +class FakeFile: + def __init__(self, name, size=10, content=b'x'): + self.name = name + self.size = size + self._content = content + + def download(self, root, replace=False, exist_ok=False): + path = Path(root) / self.name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(self._content) + return path + + +class FakeArtifact: + def __init__(self, name='model-weights:v2', type='model', size=100): + self.name = name + self.type = type + self.size = size + self.created_at = CREATED_AT + + def download(self, root): + root = Path(root) + root.mkdir(parents=True, exist_ok=True) + (root / 'model.pt').write_bytes(b'weights') + return str(root) + + +class FakeRun: + entity = 'acme' + project = 'vision' + + def __init__(self, run_id='abc123', artifacts=None, output_log=None): + self.id = run_id + self.name = 'sunny-lion-1' + self.notes = 'baseline run' + self.tags = ['baseline'] + self.state = 'finished' + self.config = {'lr': 0.1} + self.summary = FakeSummary({'loss': 0.05, '_wandb': {'runtime': 12}}) + self.created_at = CREATED_AT + self.heartbeat_at = '2025-05-01T12:00:00Z' + self.url = f'https://wandb.ai/acme/vision/runs/{run_id}' + self.metadata = {'gpu': 'NVIDIA H100'} + self._artifacts = artifacts if artifacts is not None else [FakeArtifact()] + self._files = [ + FakeFile( + 'output.log', + content=output_log + if output_log is not None + else b'starting up\nepoch 0 done\n', + ), + FakeFile('media/images/sample_3_abc.png', content=b'PNG'), + FakeFile('requirements.txt'), + ] + self.scan_history_calls = 0 + + def scan_history(self, page_size=1000): + self.scan_history_calls += 1 + return iter( + [ + {'_step': 0, '_timestamp': T0, 'loss': 1.0, 'acc': 0.1}, + {'_step': 1, '_timestamp': T0 + 1, 'loss': 0.5}, + { + '_step': 3, + '_timestamp': T0 + 3, + 'sample': { + '_type': 'image-file', + 'path': 'media/images/sample_3_abc.png', + 'caption': 'a dog', + }, + 'weights': { + '_type': 'histogram', + 'values': [1, 2, 1], + 'bins': [0, 1, 2, 3], + }, + }, + ] + ) + + def history(self, stream='default', pandas=True, samples=None): + assert stream == 'events' and pandas is False + return [ + {'_timestamp': T0, 'system.gpu.0.gpu': 55.0, 'system.cpu': 12.0}, + {'_timestamp': T0 + 2, 'system.gpu.0.gpu': 60.0}, + ] + + def files(self): + return list(self._files) + + def logged_artifacts(self): + return list(self._artifacts) + + +class FakeApi: + def __init__(self, runs): + self._runs = runs + + def runs(self, path, filters=None): + assert path == 'acme/vision' + return list(self._runs) + + +def _export(tmp_path, run=None, **kwargs): + run = run or FakeRun() + exporter = WandbExporter( + entity='acme', + project='vision', + output_dir=tmp_path, + api=FakeApi([run]), + **kwargs, + ) + summary = exporter.export() + return run, tmp_path / 'acme' / 'vision' / 'runs' / run.id, summary + + +def _rows(run_dir, attribute_type=None): + rows = [] + for table in iter_part_tables(run_dir): + rows.extend(table.to_pylist()) + if attribute_type: + rows = [r for r in rows if r['attribute_type'] == attribute_type] + return rows + + +class TestWandbExporter: + def test_run_json_manifest(self, tmp_path): + run, run_dir, summary = _export(tmp_path) + assert summary == {'exported': 1, 'skipped': 0, 'failed': []} + manifest = read_json(run_dir / 'run.json') + assert manifest['name'] == 'sunny-lion-1' + assert manifest['notes'] == 'baseline run' + assert manifest['tags'] == ['baseline'] + assert manifest['state'] == 'finished' + assert manifest['config'] == {'lr': 0.1} + assert manifest['summary'] == {'loss': 0.05} # _wandb internals dropped + assert manifest['createdAt'] == CREATED_AT_MS + assert manifest['updatedAt'] == CREATED_AT_MS + 2 * 3600 * 1000 + assert manifest['url'] == run.url + assert is_run_exported(run_dir) + + def test_metric_rows_preserve_step_and_timestamp(self, tmp_path): + _, run_dir, _ = _export(tmp_path) + metrics = _rows(run_dir, 'metric') + assert { + 'attribute_path': 'loss', + 'step': 0, + 'timestamp_ms': int(T0 * 1000), + 'float_value': 1.0, + }.items() <= metrics[0].items() + assert [m['attribute_path'] for m in metrics] == ['loss', 'acc', 'loss'] + + def test_media_and_histogram_rows(self, tmp_path): + _, run_dir, _ = _export(tmp_path) + media = {m['attribute_path']: m for m in _rows(run_dir, 'media')} + img = media['sample'] + assert img['string_value'] == 'image-file' + assert img['file_value'] == 'files/media/images/sample_3_abc.png' + assert img['caption'] == 'a dog' + assert img['step'] == 3 + assert (run_dir / 'files/media/images/sample_3_abc.png').exists() + hist = media['weights'] + assert hist['file_value'] is None + assert json.loads(hist['string_value']) == { + '_type': 'histogram', + 'values': [1, 2, 1], + 'bins': [0, 1, 2, 3], + } + + def test_system_metric_rows_keep_source_names(self, tmp_path): + # Staging is source-faithful; the loader owns the sys/ translation. + _, run_dir, _ = _export(tmp_path) + sys_rows = _rows(run_dir, 'system_metric') + assert { + 'attribute_path': 'system.gpu.0.gpu', + 'step': 0, + 'timestamp_ms': int(T0 * 1000), + 'float_value': 55.0, + }.items() <= sys_rows[0].items() + assert {r['attribute_path'] for r in sys_rows} == { + 'system.gpu.0.gpu', + 'system.cpu', + } + + def test_console_rows_from_output_log(self, tmp_path): + _, run_dir, _ = _export(tmp_path) + console = _rows(run_dir, 'console') + assert [(r['string_value'], r['step']) for r in console] == [ + ('starting up', 1), + ('epoch 0 done', 2), + ] + # no per-line timestamps in the log -> stamped with run createdAt + assert console[0]['timestamp_ms'] == CREATED_AT_MS + + def test_console_lines_with_timestamps_are_parsed(self, tmp_path): + log = b'2025-05-01T10:00:05.500Z first line\nplain line\n' + run = FakeRun(run_id='tsrun', output_log=log) + _, run_dir, _ = _export(tmp_path, run=run) + console = _rows(run_dir, 'console') + # Timestamp parsed for the row time, but the message content is + # preserved verbatim — user log lines must not be rewritten. + assert console[0]['string_value'] == '2025-05-01T10:00:05.500Z first line' + assert console[0]['timestamp_ms'] == CREATED_AT_MS + 5500 + assert console[1]['string_value'] == 'plain line' + assert console[1]['timestamp_ms'] == CREATED_AT_MS + + def test_artifact_rows_and_download(self, tmp_path): + _, run_dir, _ = _export(tmp_path) + rows = _rows(run_dir, 'artifact') + assert len(rows) == 1 + row = rows[0] + assert row['attribute_path'] == 'model-weights:v2' + assert row['file_value'] == 'artifacts/model-weights:v2/model.pt' + assert row['timestamp_ms'] == CREATED_AT_MS + meta = json.loads(row['string_value']) + assert meta['type'] == 'model' + assert (run_dir / 'artifacts/model-weights:v2/model.pt').exists() + + def test_artifact_size_cap_skips_download(self, tmp_path): + big = FakeArtifact(name='huge:v0', size=10**12) + run = FakeRun(run_id='bigrun', artifacts=[big]) + _, run_dir, _ = _export(tmp_path, run=run, artifact_max_bytes=10**6) + assert _rows(run_dir, 'artifact') == [] + assert not (run_dir / 'artifacts').exists() + + def test_resume_skips_completed_runs(self, tmp_path): + run, run_dir, _ = _export(tmp_path) + assert run.scan_history_calls == 1 + exporter = WandbExporter( + entity='acme', project='vision', output_dir=tmp_path, api=FakeApi([run]) + ) + summary = exporter.export() + assert summary == {'exported': 0, 'skipped': 1, 'failed': []} + assert run.scan_history_calls == 1 # untouched on resume + + def test_run_failure_is_recorded_not_raised(self, tmp_path): + run = FakeRun(run_id='boom') + run.scan_history = lambda page_size=1000: (_ for _ in ()).throw( + RuntimeError('api exploded') + ) + exporter = WandbExporter( + entity='acme', project='vision', output_dir=tmp_path, api=FakeApi([run]) + ) + summary = exporter.export() + assert summary['exported'] == 0 + assert summary['failed'] and summary['failed'][0]['run_id'] == 'boom' + assert not is_run_exported(tmp_path / 'acme' / 'vision' / 'runs' / 'boom') + manifest = read_json(tmp_path / 'manifest.json') + assert manifest['failed'][0]['run_id'] == 'boom' + + def test_run_ids_filter(self, tmp_path): + wanted, unwanted = FakeRun(run_id='keep'), FakeRun(run_id='drop') + exporter = WandbExporter( + entity='acme', + project='vision', + output_dir=tmp_path, + api=FakeApi([wanted, unwanted]), + run_ids=['keep'], + ) + summary = exporter.export() + assert summary['exported'] == 1 + assert (tmp_path / 'acme/vision/runs/keep').exists() + assert not (tmp_path / 'acme/vision/runs/drop').exists()