diff --git a/docs-api/media.mdx b/docs-api/media.mdx index 10988e7..6a1bbba 100644 --- a/docs-api/media.mdx +++ b/docs-api/media.mdx @@ -26,6 +26,9 @@ File(path: str, name: Optional[str] = None, **kwargs) -> None Image( data: Union[str, PILImage.Image, np.ndarray, bytes, bytearray], caption: Optional[str] = None, + boxes: Optional[Dict[str, Any]] = None, + masks: Optional[Dict[str, Any]] = None, + annotations: Optional[Union[str, Dict[str, Any]]] = None, ) -> None ``` diff --git a/docs-api/meta.json b/docs-api/meta.json index 3db89d6..a6f0bd9 100644 --- a/docs-api/meta.json +++ b/docs-api/meta.json @@ -15,6 +15,7 @@ "Table", "Text", "Video", + "agent", "alert", "finish", "generate_run_id", @@ -23,6 +24,7 @@ "logout", "query", "setup", + "sweep", "watch" ], "pages": [ 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/__init__.py b/pluto/__init__.py index 38b8679..522c1c2 100644 --- a/pluto/__init__.py +++ b/pluto/__init__.py @@ -8,6 +8,7 @@ from .file import Artifact, Audio, File, Image, Text, Video from .init import finish, init from .sets import Settings, setup +from .sweep import agent, sweep from .sys import System from .util import generate_run_id @@ -39,6 +40,8 @@ 'setup', 'query', 'generate_run_id', + 'sweep', + 'agent', ) __version__ = '0.0.27' 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/api.py b/pluto/api.py index b75b41b..4d251cc 100644 --- a/pluto/api.py +++ b/pluto/api.py @@ -82,6 +82,14 @@ def make_compat_status_v1(settings, trace=None): 'status': STATUS[settings._op_status], # "metadata": json.dumps(settings.meta), 'statusMetadata': json.dumps(trace) if trace is not None else None, + # Historical terminal-status time (epoch ms) for backfilled/migrated + # runs (pluto.migrate). Server applies it to the durable + # `statusUpdated` column so Duration = end - createdAt is correct; + # duration's `end` reads `statusUpdated ?? updatedAt`, and updatedAt + # re-bumps via Prisma @updatedAt, so statusUpdated is the field that + # must stick. None for normal runs (empty compat) -> server keeps + # now(), so behavior is unchanged. + 'statusUpdated': settings.compat.get('updatedAt'), } ).encode() 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/file.py b/pluto/file.py index d362541..79b447f 100644 --- a/pluto/file.py +++ b/pluto/file.py @@ -1,4 +1,6 @@ import hashlib +import io +import json import logging import mimetypes import os @@ -56,6 +58,13 @@ class File: # override this instance attribute in their __init__; the class-level default # ensures it always exists (e.g. on a directly-constructed File). _caption: Optional[str] = None + # Optional opaque JSON string of image annotations (wandb-shape boxes/masks), + # sent to the server as mlop_files.annotations. Only Image sets it. + _annotations: Optional[str] = None + # Optional override for the upload payload's fileType (e.g. "mask" for a + # segmentation-mask PNG, so the frontend hides it from the media grid). + # None → fileType is derived from the extension as usual. + _upload_file_type: Optional[str] = None def __init__( self, @@ -225,12 +234,25 @@ def __init__( self, data: Union[str, 'PILImage.Image', np.ndarray, bytes, bytearray], caption: Optional[str] = None, + boxes: Optional[Dict[str, Any]] = None, + masks: Optional[Dict[str, Any]] = None, + annotations: Optional[Union[str, Dict[str, Any]]] = None, ) -> None: self._name = caption + f'.{uuid.uuid4()}' if caption else f'{uuid.uuid4()}' # Preserve the raw caption separately so it can be sent to the server # as a dedicated field (mlop_files.caption); _name keeps the legacy # caption-as-filename behavior for back-compat with older servers. self._caption = caption + # Image annotations (wandb-shape boxes/masks) → mlop_files.annotations. + # ``annotations`` is a ready JSON blob (raw string/dict) — used by the + # wandb migration, which forwards wandb's own {boxes, masks} verbatim. + # ``boxes`` ({layer: {box_data, class_labels}}) is folded in with each box + # defaulted to domain="pixel". ``masks`` ({layer: {mask_data|path, + # class_labels}}) each become a separate PNG uploaded as fileType "mask" + # and referenced by fileName; the mask PNGs to upload alongside are + # collected in ``_annotation_files``. + self._annotation_files: list = [] + self._annotations = self._build_annotations(annotations, boxes, masks) self._id = f'{uuid.uuid4()}{uuid.uuid4()}'.replace('-', '') self._ext = '.png' self._image: Any = None @@ -264,6 +286,92 @@ def __init__( logger.debug(f'{self.tag}: attempted conversion from array') self._image = make_compat_image_numpy(data) + def _build_annotations( + self, + annotations: Optional[Union[str, Dict[str, Any]]], + boxes: Optional[Dict[str, Any]], + masks: Optional[Dict[str, Any]], + ) -> Optional[str]: + """Assemble the annotations JSON string sent to the server. + + ``annotations`` (a ready JSON string/dict, wandb's {boxes, masks} shape) + is forwarded as-is; ``boxes`` is folded into ``annotations.boxes`` with + each box defaulted to ``domain: "pixel"``; ``masks`` each become a PNG + (uploaded separately as fileType "mask") referenced by ``fileName``. + Returns None when there is nothing to attach. + """ + # Ready JSON string with nothing to merge → forward verbatim. + if isinstance(annotations, str) and not boxes and not masks: + return annotations or None + result: Dict[str, Any] = {} + if isinstance(annotations, str): + try: + parsed = json.loads(annotations) + result = parsed if isinstance(parsed, dict) else {} + except (ValueError, TypeError): + result = {} + elif isinstance(annotations, dict): + result = dict(annotations) + if boxes: + merged = dict(result.get('boxes') or {}) + for layer, spec in boxes.items(): + if not isinstance(spec, dict): + continue + out = dict(spec) + box_data = out.get('box_data') + if isinstance(box_data, list): + out['box_data'] = [ + {'domain': 'pixel', **b} if isinstance(b, dict) else b + for b in box_data + ] + merged[layer] = out + if merged: + result['boxes'] = merged + if masks: + merged_masks = dict(result.get('masks') or {}) + for layer, spec in masks.items(): + entry = self._stage_mask(spec) + if entry is not None: + merged_masks[layer] = entry + if merged_masks: + result['masks'] = merged_masks + return json.dumps(result) if result else None + + def _stage_mask(self, spec: Any) -> Optional[Dict[str, Any]]: + """Turn one mask-layer spec into an uploadable PNG + its annotations ref. + + ``spec`` is ``{mask_data: }`` (native — encode a PNG + with the class id in the red channel) or ``{path: }`` + (migration — forward wandb's mask file). The PNG is queued in + ``_annotation_files`` (uploaded as fileType "mask"); returns + ``{fileName, class_labels?}`` for ``annotations.masks[layer]``. + """ + if not isinstance(spec, dict): + return None + png: Union[str, bytes, None] = None + if spec.get('path'): + png = spec['path'] # migration: an existing mask PNG on disk + elif spec.get('mask_data') is not None: + arr = np.asarray(spec['mask_data']).astype('uint8') + zeros = np.zeros_like(arr) + rgb = np.stack([arr, zeros, zeros], axis=-1) # class id → red channel + buf = io.BytesIO() + PILImage.fromarray(rgb, 'RGB').save(buf, format='PNG') + png = buf.getvalue() + if png is None: + return None + mask_name = f'{uuid.uuid4()}.mask' + mask_img = Image(png) + mask_img._name = mask_name # uploaded fileName = f'{mask_name}.png' + # Marks the upload's fileType as "mask" so the frontend hides it from the + # media grid and resolves it from annotations by fileName. + mask_img._upload_file_type = 'mask' + self._annotation_files.append(mask_img) + entry: Dict[str, Any] = {'fileName': f'{mask_name}.png'} + if spec.get('class_labels'): + entry['class_labels'] = spec['class_labels'] + return entry + def load(self, dir: Optional[str] = None) -> None: if not self._path: if dir: diff --git a/pluto/iface.py b/pluto/iface.py index 7026df9..a1f2bf6 100644 --- a/pluto/iface.py +++ b/pluto/iface.py @@ -173,12 +173,23 @@ def close(self) -> None: self.client_api.close() def update_status(self, trace: Union[Any, None] = None) -> None: - """Update run status on the server (called at finish).""" + """Update run status on the server (called at finish). + + This is the run's terminal state transition — losing it strands the + run in RUNNING (the UI then renders it as FAILED with an ever-growing + duration). So, unlike fire-and-forget uploads/heartbeats, it retries + transient connection resets and raises ``PlutoRequestError`` if it + still can't confirm the update, so the caller (e.g. the migration + loader) knows the finish was not recorded. + """ self._post_v1( self.settings.url_stop, self.headers, make_compat_status_v1(self.settings, trace), client=self.client_api, + name='status', + raise_on_error=True, + retry_connection_errors=True, ) def update_tags(self, tags: List[str]) -> None: @@ -300,6 +311,7 @@ def _try( timeout: Optional[float] = None, suppress_httpx_logs: bool = False, raise_on_error: bool = False, + retry_connection_errors: bool = False, ): effective_max_retries = ( max_retries @@ -336,8 +348,17 @@ def _try( # former carries a server-provided reason worth raising. last_status # is threaded down from the last HTTP response so the exception # reports the real code (e.g. 500) rather than None. - if raise_on_error and error_info.startswith('HTTP '): - raise PlutoRequestError(error_info, status_code=last_status) + # A persistent server error carries a server-provided reason worth + # raising. When retry_connection_errors is set (critical one-shot + # requests like the terminal status update), also raise on an + # exhausted network failure so the caller can't mistake a dropped + # request for success. + if raise_on_error and ( + error_info.startswith('HTTP ') or retry_connection_errors + ): + raise PlutoRequestError( + error_info or 'request failed', status_code=last_status + ) return None @@ -401,16 +422,34 @@ def _try( httpx.RemoteProtocolError, httpx.LocalProtocolError, ) as e: - # Treat connection errors as shutdown signals - don't retry - # This prevents hanging during atexit when sockets are being torn down + if not retry_connection_errors: + # Default: treat connection errors as shutdown signals - don't + # retry. This prevents hanging during atexit when sockets are + # being torn down (heartbeat / trigger / streaming-upload spam). + logger.debug( + '%s: %s: connection error (likely shutdown): %s: %s', + tag, + name, + type(e).__name__, + e, + ) + return None + # Critical one-shot request (e.g. the terminal status update): a + # keep-alive socket dropped mid-request is transient, NOT a shutdown + # signal. Fall through to the retry/backoff path below so a single + # reset can't silently strand the run's finished status on the + # server (which then shows the run as stuck / FAILED forever). + error_info = f'{type(e).__name__}: {str(e)}' logger.debug( - '%s: %s: connection error (likely shutdown): %s: %s', + '%s: %s: attempt %s/%s: connection error from %s: %s: %s', tag, name, + retry + 1, + effective_max_retries + 1, + url, type(e).__name__, e, ) - return None except Exception as e: # Capture error info for potential failure logging error_info = f'{type(e).__name__}: {str(e)}' @@ -446,6 +485,7 @@ def _try( timeout=timeout, suppress_httpx_logs=suppress_httpx_logs, raise_on_error=raise_on_error, + retry_connection_errors=retry_connection_errors, ) def _put_v1( @@ -481,6 +521,7 @@ def _post_v1( timeout: Optional[float] = None, suppress_httpx_logs: bool = False, raise_on_error: bool = False, + retry_connection_errors: bool = False, ): # Support both queue and direct content if isinstance(q, queue.Queue): @@ -503,6 +544,7 @@ def _post_v1( timeout=timeout, suppress_httpx_logs=suppress_httpx_logs, raise_on_error=raise_on_error, + retry_connection_errors=retry_connection_errors, ) if ( diff --git a/pluto/init.py b/pluto/init.py index af36df8..26947e1 100644 --- a/pluto/init.py +++ b/pluto/init.py @@ -260,11 +260,42 @@ def init( e, ) + # Sweep: if pluto.agent set an active sweep context, merge the sampled + # hyperparameters into config (swept values win) and tag the run so it groups + # under its sweep. Mirrors how wandb.agent feeds wandb.config. The submodule + # is read via sys.modules because the pluto.sweep *function* shadows the + # pluto.sweep module attribute, and we need the live _active_sweep global. + import sys as _sys + + _sweep_mod = _sys.modules.get('pluto.sweep') + _active = getattr(_sweep_mod, '_active_sweep', None) if _sweep_mod else None + if _active is not None: + combo = _active.get('config') or {} + merged = dict(config) if isinstance(config, dict) else {} + merged.update(combo) + # Stamp the declared sweep spec (method/metric/search-space) so the + # server sees a native sweep's real declaration, like config.wandb.sweep + # does for migrated sweeps — no separate sweep entity required. + declared = getattr(_sweep_mod, '_active_declared', None) + if isinstance(declared, dict): + merged['sweep'] = declared + config = merged + sweep_tag = f'sweep:{_active["id"]}' + if sweep_tag not in normalized_tags: + normalized_tags.append(sweep_tag) + if project is None and _active.get('project'): + settings.project = get_char(_active['project']) + try: op_init = OpInit(config=config, tags=normalized_tags or None, resume=resume) op_init.setup(settings=settings) op = op_init.init() + # If this run was started under a sweep agent, hand the op back so the + # agent can read its objective metric (bayes) and finish it. + if _active is not None and _sweep_mod is not None: + setattr(_sweep_mod, '_last_run_op', op) + # Set Sentry context for this run _sentry.set_tag('project', settings.project) _sentry.set_tag('run_id', str(settings._op_id)) 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..df0cb2a --- /dev/null +++ b/pluto/migrate/cli.py @@ -0,0 +1,504 @@ +""" +CLI for pluto.migrate: `pluto migrate wandb export|load|all`. + +Three orthogonal knobs: + * phase — the subcommand: export / load / all (both) + * scope — which projects: default all, --project (repeatable), --exclude + * flags — --workers (parallelism), --cleanup, ... + +This module keeps its imports light — wandb/pyarrow (the 'migrate' extra) +load inside the handlers, so the base `pluto` CLI works without them and a +missing dep produces an install hint instead of an ImportError traceback. +""" + +from __future__ import annotations + +import argparse +import sys +import threading +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +from pluto.migrate import _INSTALL_HINT + +# Seconds a load pass waits for the export to finish before rescanning for +# newly-completed runs (the `all` pipeline). +_ALL_POLL_SECONDS = 10.0 + + +# --------------------------------------------------------------------------- # +# Argument flags +# --------------------------------------------------------------------------- # +def _add_common_flags(parser: argparse.ArgumentParser) -> None: + """Scope + parallelism — shared by export/load/all.""" + parser.add_argument( + '--project', + action='append', + dest='projects', + default=None, + help='wandb project to migrate (repeatable). Omit to migrate ALL ' + 'projects under the entity.', + ) + parser.add_argument( + '--exclude', + action='append', + default=None, + help='project to skip when migrating all projects (repeatable)', + ) + parser.add_argument( + '--workers', + type=int, + default=4, + help='parallelism (default: 4): projects migrated concurrently, and the ' + 'per-run file-download budget on export. Each concurrent project holds ' + 'its run history + upload queue in memory, so raise this only if the host ' + 'has the RAM (~2-4 GB/worker); lower it on small machines.', + ) + + +def _add_export_flags(parser: argparse.ArgumentParser) -> None: + parser.add_argument('--entity', required=True, help='wandb entity (team/user)') + 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 N 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' + ) + parser.add_argument( + '--strict', + action='store_true', + help='exit non-zero if any run had data that could not be migrated', + ) + + +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; ' + 'only valid with a single --project)', + ) + 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 (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)', + ) + parser.add_argument( + '--cleanup', + action='store_true', + help="delete each run's staged files once it's confirmed loaded", + ) + + +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) + _add_common_flags(p_export) + + p_load = actions.add_parser( + 'load', help='load a staged export directory into Pluto' + ) + _add_load_flags(p_load) + _add_common_flags(p_load) + + p_all = actions.add_parser('all', help='export + load (one project or all)') + _add_export_flags(p_all) + _add_load_flags(p_all, with_input=False) + _add_common_flags(p_all) + + +# --------------------------------------------------------------------------- # +# Per-project workers (module-level so they're picklable for ProcessPoolExecutor) +# --------------------------------------------------------------------------- # +def _export_one_project( + args: argparse.Namespace, project: str, cache_path: Optional[str] = None +) -> int: + from pluto.migrate.wandb_export import WandbExporter + + try: + exporter = WandbExporter( + entity=args.entity, + project=project, + output_dir=args.output, + api_key=args.wandb_api_key, + run_ids=getattr(args, 'run_ids', None), + 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, + download_workers=args.workers, + ) + except ValueError as e: # e.g. an unparseable --after/--before + print(f'[{project}] error: {e}') + return 2 + + summary = exporter.export() + print( + f'[{project}] export: {summary["exported"]} exported, ' + f'{summary["skipped"]} skipped, {len(summary["failed"])} failed' + ) + for failure in summary['failed']: + print(f' [{project}] failed {failure["run_id"]}: {failure["error"]}') + not_migrated = summary.get('coverage', {}).get('not_migrated', {}) + if not_migrated: + dropped = ', '.join(f'{v} {k}' for k, v in sorted(not_migrated.items())) + print(f' [{project}] NOT migrated: {dropped}') + if getattr(args, 'strict', False): + print(f'[{project}] --strict: some data could not be migrated') + return 2 + return 1 if summary['failed'] else 0 + + +def _load_one_project( + args: argparse.Namespace, project: str, cache_path: Optional[str] = None +) -> int: + from pluto.migrate.loader import PlutoLoader + + input_dir = getattr(args, 'input', None) or args.output + summary = PlutoLoader( + input_dir=input_dir, + 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, + cleanup=getattr(args, 'cleanup', False), + projects=[project], + cache_path=cache_path, + ).load() + if not args.dry_run: + print( + f'[{project}] load: {summary["loaded"]} loaded, ' + f'{summary["skipped"]} skipped, {len(summary["failed"])} failed' + ) + for failure in summary['failed']: + print(f' [{project}] failed {failure["run_id"]}: {failure["error"]}') + return 1 if summary['failed'] else 0 + + +def _all_one_project( + args: argparse.Namespace, project: str, cache_path: Optional[str] = None +) -> int: + """Export + load one project concurrently (upload while downloading).""" + from pluto.migrate.loader import PlutoLoader + + result: Dict[str, int] = {} + + def _export_worker() -> None: + try: + result['code'] = _export_one_project(args, project) + except Exception as e: # a crashed export must not read as success (code 0) + print(f'[{project}] export crashed: {type(e).__name__}: {e}') + result['code'] = 2 + + export_thread = threading.Thread(target=_export_worker, daemon=True) + export_thread.start() + + loaded_total = 0 + failed: List[Dict[str, str]] = [] + # A run that failed once is attempted at most once: skip it in later passes + # (skip_run_ids below) rather than re-running the identical staged data every + # poll. That both stops the every-poll re-attempt and keeps `failed` accurate + # — a skipped run can't later succeed, so it never lingers as a false failure. + failed_ids: set = set() + + def _load_pass() -> None: + nonlocal loaded_total + summary = PlutoLoader( + input_dir=args.output, + dest_project=args.dest_project, + flush_every=args.flush_every, + max_pending=args.max_pending, + dry_run=False, + run_ids=getattr(args, 'run_ids', None), + skip_run_ids=list(failed_ids), + force_resume=args.force_resume, + cleanup=getattr(args, 'cleanup', False), + projects=[project], + cache_path=cache_path, + ).load() + loaded_total += summary['loaded'] + for f in summary['failed']: + if f['run_id'] not in failed_ids: # count each failing run once + failed_ids.add(f['run_id']) + failed.append(f) + + while export_thread.is_alive(): + _load_pass() + export_thread.join(timeout=_ALL_POLL_SECONDS) + _load_pass() + + export_code = result.get('code', 0) + if export_code == 2: + return 2 + print(f'[{project}] all: {loaded_total} loaded, {len(failed)} failed') + for failure in failed: + print(f' [{project}] failed {failure["run_id"]}: {failure["error"]}') + return max(export_code, 1 if failed else 0) + + +# --------------------------------------------------------------------------- # +# Project resolution + orchestration +# --------------------------------------------------------------------------- # +def _apply_exclude(projects: List[str], exclude: Optional[List[str]]) -> List[str]: + excl = set(exclude or []) + return [p for p in projects if p not in excl] + + +def _resolve_export_projects(args: argparse.Namespace) -> List[str]: + """Which wandb projects to export: explicit --project list, else every + project under the entity. Minus --exclude.""" + if args.projects: + projects = list(args.projects) + else: + from pluto.migrate.wandb_export import list_wandb_projects + + projects = list_wandb_projects(args.entity, args.wandb_api_key) + return _apply_exclude(projects, args.exclude) + + +def _resolve_load_projects(args: argparse.Namespace) -> List[str]: + """Which staged projects to load, from the staging layout + input_dir/{entity}/{project}/runs/. Minus --exclude, filtered by --project.""" + input_dir = Path(getattr(args, 'input', None) or args.output) + staged = sorted({d.name for d in input_dir.glob('*/*') if (d / 'runs').is_dir()}) + if args.projects: + want = set(args.projects) + staged = [p for p in staged if p in want] + return _apply_exclude(staged, args.exclude) + + +def _cache_path_for(args: argparse.Namespace, project: str) -> str: + """Per-project resume ledger so parallel loaders don't clobber one file.""" + base = Path(getattr(args, 'input', None) or args.output) + return str(base / f'loaded_runs.{project}.json') + + +def _run_over_projects( + worker: Callable[..., int], + args: argparse.Namespace, + projects: List[str], + per_project_cache: bool, +) -> int: + """Run ``worker`` for each project — in-process for one, else serial (workers<=1) + or across a process pool (workers>1). ProcessPoolExecutor because pluto.init + uses global state and can't be run concurrently in one process.""" + if not projects: + print('migrate: no matching projects') + return 0 + + def _safe(p: str, c: Optional[str]) -> int: + try: + return worker(args, p, c) + except Exception as e: # clean message, not a raw traceback + print(f'[{p}] worker crashed: {type(e).__name__}: {e}') + return 2 + + if len(projects) == 1: + return _safe(projects[0], None) + + workers = max(1, getattr(args, 'workers', 1)) + # One process per project; more workers than projects can't fan out further. + effective = min(workers, len(projects)) + note = ( + f'workers={effective}' + if effective == workers + else f'workers={effective} (requested {workers}, capped to {len(projects)} ' + 'projects — each project runs in one process)' + ) + print(f'migrate: {len(projects)} projects [{", ".join(projects)}]; {note}') + + def cache(p: str) -> Optional[str]: + return _cache_path_for(args, p) if per_project_cache else None + + codes: List[int] = [] + if effective <= 1: + for p in projects: + codes.append(_safe(p, cache(p))) + return max(codes) + + # max_tasks_per_child (a fresh process per project, so no pluto global state + # lingers between projects) is Python 3.11+. On 3.10 the pool reuses workers, + # which is fine — the loader already does many init/finish cycles per process. + pool_kwargs: Dict[str, Any] = {'max_workers': effective} + if sys.version_info >= (3, 11): + pool_kwargs['max_tasks_per_child'] = 1 + with ProcessPoolExecutor(**pool_kwargs) as pool: + futures = {pool.submit(worker, args, p, cache(p)): p for p in projects} + for fut in as_completed(futures): + p = futures[fut] + try: + codes.append(fut.result()) + except Exception as e: # a crashed worker must not sink the rest + print(f'[{p}] worker crashed: {type(e).__name__}: {e}') + codes.append(2) + return max(codes) if codes else 0 + + +def _reject_dest_project_for_many(projects: List[str], args) -> bool: + if len(projects) > 1 and getattr(args, 'dest_project', None): + print( + 'error: --dest-project cannot be used when migrating multiple projects ' + '(each project keeps its own name). Use a single --project to rename.' + ) + return True + return False + + +def cmd_migrate(args: argparse.Namespace) -> int: + if args.action == 'export': + try: + projects = _resolve_export_projects(args) + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + return _run_over_projects( + _export_one_project, args, projects, per_project_cache=False + ) + + if args.action == 'load': + try: + from pluto.migrate.loader import PlutoLoader # noqa: F401 + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + projects = _resolve_load_projects(args) + if _reject_dest_project_for_many(projects, args): + return 2 + # Common case: one loader over the whole (optionally filtered) input — + # handles all staged projects serially. Only split into parallel + # per-project loaders when there's real parallelism to exploit. + if getattr(args, 'workers', 1) > 1 and len(projects) > 1: + return _run_over_projects( + _load_one_project, args, projects, per_project_cache=True + ) + 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 `load --dry-run`.' + ) + return 2 + try: + projects = _resolve_export_projects(args) + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + if _reject_dest_project_for_many(projects, args): + return 2 + return _run_over_projects( + _all_one_project, args, projects, per_project_cache=True + ) + + raise AssertionError(f'unknown action {args.action!r}') + + +def _run_load(args: argparse.Namespace) -> int: + """Single loader over the whole input (all staged projects, optionally + filtered by --project/--exclude). Backward-compatible default load path.""" + try: + from pluto.migrate.loader import PlutoLoader + except ImportError as e: + print(f'{_INSTALL_HINT} ({e})') + return 2 + + summary = PlutoLoader( + input_dir=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, + cleanup=getattr(args, 'cleanup', False), + projects=getattr(args, 'projects', None), + exclude_projects=getattr(args, 'exclude', None), + ).load() + if not args.dry_run: + 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 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) + 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..726daef --- /dev/null +++ b/pluto/migrate/loader.py @@ -0,0 +1,945 @@ +""" +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. ``loaded_runs.json`` records +finished loads so re-runs skip them; a run that already exists on the +server but isn't in the local cache is skipped by default rather than +re-replayed, because re-replaying would duplicate media. Pass +``force_resume`` to intentionally resume and re-replay such a run +(metric points carry identical staged timestamps, so the backend's +replace-by-time dedup keeps metrics safe; media may duplicate). + +Each run is loaded independently: a single run failing (unreadable +manifest, init error, replay error, dead sync process) is recorded in +``failed`` and never aborts the rest of the batch, so a re-run retries +only the runs that did not finish. +""" + +from __future__ import annotations + +import json +import logging +import shutil +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import httpx + +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 + +# After finishing a run we read its status back and, on a mismatch, heal it by +# resume+finish. This bounds how many heal attempts we make before giving up and +# reporting the run failed (so a later pass retries it rather than looping). +_VERIFY_MAX_ATTEMPTS = 3 + +# 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, +} + + +def _resolve_within(run_dir: Path, rel: Optional[str]) -> Optional[Path]: + """Resolve a staged file path, rejecting anything outside ``run_dir``. + + A malicious or corrupt part row could carry an absolute path or a + ``..`` sequence in ``file_value``; without this guard the loader would + happily read and upload arbitrary host files (e.g. ``/etc/passwd``). + Symlinks are resolved too, so a symlink pointing outside is rejected. + Returns None for empty/out-of-bounds paths (caller treats as missing). + """ + if not rel: + return None + base = run_dir.resolve() + candidate = (run_dir / rel).resolve() + try: + candidate.relative_to(base) + except ValueError: + logger.warning(f'{tag}: refusing staged file outside run dir: {rel!r}') + return None + return candidate + + +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, + skip_run_ids: Optional[List[str]] = None, + force_resume: bool = False, + stall_timeout: float = 600.0, + cleanup: bool = False, + projects: Optional[List[str]] = None, + exclude_projects: Optional[List[str]] = None, + cache_path: Optional[Union[str, Path]] = None, + ) -> None: + self.input_dir = Path(input_dir) + # Resume ledger. Overridable so parallel per-project loaders each get + # their own file (one shared ledger would race under concurrent writes). + self.cache_path = ( + Path(cache_path) + if cache_path is not None + else self.input_dir / LOADED_CACHE_FILENAME + ) + 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 + # Runs to skip outright (already attempted-and-failed in an earlier pass + # of the same `all` invocation). The loader only ever sees a run once + # its export is complete, so a load failure is genuine, not "needs more + # time" — re-attempting the identical staged data every poll can't help + # and risks duplicating media on a mid-replay resume. Skipping such runs + # for the rest of the run makes each failure at-most-once; the user + # re-runs `all`/`load` to retry (in-progress runs resume from the ledger). + self.skip_run_ids = set(skip_run_ids) if skip_run_ids else None + # Which staged projects to load: include-list (None = all) minus excludes. + self.projects = set(projects) if projects else None + self.exclude_projects = set(exclude_projects) if exclude_projects else set() + self.force_resume = force_resume + self.stall_timeout = stall_timeout + # Delete each run's staged files once it is confirmed loaded, so a large + # migration doesn't keep a full duplicate copy on local disk. Peak disk + # then bounds to the un-loaded backlog rather than the whole export. + self.cleanup = cleanup + + def load(self) -> Dict[str, Any]: + """Load all staged runs. Returns {'loaded', 'skipped', 'failed'}. + + Each run is isolated: any failure (unreadable manifest, init error, + replay error) is recorded in ``failed`` and the batch continues. + """ + loaded, skipped, would_load = 0, 0, 0 + failed: List[Dict[str, str]] = [] + cache = LoadedCache(self.cache_path) + + for run_dir in self._discover_runs(): + # A truncated/hand-edited run.json must not sink the whole batch. + try: + manifest = read_json(run_dir / 'run.json') + run_id = manifest['run_id'] + external_id = ( + f'wandb::{manifest["entity"]}/{manifest["project"]}/{run_id}' + ) + except Exception as e: + logger.error(f'{tag}: unreadable manifest in {run_dir}: {e}') + failed.append( + {'run_id': run_dir.name, 'error': f'{type(e).__name__}: {e}'} + ) + continue + + if self.run_ids is not None and run_id not in self.run_ids: + continue + # Already attempted-and-failed in an earlier pass of this same `all` + # run: don't re-attempt (see skip_run_ids note in __init__). + if self.skip_run_ids is not None and run_id in self.skip_run_ids: + continue + # Key the load cache on (run, destination project) so loading the + # same export into a *different* Pluto project isn't wrongly skipped. + dest_project = self.dest_project or manifest['project'] + cache_key = f'{external_id}@@{dest_project}' + if cache.is_loaded(cache_key) and not self.force_resume: + logger.info( + f'{tag}: {external_id} already loaded into {dest_project!r}, ' + 'skipping' + ) + skipped += 1 + continue + if self.dry_run: + self._print_dry_run(run_dir, manifest, external_id) + would_load += 1 + continue + + op = None + try: + try: + op = self._init_run(manifest, external_id, run_dir) + # Remember, before replaying, that we created this run. If + # replay crashes now, a later re-run sees the in_progress + # marker and resumes to complete it (rather than skipping). + cache.mark_in_progress(cache_key) + except RunExistsError: + # The run exists server-side but isn't marked done here. + if self.force_resume or cache.is_in_progress(cache_key): + # We started this run before and it didn't finish (crash + # mid-replay), or the user forced it: resume and + # re-replay to complete it. Metric points dedup by their + # identical staged timestamps; media sent before the + # crash may duplicate. + logger.warning( + f'{tag}: {external_id} exists on server but was not ' + 'finished; resuming to complete it (media may ' + 'duplicate)' + ) + cache.mark_in_progress(cache_key) + op = self._init_run(manifest, external_id, run_dir, resume=True) + else: + # Exists but we never started it here (e.g. loaded on + # another machine): skip re-replay to avoid duplicating + # media. BUT the create-with-existing above already + # reopened it to RUNNING server-side (the DDP-style + # "create an existing run = resume it" path). Left as-is + # the run is stuck RUNNING with a now() finish time, so + # re-attach and finish() — no replay — to restore its + # terminal status + historical statusUpdated, then skip. + logger.info( + f'{tag}: {external_id} already exists on server; ' + 'restoring its finished state and skipping (pass ' + '--force-resume to resume and re-replay)' + ) + try: + restore = self._init_run( + manifest, external_id, run_dir, resume=True + ) + restore.finish( + code=0 if manifest.get('state') == 'finished' else 1 + ) + if restore._status_update_error is not None: + # finish() couldn't confirm the terminal status; + # re-raise into the handler below so we record a + # failure instead of marking a still-RUNNING run + # loaded (and skipping it forever). + raise restore._status_update_error + rcode = 0 if manifest.get('state') == 'finished' else 1 + if not self._verify_and_heal( + restore, manifest, external_id, run_dir, rcode + ): + raise RuntimeError( + 'terminal status not confirmed after ' + f'{_VERIFY_MAX_ATTEMPTS} heals' + ) + except Exception as e: + # Restore failed: the run is still RUNNING server-side + # and its terminal status was never written. Do NOT + # mark it loaded — that would strand it RUNNING and + # skip it forever. Record a failure so it's reported + # and retried on the next `all`/`load` invocation. + logger.warning( + f'{tag}: could not restore terminal status for ' + f'{external_id}: {e}' + ) + failed.append( + { + 'run_id': run_id, + 'error': f'restore failed: {type(e).__name__}: {e}', + } + ) + continue + cache.mark_loaded(cache_key, {'note': 'existed-on-server'}) + skipped += 1 + continue + + self._replay_run(run_dir, op) + code = 0 if manifest.get('state') == 'finished' else 1 + op.finish(code=code) + if op._status_update_error is not None: + # finish() replayed all data but could NOT confirm the run's + # terminal status on the server (a dropped connection that + # outlasted retries). Marking it loaded here would strand it + # RUNNING/FAILED and skip it forever. Record a failure so the + # next `all`/`load` pass heals it via the RunExistsError -> + # restore path above. + err = op._status_update_error + logger.error( + f'{tag}: {external_id} replayed but terminal status was ' + f'not confirmed: {type(err).__name__}: {err}' + ) + failed.append( + { + 'run_id': run_id, + 'error': ( + f'status unconfirmed: {type(err).__name__}: {err}' + ), + } + ) + continue + # The finish reported success, but a stale-run reap can have + # marked this back-dated import FAILED and silently rejected our + # COMPLETED. Confirm it stuck (and heal a false reap) before we + # cache it as loaded — otherwise a re-run skips a stranded run. + if not self._verify_and_heal(op, manifest, external_id, run_dir, code): + logger.error( + f'{tag}: {external_id} not confirmed COMPLETED after ' + f'{_VERIFY_MAX_ATTEMPTS} heal attempts (stale-run reap?)' + ) + failed.append( + { + 'run_id': run_id, + 'error': ( + 'terminal status not confirmed after ' + f'{_VERIFY_MAX_ATTEMPTS} heals' + ), + } + ) + continue + cache.mark_loaded(cache_key, {'pluto_run_id': op.settings._op_id}) + loaded += 1 + logger.info(f'{tag}: loaded {external_id}') + if self.cleanup: + # Reclaim disk now that the run is safely loaded (and + # recorded in loaded_runs.json, so it's still skipped on a + # re-run even though its staged files are gone). + shutil.rmtree(run_dir, ignore_errors=True) + 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}'}) + if op is not None: + try: + op.finish(code=1) + except Exception: + pass + + if self.dry_run: + print( + f'[dry-run] would load {would_load} run(s); ' + f'{skipped} already loaded (would skip)' + ) + return {'loaded': loaded, 'skipped': skipped, 'failed': failed} + + def _discover_runs(self) -> List[Path]: + # Layout: input_dir/{entity}/{project}/runs/{run_id}. Filter by the + # project component so a scoped load (--project / --exclude) only picks + # up the requested projects. + 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() + and self._project_in_scope(d.parent.parent.name) + ) + + def _project_in_scope(self, project: str) -> bool: + if project in self.exclude_projects: + return False + return self.projects is None or project in self.projects + + def _read_run_status(self, op: Any) -> Optional[str]: + """Read the run's current server-side status, or None if unreadable. + + Confirms a finished import actually landed as COMPLETED. The server's + stale-run monitor can reap a freshly-created, back-dated import to + FAILED in a sub-second race; the client's COMPLETED is then silently + rejected (HTTP 200, no effect). Only a read-back catches that. Returns + None on any error so the caller degrades to best-effort rather than + failing a run because we merely couldn't verify it. + """ + s = getattr(op, 'settings', None) + run_id = getattr(s, '_op_id', None) + url_api = getattr(s, 'url_api', None) + auth = getattr(s, '_auth', None) + if run_id is None or not url_api or not auth: + return None + try: + r = httpx.get( + f'{url_api}/api/runs/details/{run_id}', + headers={'Authorization': f'Bearer {auth}'}, + timeout=15.0, + ) + if r.status_code == 200: + return r.json().get('status') + logger.debug( + '%s: status read-back for run %s returned HTTP %s', + tag, + run_id, + r.status_code, + ) + except Exception as e: + logger.debug('%s: status read-back failed for run %s: %s', tag, run_id, e) + return None + + def _verify_and_heal( + self, + op: Any, + manifest: Dict[str, Any], + external_id: str, + run_dir: Path, + code: int, + ) -> bool: + """Confirm the run reached its intended terminal status; heal a false reap. + + The stale-run monitor can mark a just-created back-dated import FAILED in + a race, after which the client's COMPLETED is silently rejected. We read + the status back; on a definite mismatch we resume+finish — reopening the + run to RUNNING and finishing again, which sidesteps the terminal-status + precedence guard the way a plain re-finish cannot, and replays nothing so + no data duplicates. Bounded to ``_VERIFY_MAX_ATTEMPTS``. + + Only a ``finished`` import (code 0) is at risk: a wandb failure that + lands FAILED is already correct, so we skip it. Returns True once + confirmed COMPLETED, or when the status simply can't be read (best-effort + — don't fail a run over our inability to verify). Returns False only on a + confirmed mismatch we could not heal, so the caller records it failed and + a later pass retries rather than caching a stranded run as loaded. + """ + if code != 0: + return True + expected = 'COMPLETED' + # Set once we heal and the heal's own finish can't confirm its terminal + # status (or the heal raises). While set, an unreadable final status must + # NOT be treated as best-effort success — the run may be stuck RUNNING, so + # we report it failed and let a re-run retry rather than cache it loaded. + unconfirmed = False + for _ in range(_VERIFY_MAX_ATTEMPTS): + status = self._read_run_status(op) + if status == expected: + return True + if status is None: + # Couldn't read (transient blip or endpoint unavailable): retry + # the read; do NOT heal a run we can't prove is broken. + continue + logger.warning( + '%s: %s came back %r, expected COMPLETED (likely a stale-run ' + 'reap); healing via resume+finish', + tag, + external_id, + status, + ) + try: + op = self._init_run(manifest, external_id, run_dir, resume=True) + op.finish(code=0) + # finish() records a dropped terminal-status update without + # raising (same as the main/restore paths), so check it + # explicitly rather than assuming the heal stuck. + unconfirmed = op._status_update_error is not None + except Exception as e: + unconfirmed = True + logger.warning( + '%s: heal attempt for %s failed: %s', tag, external_id, e + ) + final = self._read_run_status(op) + if final == expected: + return True + if final is None and not unconfirmed: + # Either we never saw a real mismatch (the status endpoint is just + # unavailable) or the last heal's finish confirmed — don't fail a run + # over our own inability to verify. + logger.warning( + '%s: could not verify terminal status for %s; proceeding as ' + 'loaded (best-effort)', + tag, + external_id, + ) + return True + return False + + def _init_run( + self, + manifest: Dict[str, Any], + external_id: str, + run_dir: Path, + resume: Optional[bool] = None, + ) -> Any: + tags = list(manifest.get('tags') or []) + if 'import:wandb' not in tags: + tags.append('import:wandb') + # Group sweep runs the same way native pluto.sweep does: tag sweep:. + sweep = manifest.get('sweep') + if isinstance(sweep, dict) and sweep.get('id'): + sweep_tag = f'sweep:{sweep["id"]}' + if sweep_tag not in tags: + tags.append(sweep_tag) + compat: Dict[str, Any] = { + 'createdAt': manifest.get('createdAt'), + 'updatedAt': manifest.get('updatedAt'), + } + # Forward the original run's own metadata (git/OS/GPU/python/args) as + # systemMetadata so repro context survives the migration. Only set when + # present so normal runs (empty compat) are unaffected. + if manifest.get('metadata') is not None: + compat['systemMetadata'] = manifest['metadata'] + settings: Dict[str, Any] = { + 'compat': compat, + # 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'), + # Custom-chart (wandb.plot.*) panel specs recovered by the + # exporter: each binds a Vega preset to a migrated backing + # table. Forwarded here so the Pluto side can rebuild the panels. + 'custom_charts': self._read_custom_charts(run_dir), + # Sweep membership + search-space config (the run is also tagged + # sweep: above), so the sweep survives the migration. + 'sweep': manifest.get('sweep'), + }.items() + if v + } + if wandb_block: + op.update_config({'wandb': wandb_block}) + return op + + @staticmethod + def _read_custom_charts(run_dir: Path) -> Optional[List[Dict[str, Any]]]: + """Load the exporter's staged custom-chart panel specs, if any.""" + path = run_dir / 'custom_charts.json' + if not path.exists(): + return None + try: + panels = read_json(path).get('panels') or None + except Exception as e: + logger.warning( + f'{tag}: could not read {path.name}: {type(e).__name__}: {e}' + ) + return None + return panels + + @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]] = [] + # Categorical/status series buffered per attribute_path for the whole + # run, then sent once at the end in a single request. Each entry is a + # list of (step, timestamp_ms, value). + string_series: Dict[str, List[Tuple[int, int, str]]] = {} + + 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) + + # Media/console/artifact rows enqueue outside the scalar-metric flush, + # so they need their own backpressure cadence — otherwise a run that is + # mostly images/logs (few scalars) never triggers _wait_for_backpressure + # and can balloon the sync queue / staged files past max_pending. + nonmetric_since_check = 0 + + def note_nonmetric(n: int = 1) -> None: + nonlocal nonmetric_since_check + nonmetric_since_check += n + if nonmetric_since_check >= self.flush_every: + nonmetric_since_check = 0 + self._wait_for_backpressure(op) + + # Media grouping: consecutive media rows sharing (name, step, timestamp) + # are one logged batch (e.g. wandb.log({"gallery": [img0, img1, ...]})). + # Logging them as a single list makes op.log assign each item its + # sampleIndex (0,1,2,...), which the server uses to preserve logged + # order — one op.log per row would make every item sampleIndex 0. + media_key: Optional[Tuple[str, int, int]] = None + media_items: List[Any] = [] + + def flush_media() -> None: + nonlocal media_key, media_items + if media_key is not None and media_items: + name, step, ts_ms = media_key + value = media_items[0] if len(media_items) == 1 else media_items + op.log({name: value}, step=step, timestamp=ts_ms / 1000) + note_nonmetric(len(media_items)) + media_key, media_items = None, [] + + 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'): + flush_media() + 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': + # One malformed media/histogram row must not abort the whole + # run's replay — skip it and keep the rest of the run's data. + try: + value = self._build_media_value(run_dir, row) + except Exception as e: + logger.warning( + f'{tag}: skipping bad media row ' + f'{row.get("attribute_path")!r} @ step ' + f'{row.get("step")}: {type(e).__name__}: {e}' + ) + value = None + if value is not None: + mkey = ( + row['attribute_path'], + row['step'], + row['timestamp_ms'], + ) + if mkey != media_key: + flush_media() + media_key = mkey + media_items.append(value) + continue + flush_media() # any non-media row ends the current media batch + if 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) + note_nonmetric(len(console_lines)) + console_lines = [] + elif attr_type == 'string_series': + string_series.setdefault(row['attribute_path'], []).append( + (row['step'], row['timestamp_ms'], row['string_value'] or '') + ) + elif attr_type == 'artifact': + try: + self._replay_artifact(run_dir, op, row) + except Exception as e: + logger.warning( + f'{tag}: skipping bad artifact row ' + f'{row.get("attribute_path")!r} @ step ' + f'{row.get("step")}: {type(e).__name__}: {e}' + ) + note_nonmetric() + + flush_media() + close_group() + flush_pending(force=True) + if console_lines: + op._log_console(console_lines) + self._wait_for_backpressure(op) + if string_series: + self._send_string_series(op, string_series) + + def _send_string_series( + self, op: Any, series: Dict[str, List[Tuple[int, int, str]]] + ) -> None: + """Send categorical/status series to the data-ingest endpoint. + + String history points have no numeric/media home; Pluto stores them as + a ``string-series`` dataType in the data ingest (rendered as a state + timeline). This posts NDJSON directly to ``url_data`` — mirroring the + sync process's data upload — rather than routing through the scalar or + media client paths. + + Every string series is sent regardless of cardinality (no distinct-value + guard). Failures are non-fatal — the run's scalars/media are already + loaded and must not be lost to an ingest hiccup here. + """ + s = op.settings + if not s.url_data or not s._op_id: + return + renderable = series + if not renderable: + return + + common = { + 'Authorization': f'Bearer {s._auth}', + 'User-Agent': str(s.tag), + 'X-Run-Id': str(s._op_id), + 'X-Run-Name': str(s._op_name or ''), + 'X-Project-Name': str(s.project or ''), + } + # NDJSON: one {time, step, dataType, logName, data} object per point. + lines = [ + json.dumps( + { + 'time': ts_ms, + 'step': step, + 'dataType': 'string-series', + 'logName': name, + 'data': value, + } + ) + for name, points in renderable.items() + for step, ts_ms, value in points + ] + body = '\n'.join(lines) + '\n' + try: + # Register the log names (logType DATA) so the server indexes them, + # then ingest the points. + resp = httpx.post( + s.url_meta, + headers={**common, 'Content-Type': 'application/json'}, + content=json.dumps( + { + 'runId': s._op_id, + 'logType': 'DATA', + 'logName': list(renderable), + } + ), + timeout=30.0, + ) + resp.raise_for_status() + resp = httpx.post( + s.url_data, + headers={**common, 'Content-Type': 'application/x-ndjson'}, + content=body, + timeout=60.0, + ) + resp.raise_for_status() + except Exception as e: + logger.warning( + f'{tag}: failed to send {len(renderable)} string-series to ' + f'ingest: {type(e).__name__}: {e}' + ) + + def _image_annotations( + self, run_dir: Path, annotation_value: Optional[str] + ) -> Tuple[Optional[str], Optional[Dict[str, Any]]]: + """Resolve staged box/mask refs for an annotated image. + + ``annotation_value`` is ``{"boxes": {layer: {path,...}}, "masks": + {layer: {path,...}}}`` (wandb's refs staged by the exporter). Boxes: read + each sidecar ``.boxes2D.json`` ({box_data, class_labels}) and inline it → + the wandb-shape ``{"boxes": {layer: ...}}`` annotations string. Masks: + resolve each ``.mask.png`` to a path so ``pluto.Image`` re-uploads it (as + fileType "mask"). Returns ``(boxes_annotations_str, masks_spec)``. + """ + if not annotation_value: + return None, None + try: + refs = json.loads(annotation_value) + except (ValueError, TypeError): + return None, None + + boxes_out: Dict[str, Any] = {} + for layer, ref in (refs.get('boxes') or {}).items(): + fp = self._resolve_ref(run_dir, ref) + if fp is None: + continue + try: + content = read_json(fp) + except Exception: + continue + if isinstance(content, dict): + boxes_out[layer] = content + boxes_str = json.dumps({'boxes': boxes_out}) if boxes_out else None + + masks_spec: Dict[str, Any] = {} + for layer, ref in (refs.get('masks') or {}).items(): + fp = self._resolve_ref(run_dir, ref) + if fp is not None: + spec: Dict[str, Any] = {'path': str(fp)} + # Carry the id→name key through so the mask renders coloured + # instead of blank (the exporter recovers it from run config). + if isinstance(ref, dict) and ref.get('class_labels'): + spec['class_labels'] = ref['class_labels'] + masks_spec[layer] = spec + return boxes_str, (masks_spec or None) + + @staticmethod + def _resolve_ref(run_dir: Path, ref: Any) -> Optional[Path]: + """Resolve a staged annotation file ref ({path: 'media/...'}) to a Path.""" + rel = ref.get('path') if isinstance(ref, dict) else None + fp = _resolve_within(run_dir, f'files/{rel}') if rel else None + return fp if fp is not None and fp.exists() else None + + def _build_media_value(self, run_dir: Path, row: Dict[str, Any]) -> Optional[Any]: + """Convert one staged media row into a pluto media value. + + Returns None to skip the row (missing file, empty histogram, unknown + inline type). The caller batches consecutive same-key values so a + multi-sample step (e.g. an image gallery) keeps its logged order. + """ + name = row['attribute_path'] + step = row['step'] + media_type = row['string_value'] or '' + + if media_type.startswith('{'): # inline JSON (histogram) + payload = json.loads(media_type) + if payload.get('_type') == 'histogram': + values = payload.get('values') + bins = payload.get('bins') + if values is None: + logger.warning( + f'{tag}: histogram {name!r} @ step {step} has no ' + 'values, skipping' + ) + return None + if bins is None: + # wandb frequently stores histogram counts without bin + # edges. pluto.Histogram's pre-binned form needs + # len(edges) == len(counts) + 1, so synthesize integer + # edges — the counts are preserved, only the x-axis is + # generic. + bins = list(range(len(values) + 1)) + return pluto.Histogram([values, bins], bins=None) + return None # unknown inline media type + + path = _resolve_within(run_dir, row['file_value']) + if path is None or not path.exists(): + logger.warning( + f'{tag}: media file missing for {name!r} ' + f'({row["file_value"]}), skipping' + ) + return None + + caption = row['caption'] + if media_type == 'table-file': + table_json = read_json(path) + return pluto.Table( + data=table_json.get('data'), + columns=table_json.get('columns', []), + ) + if media_type == 'image-file': + boxes_str, masks_spec = self._image_annotations( + run_dir, row.get('annotation_value') + ) + return pluto.Image( + str(path), caption=caption, annotations=boxes_str, masks=masks_spec + ) + make = _MEDIA_LOADERS.get(media_type) + if make is not None: + return make(str(path), caption=caption) + # plotly/html/object3D/unknown -> raw artifact + return pluto.Artifact(str(path), caption=caption) + + def _replay_artifact(self, run_dir: Path, op: Any, row: Dict[str, Any]) -> None: + path = _resolve_within(run_dir, row['file_value']) + if path is None 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. + + Fails fast if the sync subprocess has died: without this the loop + would sleep out the full ``stall_timeout`` on every flush (the + pending count never drops with no uploader), turning a large replay + into hours of pure sleeping. Raising here surfaces the dead process + as a per-run failure so the batch continues with the next run. + + For an *alive but slow* uploader (unreachable server, throttling) + it stays bounded by ``stall_timeout``: it logs and moves on rather + than hanging — data stays in the sync store either way. + """ + manager = getattr(op, '_sync_manager', None) + if manager is None: + return + deadline = time.time() + self.stall_timeout + sleep_s = 0.5 + throttled = False + while True: + proc = getattr(manager, '_process', None) + if proc is not None: + # poll() -> None while alive, an int exit code once dead. + # (isinstance keeps a mocked manager from tripping this.) + code = proc.poll() + if isinstance(code, int): + raise RuntimeError( + f'sync process exited (code {code}) with data still ' + 'pending; aborting this run (data preserved in the ' + 'sync store for a later retry)' + ) + try: + pending = 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..78c8b21 --- /dev/null +++ b/pluto/migrate/schema.py @@ -0,0 +1,169 @@ +""" +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) + 'string_series', # categorical/status history point (string_value per step) +} + +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()), + # For annotated images: JSON of the raw wandb boxes/masks refs + # ({"boxes": {layer: {path,...}}, "masks": {...}}); the loader resolves + # the referenced files into the image's annotations. + pa.field('annotation_value', 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, + annotation_value: 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, + 'annotation_value': annotation_value, + } + ) + 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..38ea69d --- /dev/null +++ b/pluto/migrate/state.py @@ -0,0 +1,107 @@ +""" +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 logging +import os +from pathlib import Path +from typing import Any, Dict, Optional, Union + +logger = logging.getLogger(f'{__name__.split(".")[0]}') +tag = 'migrate' + +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 tracking each external id's load state. + + Entries are ``{status: 'in_progress' | 'done', ...}``. ``in_progress`` is + written right after a run is created server-side but before its replay + finishes; ``done`` after ``finish()`` drains. That lets a re-run tell a + run *we* started but didn't finish (resume + complete it) from one that + merely exists server-side (loaded elsewhere → skip, no media dup). + Legacy entries without a status are treated as ``done``. + """ + + def __init__(self, path: Union[str, Path]) -> None: + self._path = Path(path) + self._entries: Dict[str, Any] = {} + if self._path.exists(): + try: + data = read_json(self._path) + if isinstance(data, dict): + self._entries = data + else: + raise ValueError('cache is not a JSON object') + except (json.JSONDecodeError, ValueError, OSError) as e: + # A corrupt/empty cache must not abort the whole load. Move it + # aside and start fresh; already-loaded runs are re-detected via + # the server-side external-id collision (RunExistsError). + logger.warning( + f'{tag}: {self._path.name} unreadable ({e}); starting with ' + 'an empty load cache (backed up as .corrupt)' + ) + try: + self._path.replace(self._path.with_suffix('.corrupt')) + except OSError: + pass + + def _status(self, external_id: str) -> Optional[str]: + entry = self._entries.get(external_id) + if entry is None: + return None + if isinstance(entry, dict): + return entry.get('status', 'done') # legacy entries -> done + return 'done' + + def is_loaded(self, external_id: str) -> bool: + return self._status(external_id) == 'done' + + def is_in_progress(self, external_id: str) -> bool: + return self._status(external_id) == 'in_progress' + + def mark_in_progress(self, external_id: str) -> None: + if self.is_loaded(external_id): + return # never downgrade a completed run + self._entries[external_id] = {'status': 'in_progress'} + write_json_atomic(self._path, self._entries) + + def mark_loaded(self, external_id: str, info: Dict[str, Any]) -> None: + self._entries[external_id] = {**info, 'status': 'done'} + 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..c595392 --- /dev/null +++ b/pluto/migrate/wandb_export.py @@ -0,0 +1,1090 @@ +""" +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 +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, TypedDict, Union + +import yaml + +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' + +# A crash-truncated wandb cache parquet is 0 bytes and stays that way; an +# in-progress download is only briefly 0 bytes at creation. The empty-cache purge +# reaps only files that have been 0 bytes at least this long, so it never deletes +# another worker's active download from the shared cache under --workers>1. +_EMPTY_CACHE_STALE_SECONDS = 60 + +# wandb encodes non-finite metric points as JSON strings; map them back to the +# real floats so they migrate instead of being dropped as "text". +_NONFINITE_STRINGS = { + 'NaN': float('nan'), + 'nan': float('nan'), + 'Infinity': float('inf'), + 'inf': float('inf'), + 'Inf': float('inf'), + '-Infinity': float('-inf'), + '-inf': float('-inf'), + '-Inf': float('-inf'), +} + +# Non-numeric string history points (status labels, phase names, etc.) migrate +# as a "string-series" (rendered as a categorical/state-timeline in Pluto). +# Guard pathological values: a single point longer than this is almost never a +# real categorical state (it's a stray log line / serialized blob), so drop it +# rather than pollute the timeline. The per-key cardinality guard lives in the +# loader (it needs the whole series to count distinct values). +_STRING_SERIES_MAX_LEN = 200 + + +class _RowBase(TypedDict): + """Identity columns shared by every staged row (see schema.write_row).""" + + project_id: str + run_id: str + + +# wandb built-in custom-chart presets (panelDefId) -> a stable short name. +# Each corresponds to a Vega-Lite template the Pluto side renders; the backing +# table (tableKey) migrates as a Table and supplies the chart's data. Presets +# outside this set (user-authored Vega) are staged but flagged: the server has +# no template to rebuild them from. +_WANDB_CHART_PRESETS = { + 'wandb/bar/v0': 'bar', + 'wandb/line/v0': 'line', + 'wandb/lineseries/v0': 'lineseries', + 'wandb/scatter/v0': 'scatter', + 'wandb/pr_curve/v0': 'pr_curve', + 'wandb/roc_curve/v0': 'roc_curve', + # area-under-curve is the SDK's shared preset behind pr_curve + roc_curve. + 'wandb/area-under-curve/v0': 'area-under-curve', + 'wandb/confusion_matrix/v0': 'confusion_matrix', + 'wandb/confusion_matrix/v1': 'confusion_matrix', + 'wandb/histogram/v0': 'histogram', +} + +# 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', +} + +# wandb column-type markers for media cells. A table with any of these columns +# loses its media on migration: we stage wandb's lossy run-files table copy, +# where media cells are text placeholders ("Image", ...), and the cell media +# migrates only as the sibling run_table artifact — unlinked from the table. +_MEDIA_COLUMN_WB_TYPES = ( + 'image-file', + 'audio-file', + 'video-file', + 'html-file', + 'object3D-file', + 'molecule-file', +) + + +def _table_media_columns(artifact_dir: Path) -> List[str]: + """Names of media columns in a downloaded table artifact, if any. + + The artifact copy of a logged ``wandb.Table`` carries ``column_types`` + (the lossy run-files copy does not). A column whose type tree references a + media ``*-file`` wb_type holds media cells that don't survive migration as + media. Returns [] for plain tables and non-table artifacts. + """ + hits: List[str] = [] + for table_json in artifact_dir.rglob('*.table.json'): + try: + with open(table_json) as fh: + data = json.load(fh) + except (OSError, ValueError): + continue + column_types = data.get('column_types') + if not isinstance(column_types, dict): + continue + type_map = column_types.get('params', {}).get('type_map', {}) + if not isinstance(type_map, dict): + continue + for col, col_type in type_map.items(): + if any(t in json.dumps(col_type) for t in _MEDIA_COLUMN_WB_TYPES): + hits.append(col) + return hits + + +# 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, + system_samples: int = 100_000, + download_workers: int = 16, + ) -> 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 + # Validate up front: an unparseable --after/--before must error, not + # silently drop the filter and export the entire project. + self.after_ms = self._parse_filter_date('--after', after) + self.before_ms = self._parse_filter_date('--before', 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 + self.system_samples = system_samples + self.download_workers = max(1, download_workers) + # Coverage: what got migrated vs. dropped, so nothing is lost silently. + self._cov_migrated: Counter = Counter() # per-run (reset in _export_run) + self._cov_skipped: Counter = Counter() + self._cov_migrated_total: Counter = Counter() + self._cov_skipped_total: Counter = Counter() + + @staticmethod + def _bins_from_packed(packed: Any) -> Optional[List[float]]: + """Reconstruct histogram bin edges from wandb's compact packedBins + ({min, size, count}); returns count+1 edges, or None if unusable.""" + if not isinstance(packed, dict): + return None + try: + mn = float(packed['min']) + sz = float(packed['size']) + cnt = int(packed['count']) + except (KeyError, TypeError, ValueError): + return None + if cnt <= 0: + return None + return [mn + i * sz for i in range(cnt + 1)] + + def _migrated(self, category: str, n: int = 1) -> None: + self._cov_migrated[category] += n + + def _skipped(self, category: str, n: int = 1) -> None: + self._cov_skipped[category] += n + + @staticmethod + def _fmt_coverage(migrated: Counter, skipped: Counter) -> str: + m = ', '.join(f'{v} {k}' for k, v in sorted(migrated.items())) or 'nothing' + line = f'migrated: {m}' + if skipped: + s = ', '.join(f'{v} {k}' for k, v in sorted(skipped.items())) + line += f'; NOT migrated: {s}' + return line + + @staticmethod + def _parse_filter_date(flag: str, value: Optional[str]) -> Optional[int]: + """None when the flag was not supplied; raise when it was supplied but + cannot be parsed (so a typo like ``2024/01/01`` fails loudly instead of + being silently ignored).""" + if value is None: + return None + ms = parse_iso_ms(value) + if ms is None: + raise ValueError( + f'{flag}: could not parse date {value!r}; use ISO-8601, ' + 'e.g. 2024-01-01 or 2024-01-01T00:00:00Z' + ) + return ms + + @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 _purge_empty_wandb_cache(self) -> int: + """Delete 0-byte parquet files from wandb's local cache. + + A crash (or OOM) mid-download truncates a run-history parquet to 0 + bytes; wandb then reuses that empty file on every later read and the + export fails with 'Parquet file too small. Size is 0' — permanently, + until the file is removed. A valid parquet is never empty, so deleting + these is always safe: wandb re-downloads the real data on next read. + Returns the number removed.""" + env = os.environ.get('WANDB_CACHE_DIR') + cache = Path(env) if env else Path.home() / '.cache' / 'wandb' + removed: List[str] = [] + now = time.time() + try: + for p in cache.rglob('*.parquet'): + try: + st = p.stat() + age = now - st.st_mtime + if st.st_size == 0 and age > _EMPTY_CACHE_STALE_SECONDS: + p.unlink() + removed.append(p.name) + except OSError: + pass + except Exception as e: # cache dir absent/unreadable — nothing to do + logger.debug(f'{tag}: wandb cache sweep skipped: {e}') + if removed: + # WARNING, not INFO: this is a non-routine recovery event (a prior + # crash left junk in wandb's cache) the user should see. Name the + # files so it's clear exactly what was cleared + re-downloaded. + shown = ', '.join(removed[:10]) + if len(removed) > 10: + shown += f' (+{len(removed) - 10} more)' + logger.warning( + f'{tag}: cleared {len(removed)} empty (crash-truncated) wandb ' + f'cache file(s); wandb will re-download them: {shown}' + ) + return len(removed) + + def export(self) -> Dict[str, Any]: + """Export all matching runs. Returns {'exported', 'skipped', 'failed'}.""" + exported, skipped = 0, 0 + failed: List[Dict[str, str]] = [] + + # A prior crash can poison wandb's cache with 0-byte parquets that make + # every history read fail; clear them up front so a crash never blocks + # a later export. + self._purge_empty_wandb_cache() + 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 + elif self.after_ms is not None or self.before_ms is not None: + # A date window was requested but this run has no usable + # created_at; don't silently widen the export past what was + # asked — skip it, loudly. + logger.warning( + f'{tag}: {run.id} has no parseable created_at; excluding it ' + 'from the date-filtered export' + ) + 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 + + # Retry once: recover from a crash-truncated cache read (purge the + # empty parquet so wandb re-downloads) or a transient network blip, + # instead of losing the run. _export_run rebuilds its tmp dir each + # attempt, so a retry is clean. + last_err: Optional[Exception] = None + for attempt in range(2): + try: + self._export_run(run, run_dir) + exported += 1 + logger.info( + f'{tag}: exported {run.id} ({run.name})' + + (' (on retry)' if attempt else '') + ) + last_err = None + break + except Exception as e: # keep going: one bad run must not stop all + last_err = e + if attempt == 0: + self._purge_empty_wandb_cache() + logger.warning( + f'{tag}: export of {run.id} failed ({e}); retrying' + ) + if last_err is not None: + logger.error(f'{tag}: export failed for {run.id}: {last_err}') + failed.append( + { + 'run_id': run.id, + 'error': f'{type(last_err).__name__}: {last_err}', + } + ) + + coverage = { + 'migrated': dict(self._cov_migrated_total), + 'not_migrated': dict(self._cov_skipped_total), + } + total_cov = self._fmt_coverage( + self._cov_migrated_total, self._cov_skipped_total + ) + log = logger.warning if self._cov_skipped_total else logger.info + log(f'{tag}: coverage across {exported} run(s) — {total_cov}') + + summary = { + 'exported': exported, + 'skipped': skipped, + 'failed': failed, + 'coverage': coverage, + } + # Per-project path: a shared output_dir/manifest.json would be clobbered + # (and its .tmp raced) when multiple projects export concurrently. + manifest_dir = self.output_dir / self.entity / self.project + manifest_dir.mkdir(parents=True, exist_ok=True) + write_json_atomic( + manifest_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) + + self._cov_migrated = Counter() + self._cov_skipped = Counter() + 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) + + # Custom-chart (wandb.plot.*) panel definitions live in the raw + # config.yaml (downloaded above), not the staging rows — recover them + # once the file is on disk. + self._export_custom_charts(run, tmp_dir) + + # Run-level context that has no home in the staging schema (sweep + # membership, input-artifact lineage) — flag it so it's not lost silently. + self._flag_run_level_omissions(run) + + # Coverage: one clear line per run of what was migrated vs. dropped + # (dropped items also surface at WARNING so they're never silent). + cov = self._fmt_coverage(self._cov_migrated, self._cov_skipped) + logger.info(f'{tag}: {run.id} coverage — {cov}') + if self._cov_skipped: + dropped = ', '.join( + f'{v} {k}' for k, v in sorted(self._cov_skipped.items()) + ) + logger.warning(f'{tag}: {run.id} NOT migrated: {dropped}') + self._cov_migrated_total.update(self._cov_migrated) + self._cov_skipped_total.update(self._cov_skipped) + + 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), + 'sweep': self._sweep_block(run), + }, + ) + + @staticmethod + def _sweep_block(run: Any) -> Optional[Dict[str, Any]]: + """Capture the run's wandb sweep — id, name, and search-space config. + + Runs that belong to a sweep carry a Sweep object; migrating it lets Pluto + group the runs (tag ``sweep:``) and keep the search space, mirroring + the native ``pluto.sweep`` data model. Best-effort: any API hiccup + returns None rather than failing the export. + """ + try: + sweep = getattr(run, 'sweep', None) + if sweep is None: + return None + sweep_id = getattr(sweep, 'id', None) + if not sweep_id: + return None + block: Dict[str, Any] = {'id': sweep_id} + name = getattr(sweep, 'name', None) + if name: + block['name'] = name + config = getattr(sweep, 'config', None) + if isinstance(config, dict) and config: + block['config'] = config + return block + except Exception: + return None + + @staticmethod + def _chart_table_key(panel_config: Dict[str, Any]) -> Optional[str]: + """Pull a custom chart's backing-table key from its userQuery. + + wandb encodes it as a ``summaryTable`` field carrying a ``tableKey`` + arg: ``userQuery.queryFields[].fields[].{name: summaryTable, + args: [{name: tableKey, value: }]}``. That value is the table's + log name, which migrates as a Table and supplies the chart data. + """ + user_query = panel_config.get('userQuery') or {} + for query_field in user_query.get('queryFields') or []: + for field in query_field.get('fields') or []: + if field.get('name') != 'summaryTable': + continue + for arg in field.get('args') or []: + if arg.get('name') == 'tableKey': + return arg.get('value') + return None + + def _export_custom_charts(self, run: Any, tmp_dir: Path) -> None: + """Recover wandb custom-chart (``wandb.plot.*``) panel definitions. + + The panels live in the run's raw config under ``_wandb.value.visualize`` + — the public API strips ``_wandb``, so we read the downloaded + ``config.yaml`` instead. Each panel binds a built-in Vega preset + (``panelDefId``) to a backing table via ``fieldSettings`` column + mappings. We stage a normalized ``custom_charts.json``; the loader + forwards it so the Pluto side can rebuild the panel from preset + + migrated table. The raw Vega spec is server-side, not reconstructed here. + """ + config_path = tmp_dir / 'files' / 'config.yaml' + if not config_path.exists(): + return # --no-files/--no-console, or a run without config.yaml + try: + with open(config_path) as f: + raw = yaml.safe_load(f) or {} + visualize = ((raw.get('_wandb') or {}).get('value') or {}).get( + 'visualize' + ) or {} + except Exception as e: + logger.warning( + f'{tag}: {run.id} could not parse config.yaml for custom ' + f'charts: {type(e).__name__}: {e}' + ) + return + if not isinstance(visualize, dict) or not visualize: + return + + panels = [] + for key, spec in visualize.items(): + if not isinstance(spec, dict): + continue + panel_config = spec.get('panel_config') or {} + panel_def = panel_config.get('panelDefId') + preset = ( + _WANDB_CHART_PRESETS.get(panel_def) + if isinstance(panel_def, str) + else None + ) + table_key = self._chart_table_key(panel_config) + field_settings = panel_config.get('fieldSettings') + fields = ( + {k: v for k, v in field_settings.items() if v is not None} + if isinstance(field_settings, dict) + else {} + ) + string_settings = panel_config.get('stringSettings') + title = ( + string_settings.get('title') + if isinstance(string_settings, dict) + else None + ) + panels.append( + { + 'key': key, + 'panelDefId': panel_def, + 'preset': preset, # None for user-authored (non-preset) Vega + 'title': title, + # Whole stringSettings dict (title + axis titles, e.g. + # x-axis-title/y-axis-title) so the renderer can substitute + # ${string:...}; falls back to raw column names if absent. + 'strings': string_settings + if isinstance(string_settings, dict) + else {}, + 'tableKey': table_key, # backing table's log name + 'fields': fields, + 'specLang': 'vega-lite' if preset else 'vega', + } + ) + if preset and table_key: + self._migrated('custom-chart') + else: + # Unknown preset or unresolved backing table: staged for + # reference, but the server has no template to rebuild it from. + self._skipped('custom-chart-unsupported') + + if panels: + write_json_atomic(tmp_dir / 'custom_charts.json', {'panels': panels}) + + def _row_base(self, run: Any) -> '_RowBase': + return {'project_id': self.project_path, 'run_id': run.id} + + def _mask_class_labels(self, run: Any) -> Dict[str, Dict[str, str]]: + """The id→name key for segmentation masks, from the run's config. + + wandb does NOT put ``class_labels`` in a mask's media descriptor (the + thing ``scan_history`` returns). It stores them in the run config under + ``_wandb.value['mask/class_labels']``, keyed by + ``'_wandb_delimeter_'``. Without this key a + migrated mask renders blank. Returns ``{that_key: {id: name}}``; parsed + once per run and cached. ``run.config`` strips ``_wandb``, so this reads + ``run.json_config`` (which keeps it). + """ + cache_id = getattr(run, 'id', None) + if getattr(self, '_mask_labels_run', None) != cache_id: + self._mask_labels_run = cache_id + out: Dict[str, Dict[str, str]] = {} + try: + raw = getattr(run, 'json_config', None) + cfg = json.loads(raw) if raw else dict(getattr(run, 'config', {})) + entries = ( + cfg.get('_wandb', {}).get('value', {}).get('mask/class_labels', {}) + ) + for k, v in (entries or {}).items(): + val = v.get('value') if isinstance(v, dict) else None + if isinstance(val, dict): + out[k] = val + except Exception as e: + logger.debug( + f'{tag}: mask class_labels parse failed for {cache_id}: {e}' + ) + self._mask_labels_cache = out + return self._mask_labels_cache + + def _image_annotation_value( + self, run: Any, key: str, boxes: Any, masks: Any + ) -> Optional[str]: + """Build the staged ``annotation_value`` JSON for one image's boxes/masks. + + Shared by the single-image and gallery (``images/separated``) paths so + list-logged images keep their annotations too. Folds mask class_labels + back in from the run config (wandb stores them there, not on the mask + descriptor) so masks render coloured. Returns None if the image has + neither boxes nor masks. + """ + if masks: + label_map = self._mask_class_labels(run) + for layer_name, layer in masks.items(): + if isinstance(layer, dict) and 'class_labels' not in layer: + cl = label_map.get(f'{key}_wandb_delimeter_{layer_name}') + if cl: + layer['class_labels'] = cl + if not (boxes or masks): + return None + if boxes: + self._migrated('image-boxes') + if masks: + self._migrated('image-masks') + return json.dumps( + {k: v for k, v in {'boxes': boxes, 'masks': masks}.items() if v} + ) + + 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 value is None: + return + # Booleans (bool is an int subclass) are real metric series in wandb + # (e.g. is_best=True/False each epoch); record them as 1.0/0.0 rather + # than silently dropping the whole series. + if isinstance(value, bool): + value = float(value) + 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), + ) + self._migrated('metric') + return + if isinstance(value, str): + # wandb serializes non-finite metric points as strings + # ('NaN'/'Infinity'/'-Infinity'); coerce them back to real floats so + # they migrate like any other metric (ClickHouse stores NaN/Inf + # natively — a run that logged nan/inf numerically already works). + nonfinite = _NONFINITE_STRINGS.get(value.strip()) + if nonfinite is not None: + writer.write_row( + **base, + attribute_path=key, + attribute_type='metric', + step=step, + timestamp_ms=timestamp_ms, + float_value=nonfinite, + ) + self._migrated('metric') + return + # Any other string (a status label, phase name, etc.) is a + # categorical series: stage it as a string_series row so it can be + # rendered as a state timeline. Over-long values are stray blobs, + # not real states -> drop with a flag. + if len(value) > _STRING_SERIES_MAX_LEN: + self._skipped('string-series-too-long') + return + writer.write_row( + **base, + attribute_path=key, + attribute_type='string_series', + step=step, + timestamp_ms=timestamp_ms, + string_value=value, + ) + self._migrated('string_series') + return + if isinstance(value, dict): + media_type = value.get('_type') + if media_type in _FILE_MEDIA_TYPES and value.get('path'): + if not self.include_files: + # --no-files: the file won't be downloaded, so don't stage a + # row pointing at it (a dangling ref would fail the loader). + self._skipped('media-file(--no-files)') + return + # Bounding boxes / segmentation masks ride on the image value as + # references to sidecar files (….boxes2D.json / ….mask.png). Stage + # those refs; the loader resolves the box JSON inline and + # re-uploads the mask PNG (as fileType "mask") into the image's + # annotations. + # Bounding boxes / segmentation masks ride on the image value as + # references to sidecar files; stage the refs (loader resolves box + # JSON inline and re-uploads the mask PNG into the annotations). + annotation_value = self._image_annotation_value( + run, key, value.get('boxes'), value.get('masks') + ) + 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'), + annotation_value=annotation_value, + ) + self._migrated('media') + elif media_type == 'histogram': + bins = value.get('bins') + if bins is None: + # wandb stores histogram edges compactly as packedBins + # ({min, size, count}), not a `bins` array. Reconstruct the + # real edges so the migrated histogram keeps its true value + # range instead of a generic 0..N axis. (Loader synthesizes + # integer edges only if this is also absent.) + bins = self._bins_from_packed(value.get('packedBins')) + 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': bins, + } + ), + ) + self._migrated('histogram') + elif media_type == 'images/separated' and value.get('filenames'): + if not self.include_files: + self._skipped('media-file(--no-files)', len(value['filenames'])) + return + captions = value.get('captions') or [] + # Per-image boxes/masks ride in parallel lists (all_boxes[i] / + # all_masks[i]), same shape as a single image's boxes/masks. Carry + # each image's annotations so galleries keep them (not just plain + # pictures). + all_boxes = value.get('all_boxes') or [] + all_masks = value.get('all_masks') or [] + for i, filename in enumerate(value['filenames']): + annotation_value = self._image_annotation_value( + run, + key, + all_boxes[i] if i < len(all_boxes) else None, + all_masks[i] if i < len(all_masks) else None, + ) + 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, + annotation_value=annotation_value, + ) + self._migrated('media', len(value['filenames'])) + elif isinstance(value.get(media_type), list) and all( + isinstance(it, dict) + and it.get('path') + and it.get('_type') in _FILE_MEDIA_TYPES + for it in value[media_type] + ): + # Media LISTS logged at one step: _type 'videos'/'audio' hold + # their items under a key matching the _type, each a full + # {path, caption, _type: -file} dict (wandb.log({"k":[v0,v1]})). + # Emit one media row per item (loader batches same name+step -> + # sampleIndex order). + if not self.include_files: + self._skipped('media-file(--no-files)', len(value[media_type])) + return + items = value[media_type] + for it in items: + writer.write_row( + **base, + attribute_path=key, + attribute_type='media', + step=step, + timestamp_ms=timestamp_ms, + string_value=it['_type'], + file_value=f'files/{it["path"]}', + caption=it.get('caption'), + ) + self._migrated('media', len(items)) + else: + # Unknown wandb media type (custom chart, molecule, bokeh, a + # partitioned/joined table, ...). Count by type so the coverage + # report says exactly what was dropped. + self._skipped(f'unsupported({media_type})') + else: + self._skipped(f'unsupported({type(value).__name__})') + + def _export_system_metrics(self, run: Any, writer: PartWriter) -> None: + base = self._row_base(run) + try: + # run.history() defaults to samples=500, which silently downsamples + # long runs (a 2h run at ~2s sampling has ~3600 points). Request a + # high sample count and warn if we still hit the cap. + events = run.history( + stream='events', pandas=False, samples=self.system_samples + ) + except Exception as e: + logger.warning(f'{tag}: system metrics unavailable for {run.id}: {e}') + return + events = list(events) + if len(events) >= self.system_samples: + logger.warning( + f'{tag}: system metrics for {run.id} hit the ' + f'{self.system_samples}-sample cap and may be downsampled; ' + 'raise system_samples to capture full resolution' + ) + 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), + ) + self._migrated('system-metric') + + def _download_files(self, run: Any, files_dir: Path) -> None: + files_dir.mkdir(parents=True, exist_ok=True) + targets = [ + f for f in run.files() if self.include_files or f.name == 'output.log' + ] + if not targets: + return + + def _download_one(f: Any) -> Optional[str]: + # Each wandb file is a separate HTTP request (file API -> storage + # redirect), so downloads are latency-bound. A media-heavy run has + # hundreds of tiny files; downloading them concurrently is a large + # (~10-30x) speedup and is the dominant cost at scale. Distinct + # filenames -> no write contention. Returns the name on failure. + try: + f.download(root=str(files_dir), exist_ok=True) + return None + except Exception as e: + logger.warning(f'{tag}: failed to download {f.name}: {e}') + return f.name + + workers = min(self.download_workers, len(targets)) + if workers <= 1: + results = [_download_one(f) for f in targets] + else: + with ThreadPoolExecutor(max_workers=workers) as pool: + # Results collected on this thread -> the coverage counter below + # is only touched here, never concurrently. + results = list(pool.map(_download_one, targets)) + failed = [name for name in results if name is not None] + if failed: + # Media rows were already staged (and counted migrated) pointing at + # these files; a failed download leaves a dangling ref the loader + # silently skips. Count it so coverage/--strict surface the loss + # instead of a false "fully migrated". + self._skipped('file-download-failed', len(failed)) + + 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) + # Fall back to the run's creation time, then its heartbeat, for lines + # without an embedded timestamp. Only if neither is parseable do we + # stamp 0 (1970) — and we warn, rather than silently dating logs to the + # epoch. + fallback_ms = created_ms + if fallback_ms is None: + fallback_ms = parse_iso_ms(getattr(run, 'heartbeat_at', None)) + if fallback_ms is None: + logger.warning( + f'{tag}: {run.id} has no parseable creation/heartbeat time; ' + 'console lines without embedded timestamps will be stamped ' + '0 (1970-01-01)' + ) + fallback_ms = 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, + ) + self._migrated('console-line') + + @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 _flag_run_level_omissions(self, run: Any) -> None: + """Surface run-level context for the coverage report / --strict. Sweep + membership now migrates (see _sweep_block -> manifest 'sweep'); input- + artifact (used_artifacts) lineage still doesn't. Both are wrapped: these + are extra API reads that must never fail an export.""" + try: + if getattr(run, 'sweep', None) is not None: + self._migrated('sweep') + except Exception as e: + logger.debug(f'{tag}: sweep check failed for {run.id}: {e}') + # Input lineage only matters when artifacts are being migrated at all. + if self.include_artifacts: + try: + if any(True for _ in run.used_artifacts()): + self._skipped('artifact-input-lineage') + except Exception as e: + logger.debug(f'{tag}: used_artifacts check failed for {run.id}: {e}') + + 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 + versioning_lost = False + for artifact in artifacts: + # Only the artifact's *files* migrate — versions, non-'latest' + # aliases, and the type/lineage graph don't. Skip wandb's internal + # per-run history artifact (always present, always v0/latest) so it + # doesn't false-trigger the flag. + if getattr(artifact, 'type', None) != 'wandb-history': + aliases = set(getattr(artifact, 'aliases', None) or []) + if getattr(artifact, 'version', None) not in (None, 'v0') or ( + aliases - {'latest'} + ): + versioning_lost = True + 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})' + ) + self._skipped('artifact-over-size-cap') + 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}') + self._skipped('artifact-download-failed') + continue + # A logged table with media columns migrates degraded: the table + # itself lands (from the run-files copy) but its media cells become + # text placeholders, and the cell media arrives as this artifact, + # unlinked. Flag it so the gap is visible (and trips --strict) + # instead of migrating silently. + media_cols = _table_media_columns(dest) + if media_cols: + logger.warning( + f'{tag}: table {artifact.name!r} has media column(s) ' + f'{media_cols}: cells migrate as text placeholders, and the ' + f'cell media migrates as this artifact but is not linked ' + f'back to the table' + ) + self._skipped('table-media-cell', len(media_cols)) + 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)), + ) + self._migrated('artifact-file') + if versioning_lost: + self._skipped('artifact-versioning') + + +def list_wandb_projects(entity: str, api_key: Optional[str] = None) -> List[str]: + """Return the names of every wandb project under ``entity`` (for migrating a + whole account at once). Imported lazily so the base package stays light.""" + import wandb + + api = wandb.Api(api_key=api_key) + return [p.name for p in api.projects(entity)] diff --git a/pluto/op.py b/pluto/op.py index e487adc..c234aa7 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 @@ -47,6 +48,12 @@ logger = logging.getLogger(f'{__name__.split(".")[0]}') tag = 'Operation' +# A single string-metric value longer than this is a stray blob (a whole log +# line, serialized object, etc.), not a categorical state label, so it is not +# sent. There is no cardinality/distinct-value guard: every string value is +# logged, regardless of how many distinct values the series has. +STRING_SERIES_MAX_LEN = 200 + def _is_distributed_environment() -> bool: """Check if running in a distributed (DDP/FSDP) environment.""" @@ -327,7 +334,16 @@ def _unregister_sigterm_handler() -> None: 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: @@ -374,16 +390,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 @@ -457,7 +476,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, @@ -497,7 +516,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. ' @@ -546,14 +565,41 @@ def __init__(self, config, settings, tags=None, resume=False) -> None: else None ) self._step = 0 + # Latest numeric value logged per key (last write wins). Lets a sweep + # agent read a run's objective metric without a server round-trip. + self._latest_metrics: Dict[str, float] = {} + # String-metric keys already warned about (over-long value) — so the + # per-value length warning fires at most once per key. + self._string_series_warned: set = set() self._queue: queue.Queue[QueueItem] = queue.Queue() self._finished = False self._finish_lock = threading.Lock() + # Set by _teardown() when the terminal status update could not be + # confirmed on the server (after retries). Callers that must know the + # run was actually finalized — notably the migration loader — check + # this instead of trusting finish() to have succeeded silently. + self._status_update_error: Union[Exception, None] = None # (log-key, exc-type) pairs already surfaced at error from a dropped log # item, so a per-step failure is shouted once then drops to debug. 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. + Migration/backfill may instead supply the original run's own + systemMetadata via ``compat['systemMetadata']`` to preserve the + historical repro context (host/git/python/gpu). + """ + override = self.settings.compat.get('systemMetadata') + if override is not None: + return override + 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 @@ -586,6 +632,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( @@ -659,7 +706,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() ) @@ -688,12 +735,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. @@ -701,23 +756,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] = [] @@ -750,11 +822,63 @@ def _log_via_sync( if metrics: self._sync_manager.enqueue_metrics(metrics, timestamp_ms, self._step) + # Cache the latest numeric value per key so a sweep agent (bayes) can + # read the objective back after the run without a server round-trip. + # Lazy-init: logging is best-effort and must never crash, and some + # call sites build a bare Op (Op.__new__) that skips __init__. + if not hasattr(self, '_latest_metrics'): + self._latest_metrics = {} + self._latest_metrics.update(metrics) # Register new metric/file names with server (required for dashboard display) 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. @@ -809,6 +933,10 @@ def _register_meta_sync( if isinstance(value, (File, Data)): new_file_meta[value.__class__.__name__].append(key) + elif isinstance(value, str): + # A string metric is a string-series (mlop_data); register it under + # the DATA log type so the server indexes it like the migration does. + new_file_meta['DATA'].append(key) elif self._is_numeric_value(value): new_metric_names.append(key) @@ -840,6 +968,40 @@ def _process_log_item_sync( timestamp_ms=timestamp_ms, step=self._step, ) + elif isinstance(value, str): + # A bare string value is a categorical "string metric" (e.g. + # phase='warmup'): route it to the string-series data path. + self._enqueue_string_series_sync(key, value, timestamp_ms) + + def _enqueue_string_series_sync( + self, key: str, value: str, timestamp_ms: int + ) -> None: + """Route a string metric value to the string-series data path. + + A string logged across steps (e.g. ``phase='warmup'``) is a categorical + state series, stored in ``mlop_data`` as ``dataType='string-series'`` and + rendered as a state timeline. Every value is sent regardless of the + series' cardinality; only a single over-long value (a stray blob, not a + state label) is skipped, with a one-time warning per key. + """ + if self._sync_manager is None: + return + if len(value) > STRING_SERIES_MAX_LEN: + if key not in self._string_series_warned: + self._string_series_warned.add(key) + logger.warning( + f'{tag}: string metric {key!r} value skipped — ' + f'{len(value)} chars (> {STRING_SERIES_MAX_LEN}); string ' + 'metrics are short state labels, not free-form text.' + ) + return + self._sync_manager.enqueue_data( + log_name=key, + data_type='string-series', + data_dict=value, # raw string; the sync sends string-series un-wrapped + timestamp_ms=timestamp_ms, + step=self._step, + ) def _enqueue_file_sync( self, @@ -869,12 +1031,16 @@ def _enqueue_file_sync( local_path=file_obj._path, file_name=file_obj._name, file_ext=file_obj._ext, - file_type=file_obj._type, + # A file may override its upload fileType (e.g. a mask PNG → + # "mask"); otherwise the server derives it from the extension. + file_type=getattr(file_obj, '_upload_file_type', None) + or file_obj._type, file_size=file_obj._stat.st_size, log_name=log_name, timestamp_ms=timestamp_ms, step=self._step, caption=file_obj._caption, + annotations=getattr(file_obj, '_annotations', None), sample_index=sample_index, ) logger.debug( @@ -885,6 +1051,12 @@ def _enqueue_file_sync( f'{tag}: Cannot enqueue file for sync - path is None after load' ) + # Upload any annotation sub-files (segmentation-mask PNGs) in the same + # log group; they carry _upload_file_type="mask" and are referenced from + # the parent image's annotations by fileName. + for extra in getattr(file_obj, '_annotation_files', None) or []: + self._enqueue_file_sync(log_name, extra, timestamp_ms, sample_index) + def finish(self, code: Union[int, None] = None) -> None: """Finish logging and mark the run as a terminal status on the server. @@ -951,6 +1123,13 @@ def _teardown(self, code: Union[int, None], update_status: bool) -> None: # because all ranks must progress together for collective operations is_distributed = _is_distributed_environment() + # Reset per teardown. status_confirmed tracks whether the terminal + # status update actually landed on the server, so (a) a later teardown + # error doesn't wrongly re-mark a finished run FAILED, and (b) the + # migration loader can tell a dropped finish from a real one. + self._status_update_error = None + status_confirmed = False + try: # Stop the monitor (system metrics and heartbeats) self._monitor.stop(code) @@ -989,9 +1168,24 @@ def _teardown(self, code: Union[int, None], update_status: bool) -> None: self._sync_manager.close() self._sync_manager = None - # Update run status on server (only when finishing) + # Update run status on server (only when finishing). This is the + # run's terminal transition; update_status() retries transient + # resets and raises if it can't confirm. Record — but don't abort + # teardown on — that failure: the store/HTTP-client cleanup below + # still has to run, and the caller inspects _status_update_error. if update_status and self._iface: - self._iface.update_status() + try: + self._iface.update_status() + status_confirmed = True + except Exception as status_exc: + self._status_update_error = status_exc + logger.error( + '%s: terminal status update not confirmed after ' + 'retries: %s: %s', + tag, + type(status_exc).__name__, + status_exc, + ) # Clean up data store if used (legacy mode) if self._store: @@ -1008,26 +1202,33 @@ def _teardown(self, code: Union[int, None], update_status: bool) -> None: logger.debug(f'{tag}: closed (run status unchanged)') except (Exception, KeyboardInterrupt) as e: _sentry.capture_exception(e) - if update_status: + # Only report FAILED if we hadn't already confirmed the terminal + # status — otherwise a teardown hiccup *after* a successful finish + # (e.g. an HTTP-client close error) would flip a COMPLETED run to + # FAILED. Guard the report itself too: it can now raise. + if update_status and not status_confirmed: self.settings._op_status = signal.SIGINT.value if self._iface: - self._iface._update_status( - self.settings, - trace={ - 'type': e.__class__.__name__, - 'message': str(e), - 'frames': [ - { - 'filename': frame.filename, - 'lineno': frame.lineno, - 'name': frame.name, - 'line': frame.line, - } - for frame in traceback.extract_tb(e.__traceback__) - ], - 'trace': traceback.format_exc(), - }, - ) + try: + self._iface._update_status( + self.settings, + trace={ + 'type': e.__class__.__name__, + 'message': str(e), + 'frames': [ + { + 'filename': frame.filename, + 'lineno': frame.lineno, + 'name': frame.name, + 'line': frame.line, + } + for frame in traceback.extract_tb(e.__traceback__) + ], + 'trace': traceback.format_exc(), + }, + ) + except Exception as report_exc: + self._status_update_error = report_exc logger.critical('%s: interrupted %s', tag, e) # Re-raise user-initiated termination so the process actually # exits as the user expects. Post-cleanup (sentry flush, diff --git a/pluto/sets.py b/pluto/sets.py index e8f475a..31ea374 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 # Capture console at the file-descriptor level (dup2 tee, like wandb's # console="redirect") so logging handlers configured before init() are diff --git a/pluto/sweep.py b/pluto/sweep.py new file mode 100644 index 0000000..cd443e6 --- /dev/null +++ b/pluto/sweep.py @@ -0,0 +1,533 @@ +"""Hyperparameter sweeps — mirrors the wandb sweep SDK. + +A *sweep* runs a training function many times, each with a different combination +of hyperparameters drawn from a search space, to find the best-performing set:: + + sweep_id = pluto.sweep( + { + "method": "grid", # grid | random (bayes: planned, via optuna) + "metric": {"name": "val_loss", "goal": "minimize"}, + "parameters": { + "lr": {"values": [0.1, 0.01]}, + "batch_size": {"values": [16, 32]}, + }, + }, + project="demo", + ) + + def train(): + run = pluto.init() # sampled config injected: run.config["lr"], ... + ... # train with those settings + run.log({"val_loss": ...}) + run.finish() + + pluto.agent(sweep_id, train, count=4) # runs the 4 grid combinations + +The config schema mirrors wandb's (``method`` / ``metric`` / ``parameters``) so +existing wandb sweep configs work verbatim and wandb sweeps migrate cleanly. + +Single-machine model: the "brain" that picks each next combination runs +client-side in :func:`agent` (grid enumerates the search space, random samples +it). Every run started under an agent is tagged ``sweep:`` and carries its +sampled hyperparameters in ``config`` — the same data model the wandb migration +produces — so a sweep dashboard can group and compare runs off of that. +""" + +from __future__ import annotations + +import itertools +import json +import logging +import math +import os +import random +import string +from typing import Any, Callable, Dict, Iterator, List, Optional + +logger = logging.getLogger(f'{__name__.split(".")[0]}') +tag = 'sweep' + +_VALID_METHODS = ('grid', 'random', 'bayes') + +# Set by agent() around each function() call so pluto.init() picks up the sampled +# hyperparameters + the sweep tag. Mirrors how wandb.agent feeds wandb.config. +_active_sweep: Optional[Dict[str, Any]] = None + +# The Op that pluto.init() created for the current sweep run — init() hands it +# back here so the agent can read the objective metric (bayes) and finish the +# run, without relying on pluto.ops (which finish() mutates + ids get reused). +_last_run_op: Optional[Any] = None + +# The active sweep's *declared* spec (id + method + metric + search space), +# constant for the whole agent() run. init() stamps it onto each run's config +# (as ``config.sweep``) so the server has the real declaration for a native +# sweep — mirroring ``config.wandb.sweep`` for migrated sweeps — instead of +# having to infer method/objective/search-space from the runs. +_active_declared: Optional[Dict[str, Any]] = None + +# id -> config, backed by an on-disk copy so an agent in a separate process can +# still load the sweep. (When the backend gains a sweep entity, _store_sweep / +# _load_sweep become the only spots that need to talk to it.) +_SWEEP_REGISTRY: Dict[str, Dict[str, Any]] = {} + + +def _gen_sweep_id(n: int = 8) -> str: + alphabet = string.ascii_lowercase + string.digits + return ''.join(random.choice(alphabet) for _ in range(n)) + + +def _sweeps_dir() -> str: + base = os.environ.get('PLUTO_DIR') or os.path.join( + os.path.expanduser('~'), '.pluto' + ) + d = os.path.join(base, 'sweeps') + os.makedirs(d, exist_ok=True) + return d + + +def _store_sweep(sweep_id: str, config: Dict[str, Any]) -> None: + _SWEEP_REGISTRY[sweep_id] = config + try: + with open(os.path.join(_sweeps_dir(), f'{sweep_id}.json'), 'w') as f: + json.dump(config, f) + except Exception as e: # persistence is best-effort; in-process still works + logger.debug(f'{tag}: could not persist sweep {sweep_id}: {e}') + + +def _load_sweep(sweep_id: str) -> Dict[str, Any]: + if sweep_id in _SWEEP_REGISTRY: + return _SWEEP_REGISTRY[sweep_id] + path = os.path.join(_sweeps_dir(), f'{sweep_id}.json') + if os.path.exists(path): + with open(path) as f: + cfg = json.load(f) + _SWEEP_REGISTRY[sweep_id] = cfg + return cfg + raise ValueError( + f'unknown sweep id {sweep_id!r}; call pluto.sweep(...) to create it first' + ) + + +def _validate_config(config: Dict[str, Any]) -> Dict[str, Any]: + if not isinstance(config, dict): + raise TypeError(f'sweep config must be a dict, got {type(config).__name__}') + method = config.get('method', 'grid') + if method not in _VALID_METHODS: + raise ValueError( + f'sweep method must be one of {_VALID_METHODS}, got {method!r}' + ) + if method == 'bayes': + metric = config.get('metric') or {} + if not isinstance(metric, dict) or not metric.get('name'): + raise ValueError( + "a bayes sweep needs metric={'name': ..., 'goal': ...} to optimize" + ) + parameters = config.get('parameters') + if not isinstance(parameters, dict) or not parameters: + raise ValueError("sweep config needs a non-empty 'parameters' dict") + for name, spec in parameters.items(): + if not isinstance(spec, dict): + raise ValueError( + f'parameter {name!r} must map to a dict ' + "(e.g. {'values': [...]} or {'min': .., 'max': ..}), " + f'got {type(spec).__name__}' + ) + return config + + +def _grid_values(name: str, spec: Dict[str, Any]) -> List[Any]: + if 'value' in spec: + return [spec['value']] + if 'values' in spec: + return list(spec['values']) + raise ValueError( + f"grid search needs discrete 'values'/'value' for parameter {name!r}, " + "but it has a range/distribution — use method='random' instead" + ) + + +def _grid_combos(parameters: Dict[str, Any]) -> Iterator[Dict[str, Any]]: + names = list(parameters) + value_lists = [_grid_values(n, parameters[n]) for n in names] + for combo in itertools.product(*value_lists): + yield dict(zip(names, combo)) + + +def _sample_param(name: str, spec: Dict[str, Any]) -> Any: + if 'value' in spec: + return spec['value'] + if 'values' in spec: + return random.choice(list(spec['values'])) + dist = spec.get('distribution') + lo, hi = spec.get('min'), spec.get('max') + if lo is not None and hi is not None: + if dist == 'int_uniform' or ( + dist is None and isinstance(lo, int) and isinstance(hi, int) + ): + return random.randint(int(lo), int(hi)) + if dist in ('log_uniform_values', 'log_uniform'): + return math.exp(random.uniform(math.log(lo), math.log(hi))) + return random.uniform(lo, hi) + if dist == 'normal': + return random.gauss(spec.get('mu', 0.0), spec.get('sigma', 1.0)) + raise ValueError( + f"cannot sample parameter {name!r}: give 'values', 'value', or " + "'min'/'max' (optionally with 'distribution')" + ) + + +def _random_combo(parameters: Dict[str, Any]) -> Dict[str, Any]: + return {name: _sample_param(name, spec) for name, spec in parameters.items()} + + +def _combo_key(config: Dict[str, Any], swept_names: List[str]) -> tuple: + """A hashable identity for a combination, over just the swept parameters.""" + return tuple((name, config.get(name)) for name in swept_names) + + +def _declared_meta(sweep_id: str, cfg: Dict[str, Any]) -> Dict[str, Any]: + """The declared sweep spec to stamp on each run (id + method/metric/params).""" + meta: Dict[str, Any] = {'id': sweep_id} + for key in ('method', 'metric', 'parameters'): + if key in cfg: + meta[key] = cfg[key] + return meta + + +def _completed_sweep_runs( + project: Optional[str], sweep_id: str +) -> List[Dict[str, Any]]: + """List already-COMPLETED runs of this sweep (for resume). Best-effort: any + query failure (no backend, project not created yet) yields [] — the agent + then simply runs everything, so resume never breaks a fresh sweep.""" + if not project: + return [] + import pluto.query as pq + + done: List[Dict[str, Any]] = [] + try: + offset = 0 + while True: + page = pq.list_runs( + project, tags=[f'sweep:{sweep_id}'], limit=200, offset=offset + ) + if not page: + break + done.extend(r for r in page if r.get('status') == 'COMPLETED') + if len(page) < 200: + break + offset += 200 + except Exception as e: + logger.debug(f'{tag}: could not list sweep {sweep_id} runs (no resume): {e}') + return [] + return done + + +def _fetch_run_config(project: Optional[str], run: Dict[str, Any]) -> Dict[str, Any]: + """Fetch a completed run's config (list_runs omits it) to match grid combos.""" + if not project: + return {} + import pluto.query as pq + + rid = run.get('displayId') or run.get('id') + if rid is None: + return {} + try: + full = pq.get_run(project, rid) + cfg = full.get('config') + return cfg if isinstance(cfg, dict) else {} + except Exception: + return {} + + +def _fetch_run_metric( + project: Optional[str], run: Dict[str, Any], metric_name: str +) -> Optional[float]: + """Fetch a completed run's final value for ``metric_name`` (bayes seeding).""" + if not project: + return None + import pluto.query as pq + + rid = run.get('displayId') or run.get('id') + if rid is None: + return None + try: + data = pq.get_metrics(project, rid, metric_names=[metric_name]) + vals = [ + p['value'] + for p in (data or []) + if isinstance(p, dict) + and p.get('metric') == metric_name + and isinstance(p.get('value'), (int, float)) + ] + return vals[-1] if vals else None # last logged value == the objective + except Exception: + return None + + +def _optuna_distribution(spec: Dict[str, Any]) -> Any: + """Translate one parameter spec into an optuna distribution (for seeding).""" + import optuna.distributions as od + + if 'values' in spec: + return od.CategoricalDistribution(list(spec['values'])) + lo, hi = spec['min'], spec['max'] + dist = spec.get('distribution') + if dist == 'int_uniform' or ( + dist is None and isinstance(lo, int) and isinstance(hi, int) + ): + return od.IntDistribution(int(lo), int(hi)) + log = dist in ('log_uniform_values', 'log_uniform') + return od.FloatDistribution(float(lo), float(hi), log=log) + + +def _optuna_suggest(trial: Any, name: str, spec: Dict[str, Any]) -> Any: + """Ask optuna for the next value of one parameter, honoring its spec.""" + if 'value' in spec: # a constant — not optimized, just passed through + return spec['value'] + if 'values' in spec: + return trial.suggest_categorical(name, list(spec['values'])) + lo, hi = spec.get('min'), spec.get('max') + if lo is None or hi is None: + raise ValueError( + f"cannot optimize parameter {name!r}: give 'values', 'value', or " + "'min'/'max'" + ) + dist = spec.get('distribution') + if dist == 'int_uniform' or ( + dist is None and isinstance(lo, int) and isinstance(hi, int) + ): + return trial.suggest_int(name, int(lo), int(hi)) + log = dist in ('log_uniform_values', 'log_uniform') + return trial.suggest_float(name, float(lo), float(hi), log=log) + + +def _seed_study( + study: Any, + project: Optional[str], + completed: List[Dict[str, Any]], + parameters: Dict[str, Any], + metric_name: str, +) -> int: + """Replay completed runs' (params -> objective) into the study so a resumed + bayes search learns from prior results. Best-effort; skips any run whose + config/objective can't be fetched.""" + import optuna + + # Only params optuna actually optimizes (constants have no distribution). + opt_params = {n: s for n, s in parameters.items() if 'value' not in s} + if not project or not opt_params: + return 0 + seeded = 0 + for run in completed: + cfg = _fetch_run_config(project, run) + value = _fetch_run_metric(project, run, metric_name) + if not cfg or value is None: + continue + try: + params = {n: cfg[n] for n in opt_params if n in cfg} + if len(params) != len(opt_params): + continue # missing a swept param; can't reconstruct the trial + distributions = {n: _optuna_distribution(opt_params[n]) for n in params} + study.add_trial( + optuna.trial.create_trial( + params=params, distributions=distributions, value=value + ) + ) + seeded += 1 + except Exception: + continue + return seeded + + +def _run_combo( + sweep_id: str, + combo: Dict[str, Any], + project: Optional[str], + function: Callable[[], Any], + metric_name: Optional[str], + index: int, + total: int, +) -> Optional[float]: + """Run ``function`` once with ``combo`` injected; return the run's objective + (its final ``metric_name`` value) if asked. Sets/clears the sweep context and + finishes any run the function left open.""" + global _active_sweep, _last_run_op + _active_sweep = {'id': sweep_id, 'config': combo, 'project': project} + _last_run_op = None + objective: Optional[float] = None + try: + function() + except Exception as e: + logger.error( + f'{tag}: sweep {sweep_id} run {index + 1}/{total} raised ' + f'{type(e).__name__}: {e}' + ) + finally: + _active_sweep = None + op = _last_run_op + _last_run_op = None + if op is not None: + if metric_name is not None: + objective = getattr(op, '_latest_metrics', {}).get(metric_name) + # Close the run if the function didn't, so the next combo starts clean. + if not getattr(op, '_finished', True): + try: + op.finish() + except Exception as e: + logger.debug(f'{tag}: error finishing sweep run: {e}') + return objective + + +def sweep( + config: Dict[str, Any], + project: Optional[str] = None, + entity: Optional[str] = None, +) -> str: + """Create a sweep from a search-space config and return its id. + + ``config`` mirrors wandb's schema — ``method`` (``grid``/``random``/ + ``bayes``), ``metric`` (``{"name", "goal"}``; required for ``bayes``), and + ``parameters`` (each a ``{"values": [...]}`` / ``{"value": x}`` / + ``{"min", "max"}`` spec). Pass the returned id to :func:`agent`. ``entity`` + is accepted for wandb parity and currently unused. + """ + cfg = dict(_validate_config(config)) + sweep_id = _gen_sweep_id() + cfg['_project'] = project # remember the target project for the agent + _store_sweep(sweep_id, cfg) + logger.info( + f'{tag}: created sweep {sweep_id} (method={cfg.get("method", "grid")}, ' + f'{len(cfg["parameters"])} parameters)' + ) + return sweep_id + + +def agent( + sweep_id: str, + function: Callable[[], Any], + count: Optional[int] = None, + project: Optional[str] = None, +) -> None: + """Run ``function`` once per hyperparameter combination in the sweep. + + ``function`` takes no arguments and should call :func:`pluto.init` inside — + the sampled hyperparameters are injected into that run's ``config`` and the + run is tagged ``sweep:`` (mirrors ``wandb.agent``). ``grid`` runs every + combination (``count`` caps it); ``random`` and ``bayes`` run ``count`` + combinations (``count`` required). ``bayes`` uses optuna to pick each next + combination from the results so far. Already-COMPLETED runs are skipped + (resume); any run the function leaves open is finished automatically. + """ + cfg = _load_sweep(sweep_id) + method = cfg.get('method', 'grid') + parameters = cfg['parameters'] + proj = project or cfg.get('_project') + metric = cfg.get('metric') or {} + metric_name = metric.get('name') if isinstance(metric, dict) else None + + # Resume: skip combinations whose run already COMPLETED. Best-effort — needs + # a project to query; a fresh sweep just sees zero done and runs everything. + completed = _completed_sweep_runs(proj, sweep_id) if proj else [] + n_done = len(completed) + + global _active_declared + _active_declared = _declared_meta(sweep_id, cfg) + try: + if method == 'bayes': + _run_bayes(sweep_id, cfg, proj, function, count, completed, metric_name) + return + + if method == 'grid': + swept = list(parameters) + done_keys = { + _combo_key(rc, swept) + for r in completed + if (rc := _fetch_run_config(proj, r)) + } + combos = [ + c + for c in _grid_combos(parameters) + if _combo_key(c, swept) not in done_keys + ] + if count is not None: + combos = combos[:count] # cap new runs this invocation + else: # random + if count is None: + raise ValueError( + "method='random' needs count= in pluto.agent(sweep_id, fn, " + 'count=n) — a random search has no natural end' + ) + remaining = max(0, count - n_done) # count is the sweep's total target + combos = [_random_combo(parameters) for _ in range(remaining)] + + if n_done: + logger.info( + f'{tag}: resuming sweep {sweep_id}: {n_done} run(s) already done, ' + f'{len(combos)} to go' + ) + logger.info(f'{tag}: agent starting {len(combos)} runs for sweep {sweep_id}') + for i, combo in enumerate(combos): + _run_combo(sweep_id, combo, proj, function, None, i, len(combos)) + logger.info(f'{tag}: agent finished {len(combos)} runs for sweep {sweep_id}') + finally: + _active_declared = None + + +def _run_bayes( + sweep_id: str, + cfg: Dict[str, Any], + project: Optional[str], + function: Callable[[], Any], + count: Optional[int], + completed: List[Dict[str, Any]], + metric_name: Optional[str], +) -> None: + """Bayesian search via optuna: ask for a combination, run it, tell optuna the + objective, repeat — learning as it goes. Resumes by (a) not exceeding the + total ``count`` and (b) seeding the study with completed runs' results.""" + try: + import optuna + except ImportError: + raise ImportError( + "sweep method='bayes' needs optuna — install it with " + '`pip install optuna` (or `pip install pluto[sweep]`).' + ) + if count is None: + raise ValueError("method='bayes' needs count= in pluto.agent(...)") + if not metric_name: + raise ValueError( + "a bayes sweep needs metric={'name': ..., 'goal': ...} to optimize" + ) + parameters = cfg['parameters'] + goal = str((cfg.get('metric') or {}).get('goal') or 'minimize').lower() + direction = 'maximize' if goal.startswith('max') else 'minimize' + + optuna.logging.set_verbosity(optuna.logging.WARNING) + study = optuna.create_study(direction=direction) + seeded = _seed_study(study, project, completed, parameters, metric_name) + remaining = max(0, count - len(completed)) + if seeded or completed: + logger.info( + f'{tag}: resuming bayes sweep {sweep_id}: {len(completed)} done ' + f'({seeded} seeded into the optimizer), {remaining} to go' + ) + logger.info(f'{tag}: agent starting {remaining} bayes runs for sweep {sweep_id}') + for i in range(remaining): + trial = study.ask() + combo = { + name: _optuna_suggest(trial, name, spec) + for name, spec in parameters.items() + } + objective = _run_combo( + sweep_id, combo, project, function, metric_name, i, remaining + ) + if objective is None: # the run didn't log the metric — skip this trial + study.tell(trial, state=optuna.trial.TrialState.FAIL) + logger.warning( + f'{tag}: bayes run {i + 1}/{remaining} logged no {metric_name!r}; ' + 'optuna cannot learn from it' + ) + else: + study.tell(trial, objective) + logger.info(f'{tag}: agent finished {remaining} bayes runs for sweep {sweep_id}') diff --git a/pluto/sync/process.py b/pluto/sync/process.py index be7da3d..e264fca 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 @@ -199,20 +199,53 @@ def stop( timeout = timeout or self.settings.get('sync_process_shutdown_timeout', 30.0) start = time.time() + completed = False while time.time() - start < timeout: pending = self.store.get_pending_count(self.run_id) if pending == 0: self.store.mark_run_synced(self.run_id) logger.info('Sync process completed successfully') - return True + completed = True + break time.sleep(0.1) - pending = self.store.get_pending_count(self.run_id) - logger.warning( - f'Sync process did not complete within {timeout}s, ' - f'{pending} records pending. Data preserved in {self.db_path}' - ) - return False + if not completed: + pending = self.store.get_pending_count(self.run_id) + logger.warning( + f'Sync process did not complete within {timeout}s, ' + f'{pending} records pending. Data preserved in {self.db_path}' + ) + + # The data is flushed, but the sync process is a persistent daemon: its + # main loop only breaks on SIGTERM or when the PARENT dies. If we just + # return here it keeps running until *this* process exits. A long-lived + # caller that spawns one sync process per run (e.g. a bulk pluto.migrate + # load of hundreds of runs in a single invocation) would then accumulate + # one live subprocess per run and exhaust memory. Terminate it now so its + # lifetime is scoped to the run, not to the whole caller. + self._terminate_process(timeout) + return completed + + def _terminate_process(self, timeout: float) -> None: + """SIGTERM the sync subprocess and reap it, escalating to SIGKILL. + + Safe to call after the data flush above: the process drains (a no-op + when pending is already 0), closes its clients, and exits. Bounded so a + subprocess that hangs on exit can't stall the caller. + """ + proc = self._process + if proc is None or proc.poll() is not None: + return + try: + proc.terminate() + try: + proc.wait(timeout=max(5.0, min(timeout, 30.0))) + except subprocess.TimeoutExpired: + logger.warning('Sync process did not exit on SIGTERM; killing it') + proc.kill() + proc.wait(timeout=5) + except Exception as e: + logger.debug(f'Failed to terminate sync process: {e}') def enqueue_metrics( self, @@ -231,6 +264,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( @@ -253,7 +306,7 @@ def enqueue_data( self, log_name: str, data_type: str, - data_dict: Dict[str, Any], + data_dict: Union[Dict[str, Any], str], timestamp_ms: int, step: Optional[int] = None, ) -> None: @@ -352,6 +405,7 @@ def enqueue_file( timestamp_ms: int, step: Optional[int] = None, caption: Optional[str] = None, + annotations: Optional[str] = None, sample_index: int = 0, ) -> None: """ @@ -371,6 +425,7 @@ def enqueue_file( timestamp_ms=timestamp_ms, step=step, caption=caption, + annotations=annotations, sample_index=sample_index, ) # Update heartbeat to show we're alive @@ -1156,6 +1211,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) @@ -1180,12 +1239,17 @@ def upload_data_batch(self, records: List[SyncRecord]) -> None: lines = [] for record in records: payload = record.payload + raw = payload.get('data', {}) + data_type = payload.get('data_type', 'UNKNOWN') + # string-series carries a raw string value (e.g. 'warmup'); every + # other data type carries a dict that must be JSON-encoded. + data_field = raw if data_type == 'string-series' else json.dumps(raw) lines.append( json.dumps( { 'time': record.timestamp_ms, - 'data': json.dumps(payload.get('data', {})), - 'dataType': payload.get('data_type', 'UNKNOWN'), + 'data': data_field, + 'dataType': data_type, 'logName': payload.get('log_name', ''), 'step': record.step or 0, } @@ -1295,7 +1359,12 @@ def _get_presigned_urls( batch = [] for f in file_records: file_ext = f.file_ext - file_type = file_ext[1:] if file_ext.startswith('.') else file_ext + # A special upload fileType (e.g. "mask") is stored on file_type and + # sent verbatim; otherwise derive it from the extension as before. + if f.file_type == 'mask': + file_type = 'mask' + else: + file_type = file_ext[1:] if file_ext.startswith('.') else file_ext entry: Dict[str, Any] = { 'fileName': f'{f.file_name}{f.file_ext}', 'fileSize': f.file_size, @@ -1312,6 +1381,10 @@ def _get_presigned_urls( # unknown fields anyway). if f.caption is not None: entry['caption'] = f.caption + # Opaque wandb-shape annotations JSON (boxes/masks); only when set so + # the payload is unchanged for un-annotated files / older servers. + if getattr(f, 'annotations', None) is not None: + entry['annotations'] = f.annotations batch.append(entry) body = json.dumps({'files': batch}) diff --git a/pluto/sync/store.py b/pluto/sync/store.py index d00e175..6257691 100644 --- a/pluto/sync/store.py +++ b/pluto/sync/store.py @@ -163,6 +163,10 @@ class FileRecord: error_message: Optional[str] presigned_url: Optional[str] caption: Optional[str] = None + # Opaque JSON string of image annotations (wandb-shape boxes/masks), + # forwarded verbatim to the server (mlop_files.annotations) and parsed by the + # frontend. None for un-annotated files. + annotations: Optional[str] = None # 0-based position of this file within a single (log_name, step) log() call # (e.g. pluto.log({"k": [a, b, c]}) → 0, 1, 2). Lets the server restore the # logged order instead of sorting by fileName. Scalars get 0. @@ -198,8 +202,9 @@ class SyncStore: - Health diagnostics (queue depth, write latency, WAL size) """ - # v2: added file_uploads.caption (backfilled via _add_column_if_missing) - SCHEMA_VERSION = 2 + # v2: file_uploads.caption; v3: file_uploads.annotations (both backfilled + # via _add_column_if_missing, so the version bump is informational). + SCHEMA_VERSION = 3 # SQLite's own busy handler wait. Kept deliberately short because we layer # application-level exponential backoff (see _retry_on_locked) on top: a @@ -327,6 +332,7 @@ def _init_schema(self) -> None: file_ext TEXT, log_name TEXT, caption TEXT, + annotations TEXT, sample_index INTEGER DEFAULT 0, timestamp_ms INTEGER NOT NULL, step INTEGER, @@ -343,6 +349,7 @@ def _init_schema(self) -> None: # CREATE TABLE IF NOT EXISTS above is a no-op on an existing table, # so backfill new columns here. Idempotent (skips if present). self._add_column_if_missing(cursor, 'file_uploads', 'caption', 'TEXT') + self._add_column_if_missing(cursor, 'file_uploads', 'annotations', 'TEXT') self._add_column_if_missing( cursor, 'file_uploads', 'sample_index', 'INTEGER DEFAULT 0' ) @@ -702,6 +709,7 @@ def enqueue_file( timestamp_ms: int, step: Optional[int] = None, caption: Optional[str] = None, + annotations: Optional[str] = None, sample_index: int = 0, ) -> int: """Add a file to the upload queue. Returns file record ID.""" @@ -713,9 +721,9 @@ def enqueue_file( INSERT INTO file_uploads ( run_id, local_path, file_type, file_size, timestamp_ms, step, created_at, file_name, file_ext, - log_name, caption, sample_index + log_name, caption, annotations, sample_index ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, @@ -729,6 +737,7 @@ def enqueue_file( file_ext, log_name, caption, + annotations, sample_index, ), ) @@ -773,6 +782,9 @@ def get_pending_files( error_message=row['error_message'], presigned_url=row['remote_url'], caption=(row['caption'] if 'caption' in row.keys() else None), + annotations=( + row['annotations'] if 'annotations' in row.keys() else None + ), sample_index=( row['sample_index'] if 'sample_index' in row.keys() diff --git a/poetry.lock b/poetry.lock index b32647f..02ab6e7 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,38 @@ -# 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.1.1 and should not be changed by hand. + +[[package]] +name = "alembic" +version = "1.18.5" +description = "A database migration tool for SQLAlchemy." +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +files = [ + {file = "alembic-1.18.5-py3-none-any.whl", hash = "sha256:06d8ba9d04558022f5395e9317de03d270f3dced49cee01f89fe7a13c26f14bc"}, + {file = "alembic-1.18.5.tar.gz", hash = "sha256:1554982221dd17e9a749b53902407578eb305e453f71999e8c7f0a48389fff8e"}, +] + +[package.dependencies] +Mako = "*" +SQLAlchemy = ">=1.4.23" +tomli = {version = "*", markers = "python_version < \"3.11\""} +typing-extensions = ">=4.12" + +[package.extras] +tz = ["tzdata"] + +[[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 +197,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 +313,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 +337,30 @@ 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 = "platform_system == \"Windows\" or sys_platform == \"win32\""} + +[[package]] +name = "colorlog" +version = "6.12.0" +description = "Add colours to the output of Python's logging module." +optional = false +python-versions = ">=3.6" +groups = ["main", "dev"] +files = [ + {file = "colorlog-6.12.0-py3-none-any.whl", hash = "sha256:30d392604e9110045a2c2aeefc27d7a017abbab63f3a8aee594eac0801df784e"}, + {file = "colorlog-6.12.0.tar.gz", hash = "sha256:2a7924c1dadf18b22a0eb8b06d1c7b01d5341707ec1641eb6fcc4fde0c3e8e5f"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "sys_platform == \"win32\""} + +[package.extras] +development = ["black", "flake8", "mypy", "pytest", "types-colorama"] [[package]] name = "contourpy" @@ -483,7 +553,7 @@ description = "Backport of PEP 654 (exception groups)" optional = false python-versions = ">=3.7" groups = ["main", "dev"] -markers = "python_version == \"3.10\"" +markers = "python_version < \"3.11\"" files = [ {file = "exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598"}, {file = "exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219"}, @@ -669,6 +739,100 @@ gitdb = ">=4.0.1,<5" doc = ["sphinx (>=7.1.2,<7.2)", "sphinx-autodoc-typehints", "sphinx_rtd_theme"] test = ["coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock ; python_version < \"3.8\"", "mypy", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar", "typing-extensions ; python_version < \"3.11\""] +[[package]] +name = "greenlet" +version = "3.5.4" +description = "Lightweight in-process concurrent programming" +optional = false +python-versions = ">=3.10" +groups = ["main", "dev"] +markers = "platform_machine == \"aarch64\" or platform_machine == \"ppc64le\" or platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"win32\" or platform_machine == \"WIN32\"" +files = [ + {file = "greenlet-3.5.4-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:ac5bf81d79d2c8eeb2ef6359b2e1687a1e9ebf46c2b1f970da9a9255df51d190"}, + {file = "greenlet-3.5.4-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:89f3738167bab8c1084b94e23023d41d247117ac149fa0fbcb5bd4cf6262b353"}, + {file = "greenlet-3.5.4-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e9a5e3406e3ed8125ae1a3b37c12f3434e2b1f0fa053197c5557895b4fb09606"}, + {file = "greenlet-3.5.4-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a2d614cb2372c7101a12ea8b96dd56f81c986d247c5a73db67063f3ed1ca4a52"}, + {file = "greenlet-3.5.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4ab9f0704bccf6d3b38e0d2130b7b33271cff11453690da074fa280c3aa8e8e7"}, + {file = "greenlet-3.5.4-cp310-cp310-manylinux_2_39_riscv64.whl", hash = "sha256:188e4d142f243051d92a1f5c244a741da02dddc070a0620c842804d7b56d008c"}, + {file = "greenlet-3.5.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:2cdaadc3d31445a8f782bde3cd37e49a2c2a9c6da6daf76a3e34c683b271a3c7"}, + {file = "greenlet-3.5.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:70bdfacdc183dac838b2a0aaff2dd6134a457c52fe68a9c6bbab435483d2b9df"}, + {file = "greenlet-3.5.4-cp310-cp310-win_amd64.whl", hash = "sha256:69173331fbc5d64bfac0065d7e22c39cfcd089e9b18d125bdcd5079363b09616"}, + {file = "greenlet-3.5.4-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:e883de250e299654b1f1680f72a1a9f9ba62c9bd1bce84099c90657349a8dfbb"}, + {file = "greenlet-3.5.4-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32802705c2c1ff25e8237b3bdacf2594fa02be80af8a66703eb7853ea7e68686"}, + {file = "greenlet-3.5.4-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:57aa201b351f7c7c75627c60d29e4d5b97a07d37efeb62b903466fca42c097d7"}, + {file = "greenlet-3.5.4-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9667862a2e38ad379f11b845daeda22c8989186def44f06962c9c4c05e556da7"}, + {file = "greenlet-3.5.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c3fe76c2cac86b4f7a1e92865ac0a54384deb05c92986287c1a7110d9bd53071"}, + {file = "greenlet-3.5.4-cp311-cp311-manylinux_2_39_riscv64.whl", hash = "sha256:ae53534b5dec0f4c2ec26f898f538dc8ea1ca3ef2927d597a9439e40a09da937"}, + {file = "greenlet-3.5.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:1e1a4a684b16c45ba324e60b32a4386a87722bcb815d2a149d2182f9b401ca72"}, + {file = "greenlet-3.5.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:e849e6e139b9671adeac505f72fc05f4af7fd1921faef40295e214fc3b361b59"}, + {file = "greenlet-3.5.4-cp311-cp311-win_amd64.whl", hash = "sha256:dc418cf4c873357964d6624445ed09472e50def990c65dd4e76fc3ba8cd9cef6"}, + {file = "greenlet-3.5.4-cp311-cp311-win_arm64.whl", hash = "sha256:c38c902a0986eba1f6e7ba1ab39ad5195926abde90f3fe080e08212db62176da"}, + {file = "greenlet-3.5.4-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3d66250e8b09f182ede05490998c818b5961f7a3640332d44c4927caec7bbfe4"}, + {file = "greenlet-3.5.4-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c90e930c9c192e5b3ee9fb8bcd920ea3926155e2e3ded39fc697323addecee17"}, + {file = "greenlet-3.5.4-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:791fdfeeb9c6e0c7b10fa151bf110d2a6974866f13dcb5b1c7efae698245893a"}, + {file = "greenlet-3.5.4-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b7c895310363f310361e0fe2072af85269d2a2a285cd04c0c59e79a5e3670dcf"}, + {file = "greenlet-3.5.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f88193799d43dbf8c8a806d6405c9c52fe2af40bf75072a606357b33cc336c7f"}, + {file = "greenlet-3.5.4-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:13b980043cb1b3134e81ea469da1250ddcc6bfe6d245bbaa59168d9cdc8f228f"}, + {file = "greenlet-3.5.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b7a5f095767c4493afcd06067f2bb3b8716e3f3f9e92b99c88e7e99f885b3d4d"}, + {file = "greenlet-3.5.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:42afdc1ab5f66da8c586c32af9224a74a706b4f0ea0dc3a4188a0860a09c65c9"}, + {file = "greenlet-3.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:60149df8f462d1b230038e6590c23c3b4768bb5d6c022b3b6e82532b34b0b8a3"}, + {file = "greenlet-3.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:77d6ce04fed0d9aeed42e0f37923cc43eba9b027bdd9c34546bb4ccd143d0fe0"}, + {file = "greenlet-3.5.4-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:24e61b88cb7e1b1d794b32a10cc346ac779681d6d74ff137a3e0a444d2bf1f02"}, + {file = "greenlet-3.5.4-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:870d730fec833f5a06906a32596cc099b9161594642a92a520b7a88911c95356"}, + {file = "greenlet-3.5.4-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ec5ff0d1878df6af3bf9b638a5a92a7d5693291de77c91bff10fa48519c604ef"}, + {file = "greenlet-3.5.4-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:07bd44616608d873d06735b63ef1a88191d6ca57c8d291d6559c71bc14c0893c"}, + {file = "greenlet-3.5.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d84d993f6e575c950d91a23c1345d18fe1a4310d447bf630849d7809196b52f0"}, + {file = "greenlet-3.5.4-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:3529a8a933582ad19e224792cac7372489526576b75b4c124e8e4f29948f4861"}, + {file = "greenlet-3.5.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:58023945f421093de5e6fa108c0985a8659d43f49e0216da25099369a121bcbd"}, + {file = "greenlet-3.5.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:bae2728e1897aa8df8cb1af38cd48b3a743aefe29372de7b8b7a9f532501e69f"}, + {file = "greenlet-3.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:ca5726c0b08ca35ae873557266a78b2c3f3b2b7d7401aa5ff886c2045dd0111c"}, + {file = "greenlet-3.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:7c1303791d603080cac6fc3b34df51c3b75b723739c282c8029e48a0d241672f"}, + {file = "greenlet-3.5.4-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:7e8afa5eac028f8140ceafe5ceec66e6aa127ddcb21452d2a564dcd2900b5f22"}, + {file = "greenlet-3.5.4-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:73b37afe369021423ea53dd3123e04bffa7e93ac64429b9f50835b2e4fcae7cf"}, + {file = "greenlet-3.5.4-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3ef964f56dfcb6f9bbef2a190d9126795eac408716aeae47b5e7c73c32aafca9"}, + {file = "greenlet-3.5.4-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:cef589bc65fae02d10bca2ac341191c5b33acc2967892ebf4fcbd10eabb7a74c"}, + {file = "greenlet-3.5.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9c53ff01a5c53a40f2c16820ebc56d7c61a77f5fbe009dadd96292d5682f80f8"}, + {file = "greenlet-3.5.4-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:dfc41ae893d9ceaf22c824f2153a88b30651b20e8758c2cd9ac143f23640563c"}, + {file = "greenlet-3.5.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ecca4d80d55a01ad6b23b33262662956149fbb7b2c6be2910f1705921958cbf3"}, + {file = "greenlet-3.5.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2ffbc533e0eaf8e80d8471411646ab88fe58f641d508c0b02b24494479f4d9ec"}, + {file = "greenlet-3.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:305f69e6c4523d7f6979ed001cff4e5853c063e5da04880296603aa0227e544c"}, + {file = "greenlet-3.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:f260930bbbbcf9caee661211235a5111c86dfe5832fdf6ae4570da1e0995320f"}, + {file = "greenlet-3.5.4-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:41ddab54e4b238f4a6c323f39b4e59e176affd5a94d461a9fb7583dac74240a3"}, + {file = "greenlet-3.5.4-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b3dabe3e2809013052c68bdf0b7fa5f5f2859c43a80803131ad61af9cabd7867"}, + {file = "greenlet-3.5.4-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:27d3f00718634d4520a3a150154ac5da36f257869d41321953375b90bfbbc72c"}, + {file = "greenlet-3.5.4-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:39169a11d87a6a263afda3e9a27d1df16d0f919d40a4837cc73986c9884c0dd8"}, + {file = "greenlet-3.5.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cbd60b5763c6543c1827e48faaf14ea9bfbad245f52b1a4d76a2a2d8884c6c66"}, + {file = "greenlet-3.5.4-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:bd3d1145f603b2db19feb9078c2e6855eb7c67e15580c010ed815cee519b86fd"}, + {file = "greenlet-3.5.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f00f910f0e7b35416c63b23ad78b769aeccfc1775f712b43c4ee525624a2eef7"}, + {file = "greenlet-3.5.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:91c26423753b92caf41ab3f98fd547d7374d4d9fc2d85be041886c1579d9255e"}, + {file = "greenlet-3.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:ee032b91fd8ec29ec6c4cea2b8c561b178435134bd0752c7334b94e9c736c132"}, + {file = "greenlet-3.5.4-cp315-cp315-macosx_11_0_universal2.whl", hash = "sha256:178111881dd7a6c946471fda85485ec796e1043c2b939f694b096e2ecf986809"}, + {file = "greenlet-3.5.4-cp315-cp315-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d92df08dd65fede97fc37aad36c2e9dcda3b31c467f8e0c2c096456cb818e927"}, + {file = "greenlet-3.5.4-cp315-cp315-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:99e8f8c4ebc4fd80aa26c1280ae9ad43a0976e786349703a181cf0bae60413e5"}, + {file = "greenlet-3.5.4-cp315-cp315-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1f17e362d78e37559e0506c5a7d066bdd45073c36a0127a543e8a0df27242ff3"}, + {file = "greenlet-3.5.4-cp315-cp315-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:394de08dad5ffcb1f50c2159d93e398d9d2da3ed437645eaa54771fa720db9f0"}, + {file = "greenlet-3.5.4-cp315-cp315-manylinux_2_39_riscv64.whl", hash = "sha256:cd320d998cbaa032932830448e39abf3c6a12901295e386e8114db926e10cffb"}, + {file = "greenlet-3.5.4-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:c883d61f2282d72c767a14936641b3efcbde9d82f1080712aaea0b1d3126cb88"}, + {file = "greenlet-3.5.4-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:2a924f15d17957e252a810acefcb5942f5ca712298e8b6fcaed9a307d357522c"}, + {file = "greenlet-3.5.4-cp315-cp315-win_amd64.whl", hash = "sha256:ed17e5f3420360d5b459de8462efb52060399a5326a613d4cde31cef63ef95da"}, + {file = "greenlet-3.5.4-cp315-cp315-win_arm64.whl", hash = "sha256:f908898d6fa484ce4b6f447ce70ea99b52c503fee419e53cf74d60a16bc9e667"}, + {file = "greenlet-3.5.4-cp315-cp315t-macosx_11_0_universal2.whl", hash = "sha256:1833637f17d5e7472548a48575c394fe39f1b1890d676d162d86593610f44d8c"}, + {file = "greenlet-3.5.4-cp315-cp315t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:12cda9122e03341f1cb6b8207a19d7a9d375e52f1b4e9243918375f40fd7b4b9"}, + {file = "greenlet-3.5.4-cp315-cp315t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d83ae0e32d14957ab7170785a20f582635c8474deab1bfbb552b17e769a6ce25"}, + {file = "greenlet-3.5.4-cp315-cp315t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:123aa379c962ed5fe90a880327e0c3066124ac64ec99e12a238be9fd8eb3db3d"}, + {file = "greenlet-3.5.4-cp315-cp315t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9f1467de1bb767f75db0aa34c195e3a496d8d1278c796e70c24ce205d3e99cde"}, + {file = "greenlet-3.5.4-cp315-cp315t-manylinux_2_39_riscv64.whl", hash = "sha256:adf2244d7f69409925a8f22ed22cc5f93cdfe5c9dc87ff3476be2c2aaae61a05"}, + {file = "greenlet-3.5.4-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:0fa53040b78b578120eecdc0265e3f1051487cc425d11a2b7c761daadf4feaa8"}, + {file = "greenlet-3.5.4-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:60e0bc961d367df506660e9ac0177a76bc6d81305300704b0977d1634f76efe2"}, + {file = "greenlet-3.5.4-cp315-cp315t-win_amd64.whl", hash = "sha256:f680e549edb3eaf21eea4e7fe101e15ec180c74b7879ab46adc080f22d4015d2"}, + {file = "greenlet-3.5.4-cp315-cp315t-win_arm64.whl", hash = "sha256:08fc36de8442d5c3e95b044550dbea9bf144d31ec0cc58e36fb241cb6ef6a994"}, + {file = "greenlet-3.5.4.tar.gz", hash = "sha256:0232ae1de90a8e07867bb127d7a6ba2301e859145489f25cda8a6096dabe1d20"}, +] + +[package.extras] +docs = ["Sphinx", "furo"] +test = ["objgraph", "psutil", "setuptools"] + [[package]] name = "griffe" version = "2.0.2" @@ -1266,6 +1430,26 @@ files = [ {file = "librt-0.7.4.tar.gz", hash = "sha256:3871af56c59864d5fd21d1ac001eb2fb3b140d52ba0454720f2e4a19812404ba"}, ] +[[package]] +name = "mako" +version = "1.3.12" +description = "A super-fast templating language that borrows the best ideas from the existing templating languages." +optional = false +python-versions = ">=3.8" +groups = ["main", "dev"] +files = [ + {file = "mako-1.3.12-py3-none-any.whl", hash = "sha256:8f61569480282dbf557145ce441e4ba888be453c30989f879f0d652e39f53ea9"}, + {file = "mako-1.3.12.tar.gz", hash = "sha256:9f778e93289bd410bb35daadeb4fc66d95a746f0b75777b942088b7fd7af550a"}, +] + +[package.dependencies] +MarkupSafe = ">=0.9.2" + +[package.extras] +babel = ["Babel"] +lingua = ["lingua"] +testing = ["pytest"] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -1296,7 +1480,7 @@ version = "3.0.3" description = "Safely add untrusted strings to HTML/XML markup." optional = false python-versions = ">=3.9" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "markupsafe-3.0.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559"}, {file = "markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419"}, @@ -1935,13 +2119,38 @@ files = [ antlr4-python3-runtime = "==4.9.*" PyYAML = ">=5.1.0" +[[package]] +name = "optuna" +version = "4.9.0" +description = "A hyperparameter optimization framework" +optional = false +python-versions = ">=3.9" +groups = ["main", "dev"] +files = [ + {file = "optuna-4.9.0-py3-none-any.whl", hash = "sha256:f52f3be6148654850c92a5860d398fd88ec6b2c84ab68d9c3d07dcff02e7afee"}, + {file = "optuna-4.9.0.tar.gz", hash = "sha256:b322e5cbdf1655fb84c37646c4a7a1f391de1b47806bbe222e015825d0a82b87"}, +] + +[package.dependencies] +alembic = ">=1.5.0" +colorlog = "*" +numpy = "*" +packaging = ">=20.0" +PyYAML = "*" +sqlalchemy = ">=1.4.2" +tqdm = "*" + +[package.extras] +document = ["ase", "cmaes (>=0.12.0)", "fvcore", "kaleido (!=0.2.1.post1,<0.4)", "lightgbm", "matplotlib (!=3.6.0)", "pandas", "pillow", "plotly (>=4.9.0)", "scikit-learn", "sphinx", "sphinx-copybutton", "sphinx-gallery", "sphinx-notfound-page", "sphinx_rtd_theme (>=1.2.0)", "torch", "torchvision"] +optional = ["boto3", "cmaes (>=0.12.0)", "google-cloud-storage", "greenlet", "grpcio", "matplotlib (!=3.6.0)", "pandas", "plotly (>=4.9.0)", "protobuf (>=5.28.1)", "redis", "scikit-learn (>=0.24.2)", "scipy", "torch"] + [[package]] name = "packaging" 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"}, @@ -2068,6 +2277,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 +2321,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 +2373,66 @@ 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"}, +] + [[package]] name = "pycparser" version = "2.23" @@ -2145,6 +2446,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 +2712,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"}, @@ -2337,11 +2795,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" @@ -2548,6 +3007,103 @@ files = [ cffi = ">=1.0" numpy = "*" +[[package]] +name = "sqlalchemy" +version = "2.0.51" +description = "Database Abstraction Library" +optional = false +python-versions = ">=3.7" +groups = ["main", "dev"] +files = [ + {file = "sqlalchemy-2.0.51-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0e8203d2fbd5c6254692ef0a72c740d75b2f3c7ca345404f4c1a4604813c77c0"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1af05726b3d0cdba1c55284bf408fd3b792e690fe2399bfb8304565551cda652"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2e54ff2dd657f2e3e0fbf2b097db1182f7bfea263eca4353f00065bae2a67c3d"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:1e47b1199c2e832e325eacabc8d32d2487f58c9358f97e9a00f5eb93c5680d84"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:c68568f3facf8f66fa76c60e0ced69b67666ffa9941d1d0a3756fda196049080"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-win32.whl", hash = "sha256:0592bdadf86ddcabfd72d9ab66ea8a5d8d2cc6be1cc51fa7e66c03868ac5eac1"}, + {file = "sqlalchemy-2.0.51-cp310-cp310-win_amd64.whl", hash = "sha256:740cf6f35351b1ac3d82369152acf1d51d37e3dcf85d4dc0a22ca01410eabe2a"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1aa10c0daee6705294d181daadaa793221e1a59ed55000a3fab1d42b088ce4ba"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a5b2ed6d828f1f09bd812861f4f59ca3bc3803f9df871f4555187f0faf018604"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:436728ce18a80f6951a1e11cc6112c2ede9faf20766f1a26195a7c441ca12dbd"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:dc261707bf5739aea8a541593f3cc1d463c2701fb05fbcbba0ce031b69a21260"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:a6d26094615306d116dd5e4a51b0304c99dd2356fc569eed6922a80a6bd3b265"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-win32.whl", hash = "sha256:ca8435d13829b92f4a97362d91975154a4015db3a2634154e1754e9a915e6b86"}, + {file = "sqlalchemy-2.0.51-cp311-cp311-win_amd64.whl", hash = "sha256:4a011ea4510683319ce4ed274b56ee05194b39b6da9d09ca7a39388f0fa84dcc"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7d78702b26ba1c18b2d0fb2ea940ba7f17a9581b42e8361ff93920ebbee1235a"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:581921d849d6e6f994d560389192955e80e2950e18fcdfe2ccea863e01158e6e"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1d21ce524ab86c23046e992a5b81cb54c21079c6df6e78b8fc77d77cac70a6b9"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c5d98a2709840027f5a347c3af0a7c3d5f6c1ff93af2ca1c54494e23cba8f389"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1181256e0f16479691b5616d36375dc2620ad8332b25978763c3d206ad3f3f1d"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-win32.whl", hash = "sha256:9f380393be5abeb6815f68fd39271b95127173511b6706b0a630a9995d53f8f5"}, + {file = "sqlalchemy-2.0.51-cp312-cp312-win_amd64.whl", hash = "sha256:2cf39aabdf48e87c1c2c2ed6d20d33ffa0733b3071ce9c5f66357947dd009080"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7c2056838b6685b72fdb36c99996cf862753461a62f2e84f4196371d3b2d6a07"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:483b11bd46bf35fc14c52faf338b04300c9e6ce554bce9b11be85bfec3bc3195"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1bed1ee8b01da6088210aa9412023326fb98a599ba502e6118308601dcbef77f"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:72ca54c952107ba5cd58854b67a5a6268631289d21651a1235396f3b98b47400"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b3e693d15533a45cd5906f0589f9c35090bef6ef45bf1e8195c424aa0ae06a8d"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-win32.whl", hash = "sha256:b93ab07b5292dbe7e6b8da89475275e7042744283921344b56105f3eeb0f828b"}, + {file = "sqlalchemy-2.0.51-cp313-cp313-win_amd64.whl", hash = "sha256:0f053118c30e53161857a953e4de667d90e274980dccbe5dd3829bbbeece72a5"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:6ea306caaae6bd5afd0a46050003c88f6bf33227377a49298c498c3cb88ff491"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c45a496d6bc05dec41dcd4c3a2b183723f47473255c159cd80b503c8f246424d"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4004ada0aafe8ae1991b2cd1d99c6d9146126e123bd6f883c260d974aa012e54"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:0f6bcad487aee1c638d707235682fc96f741de00663619881ab235400d03289e"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:39a76529db6305693d8d4affa58ad5b5e2e18edd62daea628b29b97930b3513d"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-win32.whl", hash = "sha256:08a204d8b5638717c26a24df18fcf40af45a6b22e35b70b1d62f0113c2e278e8"}, + {file = "sqlalchemy-2.0.51-cp314-cp314-win_amd64.whl", hash = "sha256:96747bfbadb055466e5b46d572618170046b45ce5a4879167f50d70a5319a499"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e5ea1a213be1fcd5e49d9904c3b9939211ded90bc2a64e93f4c01963474285de"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c6b36ed71f41942bdcd2ad2522be46bfce09d5705be5640ecf19bbc7660e4b7"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0c2c62877097e1a0db401fba5cb4debee33265e5b2a55c4ccb489c02c53b4f72"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0378d055e9e8cd6ce4d8dff683bdd3d7d413533c4ee51d67a2b1e0f9eacc0f23"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6e46fc36029eff666391e0531e5387b62ce6c4f1d8e50b3fb3099eaca1b42522"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-win32.whl", hash = "sha256:9161cfc9efce70d1715f47d6ff40f79c6778c00d53be4fbc09d70301e4b83ba7"}, + {file = "sqlalchemy-2.0.51-cp314-cp314t-win_amd64.whl", hash = "sha256:159bb6ba32059f57ad7375a8f50d844dd2f19d14954ecf820cd33e20debd46b2"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:bb1f5062f98b0b3290e72b707747fdd7e0f22d6956b236ba7ca7f5c9971d2da2"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:247acaa29ccef6250dfd6a3eedf8f94ddf23564180a39fe362e32ae9dbdbde46"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c95ef01f53233a305a874a44a63fbfb1d81cd79b49de0f8529b3548cde437e37"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:fa268106c8987639a17a18514cfe0cd9bf17420ab887e1e1bf486da8836135b1"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:b7f08588854bbb724041d9ae9d980d40040c922382e1d9a2ecb390edc4fd5032"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-win32.whl", hash = "sha256:6b588fd681ddf0c196b8df1ea49a8913514894b2b8f945a9511b4b48871f99c8"}, + {file = "sqlalchemy-2.0.51-cp38-cp38-win_amd64.whl", hash = "sha256:ca216e8af5c05e326efc7e28716ac2381a7cf9791749f5ee1849dccdc99c9b00"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:aa18ae738b5170e253ad0bb6c4b0f07585081e8a6e50893e4d911d47b39a0904"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:59cab3686b1bc039dd9cded2f8d0c08a246e84e76bd4ab5b4f18c7cdae293825"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:111604e637da87031255ddc26c7d7bc22bc6af6f5d459ccff3af1b4660233a85"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:ad30ae663711786303fbcd46a47516302d201ee49a877cb3fac61f672895110a"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:b21f0e7efc7a5c509e953784e9d1575ebb8b4318960e7e7d7a93bb803626cf64"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-win32.whl", hash = "sha256:a42ad6afcbaaa777241e347aa2e29155993045a0d6b7db74da61053ffe875fe0"}, + {file = "sqlalchemy-2.0.51-cp39-cp39-win_amd64.whl", hash = "sha256:2a97eaad21c84b4ef8010b11eeba9fe6153eb0b3df3ff8b6abc309df1b978ef7"}, + {file = "sqlalchemy-2.0.51-py3-none-any.whl", hash = "sha256:bb024d8b621d0be75f4f44ecc7c950450026e76d66dc8f791bb5331d7fed59d5"}, + {file = "sqlalchemy-2.0.51.tar.gz", hash = "sha256:804dccd8a4a6242c4e30ad961e540e18a588f6527202f2d6791b01845d59fdc9"}, +] + +[package.dependencies] +greenlet = {version = ">=1", markers = "platform_machine == \"aarch64\" or platform_machine == \"ppc64le\" or platform_machine == \"x86_64\" or platform_machine == \"amd64\" or platform_machine == \"AMD64\" or platform_machine == \"win32\" or platform_machine == \"WIN32\""} +typing-extensions = ">=4.6.0" + +[package.extras] +aiomysql = ["aiomysql (>=0.2.0)", "greenlet (>=1)"] +aioodbc = ["aioodbc", "greenlet (>=1)"] +aiosqlite = ["aiosqlite", "greenlet (>=1)", "typing_extensions (!=3.10.0.1)"] +asyncio = ["greenlet (>=1)"] +asyncmy = ["asyncmy (>=0.2.3,!=0.2.4,!=0.2.6)", "greenlet (>=1)"] +mariadb-connector = ["mariadb (>=1.0.1,!=1.1.2,!=1.1.5,!=1.1.10)"] +mssql = ["pyodbc"] +mssql-pymssql = ["pymssql"] +mssql-pyodbc = ["pyodbc"] +mypy = ["mypy (>=0.910)"] +mysql = ["mysqlclient (>=1.4.0)"] +mysql-connector = ["mysql-connector-python"] +oracle = ["cx_oracle (>=8)"] +oracle-oracledb = ["oracledb (>=1.0.1)"] +postgresql = ["psycopg2 (>=2.7)"] +postgresql-asyncpg = ["asyncpg", "greenlet (>=1)"] +postgresql-pg8000 = ["pg8000 (>=1.29.1)"] +postgresql-psycopg = ["psycopg (>=3.0.7)"] +postgresql-psycopg2binary = ["psycopg2-binary"] +postgresql-psycopg2cffi = ["psycopg2cffi"] +postgresql-psycopgbinary = ["psycopg[binary] (>=3.0.7)"] +pymysql = ["pymysql"] +sqlcipher = ["sqlcipher3_binary"] + [[package]] name = "sympy" version = "1.14.0" @@ -2572,8 +3128,8 @@ version = "2.3.0" description = "A lil' TOML parser" optional = false python-versions = ">=3.8" -groups = ["dev"] -markers = "python_version == \"3.10\"" +groups = ["main", "dev"] +markers = "python_version < \"3.11\"" files = [ {file = "tomli-2.3.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:88bd15eb972f3664f5ed4b57c1634a97153b4bac4479dcb6a495f41921eb7f45"}, {file = "tomli-2.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:883b1c0d6398a6a9d29b508c331fa56adbcdff647f6ace4dfca0f50e90dfd0ba"}, @@ -2730,7 +3286,7 @@ version = "4.67.1" description = "Fast, Extensible Progress Meter" optional = false python-versions = ">=3.7" -groups = ["dev"] +groups = ["main", "dev"] files = [ {file = "tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2"}, {file = "tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2"}, @@ -2833,7 +3389,22 @@ 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\""} + +[[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 +3424,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 +3494,10 @@ type = ["pytest-mypy"] [extras] full = ["nvidia-ml-py"] +migrate = ["pyarrow", "wandb"] +sweep = ["optuna"] [metadata] lock-version = "2.1" python-versions = "^3.10" -content-hash = "6668aa67f414632c07ec8050bd4eee21670c916759d41644f43e12cdf36680f6" +content-hash = "f35c43dde82514d7ec12ca5461fcec7e5c14b50edc603eb3211ebd33db08c71d" diff --git a/pyproject.toml b/pyproject.toml index f37c06e..22b8136 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,11 @@ 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 } +# sweep extra: Bayesian search for pluto.sweep(method="bayes") +optuna = { version = ">=3.0", optional = true } [tool.poetry.group.dev.dependencies] mypy = "^1.10.0" @@ -50,12 +55,21 @@ torchvision = "*" imageio = "^2.37.0" moviepy = "^1.0.3" omegaconf = "^2.3" +pyarrow = ">=15" # pluto.migrate tests +optuna = ">=3.0" # pluto.sweep bayes tests [tool.poetry.extras] full = [ "nvidia-ml-py", ] +migrate = [ + "wandb", + "pyarrow", +] +sweep = [ + "optuna", +] [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..f1f4c23 --- /dev/null +++ b/tests/test_auth_transient.py @@ -0,0 +1,50 @@ +""" +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.webbrowser, 'open'), # keep the test hermetic + 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_iface_errors.py b/tests/test_iface_errors.py index efdaca3..86aad27 100644 --- a/tests/test_iface_errors.py +++ b/tests/test_iface_errors.py @@ -147,6 +147,90 @@ def fake_method(url, content=None, headers=None, **kwargs): assert r is None +def test_try_connection_reset_default_is_shutdown_signal_no_retry(): + """By default a dropped keep-alive socket is treated as a shutdown signal: + no retry, no raise (so heartbeat/trigger/upload spam doesn't hang atexit).""" + iface = _make_iface() + calls = {'n': 0} + + def fake_method(url, content=None, headers=None, **kwargs): + calls['n'] += 1 + raise ConnectionResetError('peer reset') + + r = iface._try( + fake_method, 'https://x', {}, b'{}', name='trigger', raise_on_error=True + ) + assert r is None + assert calls['n'] == 1, 'default: a connection reset is not retried' + + +def test_try_connection_reset_retried_when_flagged_then_recovers(): + """For critical one-shot requests (retry_connection_errors=True, e.g. the + terminal status update) a dropped socket is transient and IS retried, so a + later attempt can succeed instead of silently giving up on attempt 1.""" + iface = _make_iface() + calls = {'n': 0} + + def fake_method(url, content=None, headers=None, **kwargs): + calls['n'] += 1 + if calls['n'] < 3: + raise ConnectionResetError('peer reset') + return _resp(200, json_body={'ok': True}) + + r = iface._try( + fake_method, + 'https://x/api/runs/status/update', + {}, + b'{}', + name='status', + raise_on_error=True, + retry_connection_errors=True, + ) + assert r is not None and r.status_code == 200 + assert calls['n'] == 3, 'a connection reset must be retried when flagged' + + +def test_try_connection_reset_raises_after_retries_when_flagged(): + """A reset that never clears raises (not None) so the caller can't mistake a + dropped terminal-status POST for success.""" + iface = _make_iface() # x_file_stream_retry_max = 2 + calls = {'n': 0} + + def fake_method(url, content=None, headers=None, **kwargs): + calls['n'] += 1 + raise ConnectionResetError('peer reset') + + with pytest.raises(PlutoRequestError): + iface._try( + fake_method, + 'https://x/api/runs/status/update', + {}, + b'{}', + name='status', + raise_on_error=True, + retry_connection_errors=True, + ) + assert calls['n'] == 3, 'initial attempt + 2 retries' + + +def test_update_status_retries_reset_and_raises_on_persistent_failure(monkeypatch): + """update_status() is the run's terminal transition: it must retry a dropped + socket and ultimately raise, never silently swallow it (the finish-strands- + run bug). Wires the retry_connection_errors + raise_on_error flags end-to-end.""" + iface = _make_iface() + calls = {'n': 0} + + def fake_post(url, content=None, headers=None, **kwargs): + calls['n'] += 1 + raise ConnectionResetError('peer reset') + + monkeypatch.setattr(iface.client_api, 'post', fake_post) + + with pytest.raises(PlutoRequestError): + iface.update_status() + assert calls['n'] >= 2, 'update_status must retry the reset before failing' + + # --- sync-process uploader (pluto/sync/process.py) ------------------------- 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..9c7d2e0 --- /dev/null +++ b/tests/test_migrate_cli.py @@ -0,0 +1,474 @@ +""" +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() + # `all` pipelines export + load, so load runs one or more passes. + assert loader.load.call_count >= 1 + + 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), + ] + ) + assert loader.load.call_count >= 1 # staged runs still load (pipelined) + assert code == 1 # but the failure is reported + + def test_all_export_crash_is_not_reported_as_success(self, tmp_path): + # A crash in the export thread must surface as a non-zero exit, not 0. + loader = _mock_loader({'loaded': 0, 'skipped': 0, 'failed': []}) + exporter = _mock_exporter() + exporter.export.side_effect = RuntimeError('wandb auth exploded') + 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 == 2 # export crash -> failure, not silent success + + def test_single_project_worker_crash_exits_clean(self, tmp_path): + # A single --project whose exporter raises should return 2 with a clean + # message, not propagate a raw traceback. + exporter = _mock_exporter() + exporter.export.side_effect = RuntimeError('wandb auth exploded') + 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 == 2 + + def test_all_forwards_run_id_filter_to_load(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) as cls, + ): + run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + '--run-id', + 'r1', + '--run-id', + 'r2', + ] + ) + assert cls.call_args.kwargs['run_ids'] == ['r1', 'r2'] + + def test_all_attempts_failed_run_once_via_skip_run_ids(self, tmp_path): + # A run that fails in one poll pass must not be re-attempted on the next + # (re-running identical staged data can't help and risks duplicate media). + # The cli feeds already-failed run-ids into the loader's skip_run_ids, so + # each failure is at-most-once and the reported failure count stays exact. + import time + + exporter = _mock_exporter() + + def _slow_export(*a, **k): + time.sleep(0.25) # stay alive across several fast poll passes + return {'exported': 1, 'skipped': 0, 'failed': []} + + exporter.export.side_effect = _slow_export + loader = _mock_loader( + {'loaded': 0, 'skipped': 0, 'failed': [{'run_id': 'boom', 'error': 'e'}]} + ) + with ( + mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ), + mock.patch('pluto.migrate.loader.PlutoLoader', return_value=loader) as cls, + mock.patch('pluto.migrate.cli._ALL_POLL_SECONDS', 0.02), + ): + code = run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + ] + ) + assert code == 1 # the single failure is reported once + # First pass skips nothing; once 'boom' has failed, every later pass + # passes it into skip_run_ids so the loader bypasses it. + assert cls.call_args_list[0].kwargs['skip_run_ids'] == [] + assert any(c.kwargs['skip_run_ids'] == ['boom'] for c in cls.call_args_list[1:]) + + 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_strict_fails_when_data_not_migrated(self, tmp_path): + cov = {'migrated': {'metric': 10}, 'not_migrated': {'unsupported(bokeh)': 3}} + exporter = _mock_exporter( + {'exported': 1, 'skipped': 0, 'failed': [], 'coverage': cov} + ) + args = [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + ] + with mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ): + # without --strict: dropped data is reported but exit stays 0 + assert run_migrate(args) == 0 + # with --strict: non-zero exit + assert run_migrate(args + ['--strict']) == 2 + + def test_strict_passes_when_full_coverage(self, tmp_path): + cov = {'migrated': {'metric': 10}, 'not_migrated': {}} + exporter = _mock_exporter( + {'exported': 1, 'skipped': 0, 'failed': [], 'coverage': cov} + ) + with mock.patch( + 'pluto.migrate.wandb_export.WandbExporter', return_value=exporter + ): + code = run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'vision', + '--output', + str(tmp_path), + '--strict', + ] + ) + assert code == 0 + + def test_export_all_projects_when_no_project_given(self, tmp_path): + with ( + mock.patch( + 'pluto.migrate.wandb_export.list_wandb_projects', + return_value=['p1', 'p2', 'p3'], + ) as lst, + mock.patch('pluto.migrate.cli._export_one_project', return_value=0) as w, + ): + code = run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--output', + str(tmp_path), + '--workers', + '1', + ] + ) + assert code == 0 + lst.assert_called_once() # listed all projects under the entity + assert {c.args[1] for c in w.call_args_list} == {'p1', 'p2', 'p3'} + + def test_exclude_drops_projects(self, tmp_path): + with ( + mock.patch( + 'pluto.migrate.wandb_export.list_wandb_projects', + return_value=['p1', 'p2', 'p3'], + ), + mock.patch('pluto.migrate.cli._export_one_project', return_value=0) as w, + ): + run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--output', + str(tmp_path), + '--workers', + '1', + '--exclude', + 'p2', + ] + ) + assert {c.args[1] for c in w.call_args_list} == {'p1', 'p3'} + + def test_single_project_skips_project_listing(self, tmp_path): + with ( + mock.patch('pluto.migrate.wandb_export.list_wandb_projects') as lst, + mock.patch('pluto.migrate.cli._export_one_project', return_value=0) as w, + ): + run_migrate( + [ + 'wandb', + 'export', + '--entity', + 'acme', + '--project', + 'only', + '--output', + str(tmp_path), + ] + ) + lst.assert_not_called() # explicit --project => no account listing + assert w.call_args.args[1] == 'only' + + def test_dest_project_rejected_for_multiple_projects(self, tmp_path): + with mock.patch( + 'pluto.migrate.wandb_export.list_wandb_projects', + return_value=['p1', 'p2'], + ): + code = run_migrate( + [ + 'wandb', + 'all', + '--entity', + 'acme', + '--output', + str(tmp_path), + '--dest-project', + 'combined', + '--workers', + '1', + ] + ) + assert code == 2 # can't rename many projects into one + + def test_workers_over_project_count_reports_cap(self, tmp_path, capsys): + # Two staged (empty) projects; asking for more workers than projects + # must announce the cap, not silently print the requested count. + for p in ('p1', 'p2'): + (tmp_path / 'acme' / p / 'runs').mkdir(parents=True) + code = run_migrate( + ['wandb', 'load', '--input', str(tmp_path), '--workers', '16'] + ) + assert code == 0 # nothing staged inside -> clean no-op + out = capsys.readouterr().out + assert 'requested 16' in out + assert 'capped to 2 projects' in out + + 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..bc39561 --- /dev/null +++ b/tests/test_migrate_loader.py @@ -0,0 +1,1111 @@ +""" +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, + read_json, + write_json_atomic, +) + +CREATED_AT_MS = 1746093600000 +UPDATED_AT_MS = CREATED_AT_MS + 7200000 +T0_MS = 1746093601000 +EXTERNAL_ID = 'wandb::acme/vision/abc123' +# The load cache is keyed by (external id, destination project); _stage_run +# uses project 'vision' and tests don't override dest_project. +CACHE_KEY = f'{EXTERNAL_ID}@@vision' + + +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 + + +def _stage_media_table_run(tmp_path, run_id='mtable'): + """Stage a run with a wandb media table (an image column). + + Mirrors what the exporter actually stages today: the ``table-file`` row + points at wandb's *run-files* copy of the table, where media cells are + collapsed to the literal string ``"Image"`` (the full-fidelity + ``{_type: image-file, path, sha256}`` refs live only in the separate + artifact copy). The cell images are staged as their own artifact. + """ + 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': 'rich-table-1', + 'notes': '', + 'tags': [], + 'state': 'finished', + 'config': {}, + 'summary': {}, + 'createdAt': CREATED_AT_MS, + 'updatedAt': UPDATED_AT_MS, + 'url': f'https://wandb.ai/acme/vision/runs/{run_id}', + }, + ) + # wandb's lossy run-files table copy: image cells are the string "Image". + table_rel = 'files/media/table/media_table.table.json' + table_file = run_dir / table_rel + table_file.parent.mkdir(parents=True) + write_json_atomic( + table_file, + { + 'columns': ['idx', 'img', 'score'], + 'data': [[0, 'Image', 0.0], [1, 'Image', 0.25]], + }, + ) + # the cell images arrive separately, as the auto-created run-table artifact. + img_rel = 'artifacts/run-mtable-media_table:v0/media/images/a.png' + img_file = run_dir / img_rel + img_file.parent.mkdir(parents=True) + img_file.write_bytes(b'PNG') + + base = dict(project_id='acme/vision', run_id=run_id) + with PartWriter(run_dir) as w: + w.write_row( + **base, + attribute_path='media_table', + attribute_type='media', + step=0, + timestamp_ms=T0_MS, + string_value='table-file', + file_value=table_rel, + ) + w.write_row( + **base, + attribute_path='run-mtable-media_table:v0', + attribute_type='artifact', + step=0, + timestamp_ms=CREATED_AT_MS, + string_value=json.dumps( + {'name': 'run-mtable-media_table:v0', 'type': 'run_table'} + ), + file_value=img_rel, + ) + mark_run_exported(run_dir, {'rows': 2}) + return run_dir + + +def _stage_string_series_run(tmp_path, values, run_id='ssrun', name='phase'): + """Stage a minimal run whose only history is one string_series attribute.""" + 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', + 'state': 'finished', + 'config': {}, + 'summary': {name: values[-1]}, + 'createdAt': CREATED_AT_MS, + 'updatedAt': UPDATED_AT_MS, + 'url': f'https://wandb.ai/acme/vision/runs/{run_id}', + }, + ) + base = dict(project_id='acme/vision', run_id=run_id) + with PartWriter(run_dir) as w: + for i, v in enumerate(values): + w.write_row( + **base, + attribute_path=name, + attribute_type='string_series', + step=i, + timestamp_ms=T0_MS + i * 1000, + string_value=v, + ) + mark_run_exported(run_dir, {'rows': len(values)}) + 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 + # A real Op sets this to None on a confirmed finish; without it a bare + # MagicMock auto-vivifies a truthy attribute and looks like a failed + # status update to the loader. + op._status_update_error = None + 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_custom_charts_forwarded_to_wandb_config(self, tmp_path, mock_init): + _, op = mock_init + run_dir = _stage_run(tmp_path) + panels = [ + { + 'key': 'bar', + 'preset': 'bar', + 'title': 'per-class', + 'tableKey': 'bar_table', + 'fields': {'label': 'label', 'value': 'value'}, + 'specLang': 'vega-lite', + } + ] + write_json_atomic(run_dir / 'custom_charts.json', {'panels': panels}) + PlutoLoader(tmp_path).load() + wandb_block = op.update_config.call_args.args[0]['wandb'] + assert wandb_block['custom_charts'] == panels + + def test_image_annotations_resolves_boxes_and_masks(self, tmp_path): + # Boxes: the .boxes2D.json is inlined into the annotations JSON. Masks: + # the .mask.png is resolved to a path spec for pluto.Image to re-upload. + run_dir = tmp_path / 'run' + box_dir = run_dir / 'files' / 'media' / 'metadata' / 'boxes2D' + box_dir.mkdir(parents=True) + box_content = {'box_data': [{'class_id': 1}], 'class_labels': {'1': 'cat'}} + (box_dir / 'a.boxes2D.json').write_text(json.dumps(box_content)) + mask_dir = run_dir / 'files' / 'media' / 'images' / 'mask' + mask_dir.mkdir(parents=True) + (mask_dir / 'a.mask.png').write_bytes(b'PNG') + annotation_value = json.dumps( + { + 'boxes': {'pred': {'path': 'media/metadata/boxes2D/a.boxes2D.json'}}, + 'masks': {'pred': {'path': 'media/images/mask/a.mask.png'}}, + } + ) + boxes_str, masks_spec = PlutoLoader(tmp_path)._image_annotations( + run_dir, annotation_value + ) + assert json.loads(boxes_str) == {'boxes': {'pred': box_content}} + assert masks_spec['pred']['path'].endswith('media/images/mask/a.mask.png') + + def test_mask_class_labels_preserved(self, tmp_path): + # A staged mask carrying class_labels (folded in by the exporter) keeps + # them in the mask spec, so pluto.Image re-uploads a coloured mask + # instead of a blank one. + run_dir = tmp_path / 'run' + mask_dir = run_dir / 'files' / 'media' / 'images' / 'mask' + mask_dir.mkdir(parents=True) + (mask_dir / 'a.mask.png').write_bytes(b'PNG') + annotation_value = json.dumps( + { + 'masks': { + 'pred': { + 'path': 'media/images/mask/a.mask.png', + 'class_labels': {'0': 'bg', '1': 'cat'}, + } + } + } + ) + _, masks_spec = PlutoLoader(tmp_path)._image_annotations( + run_dir, annotation_value + ) + assert masks_spec['pred']['class_labels'] == {'0': 'bg', '1': 'cat'} + assert masks_spec['pred']['path'].endswith('a.mask.png') + + def test_image_annotations_none_when_absent(self, tmp_path): + assert PlutoLoader(tmp_path)._image_annotations(tmp_path, None) == (None, None) + + def test_sweep_membership_migrated(self, tmp_path, mock_init): + init, op = mock_init + run_dir = _stage_run(tmp_path) + # inject a sweep block into the staged manifest (as the exporter would) + manifest = read_json(run_dir / 'run.json') + manifest['sweep'] = { + 'id': 'swp1', + 'name': 'my-sweep', + 'config': {'method': 'grid', 'parameters': {'lr': {'values': [0.1]}}}, + } + write_json_atomic(run_dir / 'run.json', manifest) + PlutoLoader(tmp_path).load() + # run is tagged so it groups under its sweep (like native pluto.sweep) + assert 'sweep:swp1' in init.call_args.kwargs['tags'] + # search space survives in the wandb config block + wandb_block = op.update_config.call_args.args[0]['wandb'] + assert wandb_block['sweep']['id'] == 'swp1' + assert wandb_block['sweep']['config']['method'] == 'grid' + + def test_no_custom_charts_key_when_absent(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 'custom_charts' not in wandb_block + + 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_media_in_table_migrates_degraded_not_dropped(self, tmp_path, mock_init): + # Characterization test for the known media-in-table gap. wandb keeps + # two copies of a media table; the exporter stages the lossy run-files + # copy, so image cells load as the literal text "Image". The table is + # NOT dropped and does NOT crash, and the cell images arrive as a + # separate, unlinked artifact. When we wire cells to real images this + # test should flip to assert image cells. + _, op = mock_init + _stage_media_table_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + + table_calls = _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Table) for v in d.values()) + ) + assert len(table_calls) == 1 + table = table_calls[0].args[0]['media_table'] + assert table._col == ['idx', 'img', 'score'] + # image column (index 1) is the degraded literal string, not a picture + assert [row[1] for row in table._table] == ['Image', 'Image'] + + # the cell images still migrate, but as a disconnected artifact + art_calls = _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Artifact) for v in d.values()) + ) + assert len(art_calls) == 1 + + 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_string_series_sent_to_ingest_with_data_logtype(self, tmp_path, mock_init): + _, op = mock_init + op.settings.url_data = 'http://ingest/data' + op.settings.url_meta = 'http://api/logName/add' + _stage_string_series_run(tmp_path, values=['warmup', 'train', 'train', 'eval']) + with mock.patch('pluto.migrate.loader.httpx') as httpx_mock: + PlutoLoader(tmp_path).load() + calls = httpx_mock.post.call_args_list + assert len(calls) == 2 + # 1st POST registers the log name with logType DATA. + meta = json.loads(calls[0].kwargs['content']) + assert meta['logType'] == 'DATA' + assert meta['logName'] == ['phase'] + # 2nd POST ingests the points as NDJSON string-series (raw data). + lines = [json.loads(x) for x in calls[1].kwargs['content'].strip().split('\n')] + assert [x['step'] for x in lines] == [0, 1, 2, 3] + assert [x['data'] for x in lines] == ['warmup', 'train', 'train', 'eval'] + assert all( + x['dataType'] == 'string-series' and x['logName'] == 'phase' for x in lines + ) + + def test_string_series_high_cardinality_still_sent(self, tmp_path, mock_init): + _, op = mock_init + op.settings.url_data = 'http://ingest/data' + op.settings.url_meta = 'http://api/logName/add' + # No cardinality guard: even an all-distinct series is sent in full + # (no data loss). + _stage_string_series_run(tmp_path, values=[f'v{i}' for i in range(60)]) + with mock.patch('pluto.migrate.loader.httpx') as httpx_mock: + PlutoLoader(tmp_path).load() + calls = httpx_mock.post.call_args_list + assert len(calls) == 2 # logName/add + ingest + lines = [json.loads(x) for x in calls[1].kwargs['content'].strip().split('\n')] + assert len(lines) == 60 # every point sent + + 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(CACHE_KEY) + + init.reset_mock() + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 0, 'skipped': 1, 'failed': []} + init.assert_not_called() + + def _stage_min_run(self, tmp_path, project, run_id): + d = tmp_path / 'acme' / project / 'runs' / run_id + d.mkdir(parents=True) + write_json_atomic( + d / 'run.json', + { + 'entity': 'acme', + 'project': project, + 'run_id': run_id, + 'name': run_id, + 'state': 'finished', + }, + ) + with PartWriter(d) as w: + w.write_row( + project_id=f'acme/{project}', + run_id=run_id, + attribute_path='loss', + attribute_type='metric', + step=0, + timestamp_ms=T0_MS, + float_value=1.0, + ) + mark_run_exported(d, {'rows': 1}) + + def test_discover_respects_project_scope(self, tmp_path): + self._stage_min_run(tmp_path, 'vision', 'v1') + self._stage_min_run(tmp_path, 'audio', 'a1') + self._stage_min_run(tmp_path, 'text', 't1') + # include only + got = { + d.parent.parent.name + for d in PlutoLoader( + tmp_path, projects=['vision', 'audio'] + )._discover_runs() + } + assert got == {'vision', 'audio'} + # exclude + got2 = { + d.parent.parent.name + for d in PlutoLoader(tmp_path, exclude_projects=['vision'])._discover_runs() + } + assert got2 == {'audio', 'text'} + # default = all + got3 = {d.parent.parent.name for d in PlutoLoader(tmp_path)._discover_runs()} + assert got3 == {'vision', 'audio', 'text'} + + def test_custom_cache_path(self, tmp_path, mock_init): + _stage_run(tmp_path) + custom = tmp_path / 'ledger-vision.json' + PlutoLoader(tmp_path, cache_path=custom).load() + assert custom.exists() # resume ledger went to the custom path + assert not (tmp_path / 'loaded_runs.json').exists() + + def test_external_id_collision_restores_and_skips_by_default( + self, tmp_path, mock_init + ): + # A collision means the run already exists server-side but isn't in the + # local cache. Re-replaying would duplicate media, so the loader skips + # the replay. BUT the create-with-existing already reopened the run to + # RUNNING (DDP-style resume), so it re-attaches (resume) and finish()es + # to restore the terminal status + historical timestamp — without + # replaying — then skips. + from pluto.op import RunExistsError + + init, _ = mock_init + restore_op = mock.MagicMock() + restore_op._status_update_error = None # finish() confirmed the status + init.side_effect = [ + RunExistsError( + "Run with externalId 'wandb::acme/vision/abc123' already exists." + ), + restore_op, # the resume re-attach used to restore terminal status + ] + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 0, 'skipped': 1, 'failed': []} + assert init.call_count == 2 # plain init (raises) + resume to restore + restore_op.finish.assert_called_once() # restored to terminal status + restore_op._log_metrics_batch.assert_not_called() # NOT re-replayed + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_external_id_collision_replays_with_force_resume(self, tmp_path, mock_init): + # force_resume opts into healing a genuinely interrupted load: resume + # the existing run and re-replay (accepting possible media duplication). + 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, force_resume=True).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + assert init.call_count == 2 + 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(CACHE_KEY) + + def test_interrupted_run_resumes_and_completes_on_rerun(self, tmp_path, mock_init): + # A run created server-side that then crashes mid-replay is marked + # in_progress (not done). On a plain re-run the loader must recognize it + # started that run and resume to COMPLETE it — not skip it as though it + # were fully loaded elsewhere. + from pluto.op import RunExistsError + + init, op = mock_init + _stage_run(tmp_path) + # 1st init succeeds (run created); 2nd re-run init collides; 3rd is the + # resume that completes it. + init.side_effect = [op, RunExistsError('exists'), op] + with mock.patch.object( + PlutoLoader, + '_replay_run', + side_effect=[RuntimeError('crash mid-replay'), None], + ): + s1 = PlutoLoader(tmp_path).load() # crashes mid-replay + assert s1['loaded'] == 0 and len(s1['failed']) == 1 + cache = LoadedCache(tmp_path / 'loaded_runs.json') + assert cache.is_in_progress(CACHE_KEY) # marker persisted + assert not cache.is_loaded(CACHE_KEY) + + s2 = PlutoLoader(tmp_path).load() # re-run resumes + completes + assert s2 == {'loaded': 1, 'skipped': 0, 'failed': []} + assert init.call_args_list[-1].kwargs['resume'] is True # healed via resume + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_collision_restore_failure_reported_not_marked_loaded( + self, tmp_path, mock_init + ): + # Collision restore path: if re-attaching to finish() the existing run + # throws, the run is still RUNNING server-side with no terminal status. + # It must NOT be marked loaded (that would strand it RUNNING and skip it + # forever) — it's reported failed so a later run retries it. + from pluto.op import RunExistsError + + init, _ = mock_init + restore_op = mock.MagicMock() + restore_op.finish.side_effect = RuntimeError('server rejected finish') + init.side_effect = [ + RunExistsError( + "Run with externalId 'wandb::acme/vision/abc123' already exists." + ), + restore_op, # resume re-attach; its finish() blows up below + ] + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 0 and summary['skipped'] == 0 + assert [f['run_id'] for f in summary['failed']] == ['abc123'] + assert 'restore failed' in summary['failed'][0]['error'] + # Crucially: NOT marked loaded, so a re-run gets another shot. + assert not LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_unconfirmed_terminal_status_reported_not_marked_loaded( + self, tmp_path, mock_init + ): + # finish() replayed everything but could NOT confirm the run's terminal + # status on the server (a dropped connection that outlasted retries, now + # surfaced via op._status_update_error). The run must be recorded as + # failed and NOT cached as loaded — else a re-run skips it and it stays + # stranded RUNNING/FAILED forever (the real-world bug this guards). + _, op = mock_init + op._status_update_error = ConnectionResetError('peer reset during finish') + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 0 and summary['skipped'] == 0 + assert [f['run_id'] for f in summary['failed']] == ['abc123'] + assert 'status unconfirmed' in summary['failed'][0]['error'] + assert not LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_restore_unconfirmed_terminal_status_reported_not_marked_loaded( + self, tmp_path, mock_init + ): + # Same guarantee on the collision→restore path: if the restoring + # finish() couldn't confirm the terminal status, don't mark it loaded. + from pluto.op import RunExistsError + + init, _ = mock_init + restore_op = mock.MagicMock() + restore_op._status_update_error = ConnectionResetError('reset during restore') + init.side_effect = [ + RunExistsError( + "Run with externalId 'wandb::acme/vision/abc123' already exists." + ), + restore_op, + ] + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 0 and summary['skipped'] == 0 + assert [f['run_id'] for f in summary['failed']] == ['abc123'] + assert not LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_stale_reap_healed_via_resume(self, tmp_path, mock_init): + # A finished import comes back FAILED (a stale-run reap silently rejected + # our COMPLETED). The loader must detect it via read-back and heal it with + # a resume+finish, ending up loaded — not stranded. + init, op = mock_init + _stage_run(tmp_path) + with mock.patch.object( + PlutoLoader, '_read_run_status', side_effect=['FAILED', 'COMPLETED'] + ): + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + # the heal reopened the run: a second init with resume=True + assert any(c.kwargs.get('resume') is True for c in init.call_args_list) + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_stale_reap_unhealable_reported_not_loaded(self, tmp_path, mock_init): + # If the run stays FAILED after the bounded heal attempts, it is recorded + # failed (a later pass retries) and NOT cached as loaded. + _, op = mock_init + _stage_run(tmp_path) + with mock.patch.object(PlutoLoader, '_read_run_status', return_value='FAILED'): + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 0 + assert [f['run_id'] for f in summary['failed']] == ['abc123'] + assert 'not confirmed' in summary['failed'][0]['error'] + assert not LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_unverifiable_status_proceeds_best_effort(self, tmp_path, mock_init): + # If we can't read the status back at all (endpoint down / network) and no + # heal was attempted, don't fail a run over our inability to verify. + _, op = mock_init + _stage_run(tmp_path) + with mock.patch.object(PlutoLoader, '_read_run_status', return_value=None): + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + + def test_heal_finish_unconfirmed_reported_not_loaded(self, tmp_path, mock_init): + # The heal's OWN finish can drop its terminal-status update (recorded, not + # raised). If it does AND we then can't read the status back, the run may + # still be RUNNING — it must be reported failed (a re-run retries), NOT + # best-effort marked loaded. Guards the Cursor "heal ignores status update + # failure" finding. + _, op = mock_init + calls = {'n': 0} + + def _finish(code=None): + calls['n'] += 1 + if calls['n'] >= 2: # the heal's finish (2nd overall) can't confirm + op._status_update_error = ConnectionResetError('heal finish dropped') + + op.finish.side_effect = _finish + _stage_run(tmp_path) + # first read: FAILED (triggers heal); then unreadable (can't confirm) + with mock.patch.object( + PlutoLoader, '_read_run_status', side_effect=['FAILED', None, None, None] + ): + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 0 + assert [f['run_id'] for f in summary['failed']] == ['abc123'] + assert not LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_wandb_failed_run_not_status_verified(self, tmp_path, mock_init): + # A run that failed on wandb is intended FAILED; a FAILED status is + # correct, so the loader must not try to "heal" it to COMPLETED — the + # read-back is skipped entirely for code=1. + _, op = mock_init + _stage_run(tmp_path, run_id='crashed1', state='crashed') + with mock.patch.object( + PlutoLoader, '_read_run_status', return_value='FAILED' + ) as read: + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 1 + op.finish.assert_called_once_with(code=1) + read.assert_not_called() + + def test_read_run_status_parses_and_degrades(self, tmp_path): + loader = PlutoLoader(tmp_path) + op = mock.MagicMock() + op.settings._op_id = 7 + op.settings.url_api = 'http://api' + op.settings._auth = 'tok' + ok = mock.MagicMock(status_code=200) + ok.json.return_value = {'status': 'COMPLETED'} + with mock.patch('pluto.migrate.loader.httpx.get', return_value=ok) as g: + assert loader._read_run_status(op) == 'COMPLETED' + assert 'api/runs/details/7' in g.call_args.args[0] + # a non-200 and an exception both degrade to None (best-effort) + with mock.patch( + 'pluto.migrate.loader.httpx.get', + return_value=mock.MagicMock(status_code=404), + ): + assert loader._read_run_status(op) is None + with mock.patch('pluto.migrate.loader.httpx.get', side_effect=Exception('x')): + assert loader._read_run_status(op) is None + + def test_skip_run_ids_bypasses_run_without_attempting(self, tmp_path, mock_init): + # A run whose id is in skip_run_ids (already attempted-and-failed in an + # earlier `all` pass) is skipped outright — never re-initialized, never + # counted as loaded/skipped/failed. + init, _ = mock_init + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path, skip_run_ids=['abc123']).load() + assert summary == {'loaded': 0, 'skipped': 0, 'failed': []} + init.assert_not_called() + + 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()) + ) + + def test_unreadable_manifest_recorded_and_batch_continues( + self, tmp_path, mock_init + ): + # One run's run.json is corrupt; the other must still load and the bad + # one is recorded in failed[] rather than aborting the whole batch. + _stage_run(tmp_path, run_id='good1') + bad = _stage_run(tmp_path, run_id='bad1') + (bad / 'run.json').write_text('{ truncated') # invalid JSON + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 1 + assert len(summary['failed']) == 1 + assert summary['failed'][0]['run_id'] == 'bad1' + + def test_init_failure_recorded_and_batch_continues(self, tmp_path, mock_init): + # A non-RunExistsError from pluto.init on one run must not abort the + # batch; it is recorded as failed and the next run still loads. + init, op = mock_init + _stage_run(tmp_path, run_id='aaa1') + _stage_run(tmp_path, run_id='bbb2') + init.side_effect = [ConnectionError('server down'), op] + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 1 + assert len(summary['failed']) == 1 + assert 'ConnectionError' in summary['failed'][0]['error'] + + def test_dead_sync_process_fails_fast(self, tmp_path, mock_init): + # If the sync subprocess has exited, backpressure must raise (recorded + # as a failed run) instead of sleeping out the full stall_timeout. + _, op = mock_init + op._sync_manager.get_pending_count.return_value = 999 + op._sync_manager._process.poll.return_value = 1 # exited, code 1 + _stage_run(tmp_path) + summary = PlutoLoader(tmp_path, max_pending=5, stall_timeout=600).load() + assert summary['loaded'] == 0 + assert len(summary['failed']) == 1 + assert 'sync process exited' in summary['failed'][0]['error'] + + def test_path_traversal_media_refused(self, tmp_path, mock_init): + # A media row whose file_value escapes the run dir must be refused, not + # read+uploaded off the host. + _, op = mock_init + run_dir = tmp_path / 'acme' / 'vision' / 'runs' / 'evil1' + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': 'vision', + 'run_id': 'evil1', + 'name': 'evil', + 'state': 'finished', + }, + ) + outside = tmp_path / 'secret.txt' + outside.write_bytes(b'top secret') + with PartWriter(run_dir) as w: + w.write_row( + project_id='acme/vision', + run_id='evil1', + attribute_path='sneaky', + attribute_type='media', + step=0, + timestamp_ms=T0_MS, + string_value='image-file', + file_value='../../../../secret.txt', + ) + mark_run_exported(run_dir, {'rows': 1}) + summary = PlutoLoader(tmp_path).load() + assert summary['loaded'] == 1 # run itself loads + # ...but the out-of-bounds file was never turned into an upload. + assert not _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Image) for v in d.values()) + ) + + def test_dry_run_reports_would_load(self, tmp_path, mock_init, capsys): + _stage_run(tmp_path) + PlutoLoader(tmp_path, dry_run=True).load() + out = capsys.readouterr().out + assert 'would load 1' in out + + def test_histogram_with_null_bins_loads(self, tmp_path, mock_init): + # wandb often stores histogram counts with bins=None; this must not + # crash the run (real data hit `len(None)` in pluto.Histogram). + _, op = mock_init + run_dir = tmp_path / 'acme' / 'vision' / 'runs' / 'histrun' + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': 'vision', + 'run_id': 'histrun', + 'name': 'h', + 'state': 'finished', + }, + ) + with PartWriter(run_dir) as w: + w.write_row( + project_id='acme/vision', + run_id='histrun', + attribute_path='weights', + attribute_type='media', + step=0, + timestamp_ms=T0_MS, + string_value=json.dumps( + {'_type': 'histogram', 'values': [3, 1, 4, 1], 'bins': None} + ), + ) + mark_run_exported(run_dir, {'rows': 1}) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + # the histogram was logged (synthesized bins), not dropped + assert _log_calls_with( + op, lambda d: any(isinstance(v, pluto.Histogram) for v in d.values()) + ) + + def test_multi_image_step_logged_as_ordered_list(self, tmp_path, mock_init): + # An image gallery (many media rows sharing name+step+timestamp) must be + # logged in ONE op.log call as a list, so each image gets its sampleIndex + # (0,1,2) and the server preserves logged order — not 3 separate calls + # that all collapse to sampleIndex 0. + _, op = mock_init + run_dir = tmp_path / 'acme' / 'vision' / 'runs' / 'galrun' + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': 'vision', + 'run_id': 'galrun', + 'name': 'g', + 'state': 'finished', + }, + ) + for i in range(3): + f = run_dir / 'files' / f'idx{i}.png' + f.parent.mkdir(parents=True, exist_ok=True) + f.write_bytes(b'PNG') + with PartWriter(run_dir) as w: + for i in range(3): + w.write_row( + project_id='acme/vision', + run_id='galrun', + attribute_path='gallery', + attribute_type='media', + step=5, + timestamp_ms=T0_MS, + string_value='image-file', + file_value=f'files/idx{i}.png', + caption=f'c{i}', + ) + mark_run_exported(run_dir, {'rows': 3}) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + gallery_calls = _log_calls_with(op, lambda d: 'gallery' in d) + assert len(gallery_calls) == 1 # ONE batched call, not three + value = gallery_calls[0].args[0]['gallery'] + assert isinstance(value, list) and len(value) == 3 + # order preserved -> sampleIndex 0,1,2 assigned by op.log's enumerate + assert [img._caption for img in value] == ['c0', 'c1', 'c2'] + + def test_cleanup_removes_staged_files_after_load(self, tmp_path, mock_init): + # --cleanup frees each run's staged files once it's loaded, but the run + # stays recorded as loaded (so a re-run still skips it). + run_dir = _stage_run(tmp_path) + assert run_dir.exists() + summary = PlutoLoader(tmp_path, cleanup=True).load() + assert summary['loaded'] == 1 + assert not run_dir.exists() # staged files reclaimed + assert LoadedCache(tmp_path / 'loaded_runs.json').is_loaded(CACHE_KEY) + + def test_no_cleanup_keeps_staged_files(self, tmp_path, mock_init): + run_dir = _stage_run(tmp_path) + PlutoLoader(tmp_path).load() # cleanup defaults off + assert run_dir.exists() + + def test_run_metadata_forwarded_as_system_metadata(self, tmp_path, mock_init): + # run.metadata (git/OS/GPU) is staged; the loader forwards it via + # compat['systemMetadata'] so repro context survives the migration. + init, _ = mock_init + run_dir = _stage_run(tmp_path) + manifest = json.loads((run_dir / 'run.json').read_text()) + manifest['metadata'] = {'gpu': 'H100', 'python': '3.12'} + write_json_atomic(run_dir / 'run.json', manifest) + PlutoLoader(tmp_path).load() + compat = init.call_args.kwargs['settings']['compat'] + assert compat['systemMetadata'] == {'gpu': 'H100', 'python': '3.12'} + + def test_no_metadata_keeps_compat_clean(self, tmp_path, mock_init): + # Runs without metadata must not get a systemMetadata key (normal runs + # send empty compat; only migration populates it). + init, _ = mock_init + _stage_run(tmp_path) # no 'metadata' in the staged run.json + PlutoLoader(tmp_path).load() + assert 'systemMetadata' not in init.call_args.kwargs['settings']['compat'] + + def test_cache_key_includes_dest_project(self, tmp_path, mock_init): + # Loading the same export into a different dest project must NOT be + # skipped just because it was loaded into another project. + _stage_run(tmp_path) + PlutoLoader(tmp_path, dest_project='proj-a').load() + cache = LoadedCache(tmp_path / 'loaded_runs.json') + assert cache.is_loaded(f'{EXTERNAL_ID}@@proj-a') + assert not cache.is_loaded(f'{EXTERNAL_ID}@@proj-b') # different dest + # a second load into proj-b actually runs (not skipped) + s = PlutoLoader(tmp_path, dest_project='proj-b').load() + assert s['loaded'] == 1 and s['skipped'] == 0 + + def test_bad_media_row_does_not_fail_whole_run(self, tmp_path, mock_init): + # A single unparseable inline-media payload is skipped; the run's other + # data still loads and the run is not marked failed. + _, op = mock_init + run_dir = tmp_path / 'acme' / 'vision' / 'runs' / 'mixrun' + run_dir.mkdir(parents=True) + write_json_atomic( + run_dir / 'run.json', + { + 'entity': 'acme', + 'project': 'vision', + 'run_id': 'mixrun', + 'name': 'm', + 'state': 'finished', + }, + ) + with PartWriter(run_dir) as w: + w.write_row( + project_id='acme/vision', + run_id='mixrun', + attribute_path='broken', + attribute_type='media', + step=0, + timestamp_ms=T0_MS, + string_value='{not valid json', # json.loads raises + ) + w.write_row( + project_id='acme/vision', + run_id='mixrun', + attribute_path='loss', + attribute_type='metric', + step=1, + timestamp_ms=T0_MS + 1000, + float_value=0.5, + ) + mark_run_exported(run_dir, {'rows': 2}) + summary = PlutoLoader(tmp_path).load() + assert summary == {'loaded': 1, 'skipped': 0, 'failed': []} + op._log_metrics_batch.assert_called() # the good metric still replayed diff --git a/tests/test_migrate_schema.py b/tests/test_migrate_schema.py new file mode 100644 index 0000000..ac0b63f --- /dev/null +++ b/tests/test_migrate_schema.py @@ -0,0 +1,159 @@ +""" +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', + 'string_series', + } + + +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') + + def test_loaded_cache_in_progress_vs_done(self, tmp_path): + path = tmp_path / 'loaded_runs.json' + c = LoadedCache(path) + c.mark_in_progress('wandb::acme/vision/x') + assert c.is_in_progress('wandb::acme/vision/x') + assert not c.is_loaded('wandb::acme/vision/x') # in_progress != done + # survives a reopen, then promotes to done + assert LoadedCache(path).is_in_progress('wandb::acme/vision/x') + c.mark_loaded('wandb::acme/vision/x', {'pluto_run_id': 1}) + assert c.is_loaded('wandb::acme/vision/x') + assert not c.is_in_progress('wandb::acme/vision/x') + + def test_legacy_cache_entries_treated_as_done(self, tmp_path): + # Pre-existing loaded_runs.json entries have no 'status' field. + path = tmp_path / 'loaded_runs.json' + write_json_atomic(path, {'wandb::acme/vision/old': {'pluto_run_id': 9}}) + assert LoadedCache(path).is_loaded('wandb::acme/vision/old') + + def test_corrupt_loaded_cache_does_not_raise(self, tmp_path): + # A truncated/empty cache must not abort load(); start fresh + back up. + path = tmp_path / 'loaded_runs.json' + path.write_text('{ truncated') + cache = LoadedCache(path) # must not raise + assert not cache.is_loaded('anything') + assert (tmp_path / 'loaded_runs.corrupt').exists() 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..01e70ca --- /dev/null +++ b/tests/test_migrate_wandb_export.py @@ -0,0 +1,976 @@ +""" +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 +import os +import time +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, + version='v0', + aliases=None, + ): + self.name = name + self.type = type + self.size = size + self.version = version + self.aliases = aliases if aliases is not None else ['latest'] + 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) + + +def _table_artifact_json(media=True): + """A wandb run_table artifact's .table.json, with or without a media column.""" + number = {'wb_type': 'number', 'params': {}} + img = { + 'wb_type': 'union', + 'params': { + 'allowed_types': [ + {'wb_type': 'none', 'params': {}}, + {'wb_type': 'image-file', 'params': {}}, + ] + }, + } + type_map = {'idx': number, 'img': img if media else number, 'score': number} + return { + '_type': 'table', + 'columns': ['idx', 'img', 'score'], + 'column_types': {'wb_type': 'typedDict', 'params': {'type_map': type_map}}, + 'data': [[0, 'Image' if media else 1.0, 0.0]], + } + + +class FakeTableArtifact: + """A run_table artifact whose download writes a .table.json (media or not).""" + + def __init__(self, name='run-abc-media_table:v0', media=True): + self.name = name + self.type = 'run_table' + self.size = 200 + self.version = 'v0' + self.aliases = ['latest'] + self.created_at = CREATED_AT + self._media = media + + def download(self, root): + root = Path(root) + root.mkdir(parents=True, exist_ok=True) + with open(root / 'media_table.table.json', 'w') as fh: + json.dump(_table_artifact_json(self._media), fh) + return str(root) + + +class FakeSweep: + def __init__(self, id='sweep-abc', name='my-sweep', config=None): + self.id = id + self.name = name + self.config = ( + config + if config is not None + else {'method': 'grid', 'parameters': {'lr': {'values': [0.1, 0.01]}}} + ) + + +class FakeRun: + entity = 'acme' + project = 'vision' + + def __init__( + self, + run_id='abc123', + artifacts=None, + output_log=None, + sweep=None, + used_artifacts=None, + history_rows=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.sweep = sweep + self._history_rows = history_rows + self._used_artifacts = used_artifacts if used_artifacts is not None else [] + 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 + if self._history_rows is not None: + return iter(self._history_rows) + 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) + + def used_artifacts(self): + return list(self._used_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'], summary['skipped'], summary['failed']) == ( + 1, + 0, + [], + ) + 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_no_files_skips_file_media_but_keeps_histogram(self, tmp_path): + # --no-files (include_files=False) must not stage file-backed media rows: + # the files aren't downloaded, so a file_value pointer would dangle and + # fail the loader. Histograms carry no file and stay. + _, run_dir, summary = _export(tmp_path, include_files=False) + media = {m['attribute_path']: m for m in _rows(run_dir, 'media')} + assert 'sample' not in media # image-file row dropped + assert 'weights' in media # histogram kept + assert not (run_dir / 'files/media/images/sample_3_abc.png').exists() + # Dropped media is surfaced in coverage, not silently omitted. + assert 'media-file(--no-files)' in summary['coverage']['not_migrated'] + + def test_failed_download_is_counted_for_strict(self, tmp_path): + # A media file that fails to download leaves a dangling ref; it must be + # counted so coverage/--strict surface the loss (not a false success). + class _FailingFile: + name = 'media/images/sample_3_abc.png' + size = 10 + + def download(self, root, replace=False, exist_ok=False): + raise RuntimeError('storage 503') + + run = FakeRun() + run._files = [f for f in run._files if 'sample_3' not in f.name] + run._files.append(_FailingFile()) + _, _, summary = _export(tmp_path, run=run) + assert 'file-download-failed' in summary['coverage']['not_migrated'] + + def test_failed_artifact_download_is_counted(self, tmp_path): + # An artifact whose download raises (transient storage error) must be + # flagged, not silently dropped, so coverage/--strict surface the loss. + class _FailingArtifact: + name = 'model:v0' + type = 'model' + size = 10 + version = 'v0' + aliases = ['latest'] + created_at = CREATED_AT + + def download(self, root): + raise RuntimeError('storage 503') + + run = FakeRun(artifacts=[_FailingArtifact()]) + _, _, summary = _export(tmp_path, run=run) + assert 'artifact-download-failed' in summary['coverage']['not_migrated'] + + def test_sweep_is_migrated(self, tmp_path): + # A run in a sweep migrates its sweep membership + search-space config + # into the manifest (id/name/config), counted as migrated. + run = FakeRun(sweep=FakeSweep(id='swp1')) + _, run_dir, summary = _export(tmp_path, run=run) + assert summary['coverage']['migrated'].get('sweep') == 1 + manifest = read_json(run_dir / 'run.json') + assert manifest['sweep']['id'] == 'swp1' + assert manifest['sweep']['name'] == 'my-sweep' + assert manifest['sweep']['config']['method'] == 'grid' + + def test_input_artifact_lineage_is_flagged(self, tmp_path): + # A run that consumed artifacts (used_artifacts) loses that input lineage. + run = FakeRun(used_artifacts=[FakeArtifact(name='dataset:v0')]) + _, _, summary = _export(tmp_path, run=run) + assert 'artifact-input-lineage' in summary['coverage']['not_migrated'] + + def test_artifact_versioning_is_flagged(self, tmp_path): + # A logged artifact with version history / non-'latest' aliases: only its + # files migrate, the versioning doesn't. + run = FakeRun(artifacts=[FakeArtifact(name='model:v3', version='v3')]) + _, _, summary = _export(tmp_path, run=run) + assert 'artifact-versioning' in summary['coverage']['not_migrated'] + + def test_plain_run_has_no_false_lineage_flags(self, tmp_path): + # No sweep, no used artifacts, a v0/latest artifact, and wandb's internal + # history artifact (v-bumped) must NOT trip any lineage/sweep flag. + run = FakeRun( + artifacts=[ + FakeArtifact(name='model:v0'), + FakeArtifact( + name='run-abc-history:v9', type='wandb-history', version='v9' + ), + ] + ) + _, _, summary = _export(tmp_path, run=run) + nm = summary['coverage']['not_migrated'] + assert 'sweep-metadata' not in nm + assert 'artifact-input-lineage' not in nm + assert 'artifact-versioning' not in nm + + def test_media_table_is_flagged(self, tmp_path): + # A logged table with an image column migrates degraded (cells become + # text placeholders); flag it as table-media-cell so the gap is visible + # and trips --strict instead of migrating silently. + run = FakeRun(artifacts=[FakeTableArtifact(media=True)]) + _, _, summary = _export(tmp_path, run=run) + assert summary['coverage']['not_migrated'].get('table-media-cell') == 1 + + def test_plain_table_not_flagged(self, tmp_path): + # A table with only scalar columns is fully migrated — no media flag. + run = FakeRun(artifacts=[FakeTableArtifact(media=False)]) + _, _, summary = _export(tmp_path, run=run) + assert 'table-media-cell' not in summary['coverage']['not_migrated'] + + def test_unknown_media_types_flagged(self, tmp_path): + # Any media _type the exporter doesn't handle (bokeh, molecule, joined/ + # partitioned tables, ...) is dropped with an unsupported() flag + # naming exactly what was lost, and trips --strict. + run = FakeRun( + history_rows=[ + {'_step': 0, '_timestamp': T0, 'viz': {'_type': 'bokeh-file'}}, + {'_step': 1, '_timestamp': T0 + 1, 'jt': {'_type': 'joined-table'}}, + ] + ) + _, _, summary = _export(tmp_path, run=run) + nm = summary['coverage']['not_migrated'] + assert nm.get('unsupported(bokeh-file)') == 1 + assert nm.get('unsupported(joined-table)') == 1 + + def test_string_series_too_long_dropped(self, tmp_path): + # A per-step string over the cap is a stray blob, not a real state: + # drop it with string-series-too-long while normal labels still migrate. + run = FakeRun( + history_rows=[ + {'_step': 0, '_timestamp': T0, 'phase': 'train'}, + {'_step': 1, '_timestamp': T0 + 1, 'blob': 'x' * 250}, + ] + ) + _, run_dir, summary = _export(tmp_path, run=run) + assert summary['coverage']['not_migrated'].get('string-series-too-long') == 1 + series = _rows(run_dir, 'string_series') + assert any(r['attribute_path'] == 'phase' for r in series) + assert all(r['attribute_path'] != 'blob' for r in series) + + 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'], summary['skipped'], summary['failed']) == ( + 0, + 1, + [], + ) + 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 / 'acme' / 'vision' / 'manifest.json') + assert manifest['failed'][0]['run_id'] == 'boom' + + def test_export_retries_after_transient_history_failure(self, tmp_path): + # A crash-truncated cache read (or a network blip) fails the first + # attempt; the exporter purges + retries once and the run still exports. + run = FakeRun() + orig = run.scan_history + calls = {'n': 0} + + def flaky(page_size=1000): + calls['n'] += 1 + if calls['n'] == 1: + raise RuntimeError('Parquet file too small. Size is 0 but need 8') + return orig(page_size=page_size) + + run.scan_history = flaky + _, run_dir, summary = _export(tmp_path, run=run) + assert summary['exported'] == 1 and summary['failed'] == [] + assert calls['n'] == 2 # failed once, retried, succeeded + assert is_run_exported(run_dir) + + def test_purge_empty_wandb_cache_removes_only_zero_byte_parquets( + self, tmp_path, monkeypatch + ): + cache = tmp_path / 'wandbcache' / 'runhistory' + cache.mkdir(parents=True) + good = cache / 'good.parquet' + good.write_bytes(b'PAR1realdata') + empty = cache / 'empty.parquet' # crash-truncated, and stale + empty.write_bytes(b'') + old = time.time() - 3600 + os.utime(empty, (old, old)) # aged past the stale threshold + # A fresh 0-byte file = another worker's in-progress download; the age + # gate must spare it so parallel projects don't delete each other's data. + fresh = cache / 'fresh.parquet' + fresh.write_bytes(b'') + monkeypatch.setenv('WANDB_CACHE_DIR', str(tmp_path / 'wandbcache')) + exporter = WandbExporter( + entity='acme', project='vision', output_dir=tmp_path, api=FakeApi([]) + ) + assert exporter._purge_empty_wandb_cache() == 1 # only the stale empty one + assert good.exists() and not empty.exists() + assert fresh.exists() # active download spared + + 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() + + @pytest.mark.parametrize('flag', ['after', 'before']) + def test_unparseable_date_filter_raises(self, tmp_path, flag): + # A typo like 2024/01/01 must fail loudly, not silently drop the filter + # and export everything. + with pytest.raises(ValueError, match='could not parse date'): + WandbExporter( + entity='acme', + project='vision', + output_dir=tmp_path, + api=FakeApi([]), + **{flag: '2024/01/01'}, + ) + + def test_missing_date_filter_is_none(self, tmp_path): + # Not supplying the flag is fine (no filter), distinct from a bad value. + exp = WandbExporter( + entity='acme', project='vision', output_dir=tmp_path, api=FakeApi([]) + ) + assert exp.after_ms is None and exp.before_ms is None + + def test_undated_run_excluded_when_date_filter_set(self, tmp_path): + # A run with no parseable created_at must NOT silently slip through a + # requested date window — it's excluded, not exported. + run = FakeRun() + run.created_at = 'not-a-date' # parse_iso_ms -> None + _, run_dir, summary = _export(tmp_path, run=run, after='2020-01-01') + assert summary['exported'] == 0 + assert not run_dir.exists() + + def test_undated_run_kept_when_no_date_filter(self, tmp_path): + # Without any date filter, an undated run still exports (unchanged). + run = FakeRun() + run.created_at = 'not-a-date' + _, run_dir, summary = _export(tmp_path, run=run) + assert summary['exported'] == 1 + assert run_dir.exists() + + def test_boolean_metric_recorded_as_float(self, tmp_path): + run = FakeRun() + run.scan_history = lambda page_size=1000: iter( + [{'_step': 0, '_timestamp': T0, 'is_best': True, 'was_worse': False}] + ) + _, run_dir, _ = _export(tmp_path, run=run) + metrics = { + r['attribute_path']: r['float_value'] for r in _rows(run_dir, 'metric') + } + assert metrics['is_best'] == 1.0 + assert metrics['was_worse'] == 0.0 + + def test_console_falls_back_to_heartbeat_when_created_at_missing(self, tmp_path): + run = FakeRun(output_log=b'no timestamp here\n') + run.created_at = 'not-a-date' # unparseable -> None + # heartbeat_at stays valid; console lines should use it, not epoch 0. + _, run_dir, _ = _export(tmp_path, run=run) + console = _rows(run_dir, 'console') + heartbeat_ms = 1746100800000 # 2025-05-01T12:00:00Z + assert console and all(r['timestamp_ms'] == heartbeat_ms for r in console) + + def test_coverage_reports_migrated_and_dropped(self, tmp_path): + run = FakeRun() + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'loss': 1.0, # migrated metric + 'status': 'running', # string -> migrated string_series + 'chart': {'_type': 'bokeh-file'}, # unsupported -> not migrated + 'img': { # image with bounding boxes (migrated) + a mask + '_type': 'image-file', + 'path': 'media/images/sample_3_abc.png', + 'boxes': {'predictions': {'path': 'x.boxes2D.json'}}, + 'masks': {'predictions': {'path': 'x.mask.png'}}, + }, + } + ] + ) + _, _, summary = _export(tmp_path, run=run) + cov = summary['coverage'] + assert cov['migrated'].get('metric') == 1 + assert cov['migrated'].get('media') == 1 + assert cov['migrated'].get('string_series') == 1 + assert cov['migrated'].get('image-boxes') == 1 # boxes migrate + assert cov['migrated'].get('image-masks') == 1 # masks migrate too + assert cov['not_migrated'].get('unsupported(bokeh-file)') == 1 + + def test_annotated_image_stages_box_refs(self, tmp_path): + # An image with boxes stages the wandb box refs in annotation_value; the + # loader later resolves them into the image's annotations. + run = FakeRun() + boxes = { + 'pred': { + 'path': 'media/metadata/boxes2D/a.boxes2D.json', + '_type': 'boxes2D', + } + } + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'annotated': { + '_type': 'image-file', + 'path': 'media/images/sample_3_abc.png', + 'boxes': boxes, + }, + } + ] + ) + _, run_dir, _ = _export(tmp_path, run=run) + media = next( + r for r in _rows(run_dir, 'media') if r['attribute_path'] == 'annotated' + ) + assert json.loads(media['annotation_value']) == {'boxes': boxes} + + def test_mask_class_labels_folded_from_config(self, tmp_path): + # wandb keeps a mask's class_labels in the run config (not the mask + # descriptor). The exporter must fold them into the mask annotation, + # keyed '_wandb_delimeter_', or the mask renders blank. + run = FakeRun() + run.json_config = json.dumps( + { + '_wandb': { + 'value': { + 'mask/class_labels': { + 'seg_wandb_delimeter_predictions': { + 'value': {'0': 'bg', '1': 'cat'} + } + } + } + } + } + ) + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'seg': { + '_type': 'image-file', + 'path': 'media/images/seg.png', + 'masks': { + 'predictions': { + '_type': 'mask', + 'path': 'media/images/mask/seg.mask.png', + } + }, + }, + } + ] + ) + _, run_dir, _ = _export(tmp_path, run=run) + media = next(r for r in _rows(run_dir, 'media') if r['attribute_path'] == 'seg') + ann = json.loads(media['annotation_value']) + assert ann['masks']['predictions']['class_labels'] == {'0': 'bg', '1': 'cat'} + + def test_gallery_images_keep_per_image_boxes_and_masks(self, tmp_path): + # Images logged as a LIST (wandb type images/separated) carry their + # boxes/masks in parallel all_boxes[i]/all_masks[i] arrays. Each staged + # image row must get its own annotation_value — else galleries migrate + # as plain pictures with the annotations silently dropped. + run = FakeRun() + b0 = {'pred': {'path': 'meta/g0.boxes2D.json', '_type': 'boxes2D'}} + b1 = {'pred': {'path': 'meta/g1.boxes2D.json', '_type': 'boxes2D'}} + m0 = {'pred': {'path': 'mask/g0.mask.png', '_type': 'mask'}} + m1 = {'pred': {'path': 'mask/g1.mask.png', '_type': 'mask'}} + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'val/annotated': { + '_type': 'images/separated', + 'filenames': [ + 'media/images/val/annotated_0.png', + 'media/images/val/annotated_1.png', + ], + 'captions': ['idx0', 'idx1'], + 'all_boxes': [b0, b1], + 'all_masks': [m0, m1], + }, + } + ] + ) + _, run_dir, _ = _export(tmp_path, run=run) + rows = [ + r for r in _rows(run_dir, 'media') if r['attribute_path'] == 'val/annotated' + ] + assert len(rows) == 2 + rows.sort(key=lambda r: r['file_value']) + ann0 = json.loads(rows[0]['annotation_value']) + ann1 = json.loads(rows[1]['annotation_value']) + assert ann0 == {'boxes': b0, 'masks': m0} + assert ann1 == {'boxes': b1, 'masks': m1} + + def test_histogram_bins_reconstructed_from_packed(self, tmp_path): + # Real wandb histograms carry edges in packedBins ({min,size,count}), + # not a `bins` array. Reconstruct the true edges so the migrated + # histogram keeps its value range, not a generic 0..N axis. + run = FakeRun() + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'weights': { + '_type': 'histogram', + 'values': [5, 10, 25], + 'bins': None, + 'packedBins': {'min': -1.5, 'size': 0.5, 'count': 3}, + }, + } + ] + ) + _, run_dir, _ = _export(tmp_path, run=run) + hist = next( + m for m in _rows(run_dir, 'media') if m['attribute_path'] == 'weights' + ) + payload = json.loads(hist['string_value']) + # 3 counts -> 4 edges spanning the real range, not [0,1,2,3] + assert payload['values'] == [5, 10, 25] + assert payload['bins'] == [-1.5, -1.0, -0.5, 0.0] + + def test_nonfinite_string_metrics_are_coerced_to_floats(self, tmp_path): + # wandb hands NaN/Inf back as strings; they must migrate as real floats, + # not as a categorical series. A genuine text value ('running') migrates + # as a string_series instead. + import math + + run = FakeRun() + run.scan_history = lambda page_size=1000: iter( + [ + {'_step': 0, '_timestamp': T0, 'edge': 'NaN', 'status': 'running'}, + {'_step': 1, '_timestamp': T0 + 1, 'edge': 'Infinity'}, + {'_step': 2, '_timestamp': T0 + 2, 'edge': '-Infinity'}, + ] + ) + _, run_dir, summary = _export(tmp_path, run=run) + vals = {r['step']: r['float_value'] for r in _rows(run_dir, 'metric')} + assert math.isnan(vals[0]) + assert vals[1] == float('inf') + assert vals[2] == float('-inf') + assert summary['coverage']['migrated'].get('metric') == 3 + assert summary['coverage']['migrated'].get('string_series') == 1 # 'running' + + def test_string_history_migrated_as_string_series(self, tmp_path): + # A non-numeric, non-media string history value ('phase') is a + # categorical series: staged as string_series rows (one per step), + # preserving the raw label. An over-long value is a stray blob, dropped. + run = FakeRun() + long_blob = 'x' * 250 + run.scan_history = lambda page_size=1000: iter( + [ + {'_step': 0, '_timestamp': T0, 'phase': 'warmup', 'blob': long_blob}, + {'_step': 1, '_timestamp': T0 + 1, 'phase': 'train'}, + {'_step': 2, '_timestamp': T0 + 2, 'phase': 'done'}, + ] + ) + _, run_dir, summary = _export(tmp_path, run=run) + ss = _rows(run_dir, 'string_series') + assert [(r['step'], r['attribute_path'], r['string_value']) for r in ss] == [ + (0, 'phase', 'warmup'), + (1, 'phase', 'train'), + (2, 'phase', 'done'), + ] + cov = summary['coverage'] + assert cov['migrated'].get('string_series') == 3 + assert cov['not_migrated'].get('string-series-too-long') == 1 + + def test_custom_charts_extracted_from_config_yaml(self, tmp_path): + # wandb.plot.* panels live in the raw config.yaml under + # _wandb.value.visualize (the API strips _wandb). The exporter recovers + # each panel's preset, title, backing-table key, and column mappings. + run = FakeRun() + config_yaml = ( + '_wandb:\n' + ' value:\n' + ' visualize:\n' + ' bar:\n' + ' panel_type: Vega2\n' + ' panel_config:\n' + ' panelDefId: wandb/bar/v0\n' + ' fieldSettings: {label: label, value: value}\n' + ' stringSettings: {title: per-class}\n' + ' userQuery:\n' + ' queryFields:\n' + ' - name: runSets\n' + ' fields:\n' + ' - name: summaryTable\n' + ' args:\n' + ' - {name: tableKey, value: bar_table}\n' + ' pr:\n' + ' panel_type: Vega2\n' + ' panel_config:\n' + ' panelDefId: wandb/area-under-curve/v0\n' + ' fieldSettings: {x: recall, y: precision}\n' + ' stringSettings:\n' + ' title: PR\n' + ' x-axis-title: Recall\n' + ' y-axis-title: Precision\n' + ' userQuery:\n' + ' queryFields:\n' + ' - name: runSets\n' + ' fields:\n' + ' - name: summaryTable\n' + ' args:\n' + ' - {name: tableKey, value: pr_table}\n' + ' weird:\n' + ' panel_type: Vega2\n' + ' panel_config:\n' + ' panelDefId: wandb/custom/v0\n' + ) + run._files.append(FakeFile('config.yaml', content=config_yaml.encode())) + _, run_dir, summary = _export(tmp_path, run=run) + panels = { + p['key']: p for p in read_json(run_dir / 'custom_charts.json')['panels'] + } + bar = panels['bar'] + assert bar['preset'] == 'bar' + assert bar['tableKey'] == 'bar_table' + assert bar['title'] == 'per-class' + assert bar['fields'] == {'label': 'label', 'value': 'value'} + assert bar['specLang'] == 'vega-lite' + # The whole stringSettings dict is forwarded (axis titles included), not + # just the title — the renderer substitutes ${string:x-axis-title} etc. + pr = panels['pr'] + assert pr['preset'] == 'area-under-curve' # newly-mapped preset id + assert pr['strings'] == { + 'title': 'PR', + 'x-axis-title': 'Recall', + 'y-axis-title': 'Precision', + } + assert bar['strings'] == {'title': 'per-class'} + # Unknown preset: staged for reference but flagged, and marked raw Vega. + assert panels['weird']['preset'] is None + assert panels['weird']['specLang'] == 'vega' + assert panels['weird']['strings'] == {} + cov = summary['coverage'] + assert cov['migrated'].get('custom-chart') == 2 # bar + pr + assert cov['not_migrated'].get('custom-chart-unsupported') == 1 + + def test_media_lists_videos_audio_are_migrated(self, tmp_path): + # wandb.log({"rollouts": [Video, Video]}) => _type 'videos' with the + # items under a matching key, each a {path, caption, _type} dict. + # These were silently dropped before; now each item is a media row. + run = FakeRun() + run.scan_history = lambda page_size=1000: iter( + [ + { + '_step': 0, + '_timestamp': T0, + 'rollouts': { + '_type': 'videos', + 'count': 2, + 'videos': [ + { + '_type': 'video-file', + 'path': 'media/v0.gif', + 'caption': '0', + }, + { + '_type': 'video-file', + 'path': 'media/v1.gif', + 'caption': '1', + }, + ], + 'captions': ['0', '1'], + }, + 'clips': { + '_type': 'audio', + 'count': 1, + 'audio': [ + { + '_type': 'audio-file', + 'path': 'media/a0.wav', + 'caption': 'c', + }, + ], + }, + } + ] + ) + _, run_dir, summary = _export(tmp_path, run=run) + media = _rows(run_dir, 'media') + vids = [m for m in media if m['string_value'] == 'video-file'] + auds = [m for m in media if m['string_value'] == 'audio-file'] + assert [(m['file_value'], m['caption']) for m in vids] == [ + ('files/media/v0.gif', '0'), + ('files/media/v1.gif', '1'), + ] + assert [(m['file_value'], m['caption']) for m in auds] == [ + ('files/media/a0.wav', 'c'), + ] + assert summary['coverage']['migrated'].get('media') == 3 + assert 'unsupported(videos)' not in summary['coverage']['not_migrated'] + + def test_many_files_all_downloaded_concurrently(self, tmp_path): + # Media-heavy runs have hundreds of files; the exporter downloads them + # concurrently (latency-bound). All must land regardless of worker count. + run = FakeRun() + run._files = [FakeFile('output.log', content=b'x')] + [ + FakeFile(f'media/images/img_{i}.png', content=b'PNG') for i in range(40) + ] + _, run_dir, _ = _export(tmp_path, run=run, download_workers=8) + got = list((run_dir / 'files' / 'media' / 'images').glob('img_*.png')) + assert len(got) == 40 + assert (run_dir / 'files' / 'output.log').exists() + + def test_download_workers_one_still_downloads(self, tmp_path): + # Serial fallback (workers=1) must still fetch everything. + run = FakeRun() + run._files = [FakeFile(f'f_{i}.png', content=b'PNG') for i in range(5)] + _, run_dir, _ = _export(tmp_path, run=run, download_workers=1) + assert len(list((run_dir / 'files').glob('f_*.png'))) == 5 + + def test_metadata_staged_for_systemMetadata_forwarding(self, tmp_path): + _, run_dir, _ = _export(tmp_path) + # run.metadata is staged in run.json; the loader forwards it as + # systemMetadata on create (see loader test). + assert read_json(run_dir / 'run.json')['metadata'] == {'gpu': 'NVIDIA H100'} diff --git a/tests/test_run_status.py b/tests/test_run_status.py index 4f1053d..5e3d435 100644 --- a/tests/test_run_status.py +++ b/tests/test_run_status.py @@ -45,6 +45,44 @@ def test_status_map_has_terminated_sigterm(self): assert STATUS[signal.SIGTERM.value] == 'TERMINATED' +class TestStatusPayloadBackfill: + """The status-update payload carries a historical terminal-status time + for backfilled/migrated runs (pluto.migrate) so the server can keep + Duration = end - createdAt correct; normal runs send nothing and the + server keeps its now() default. + """ + + def _settings(self): + from pluto.sets import Settings + + s = Settings() + s._op_id = 42 + s._op_status = 0 # COMPLETED + return s + + def test_normal_run_sends_null_status_updated(self): + import json + + from pluto.api import make_compat_status_v1 + + s = self._settings() # compat defaults to {} + payload = json.loads(make_compat_status_v1(s).decode()) + assert 'statusUpdated' in payload + assert payload['statusUpdated'] is None + + def test_backfilled_run_sends_historical_status_updated(self): + import json + + from pluto.api import make_compat_status_v1 + + s = self._settings() + s.compat = {'createdAt': 1600000000000, 'updatedAt': 1600000003600} + payload = json.loads(make_compat_status_v1(s).decode()) + # statusUpdated mirrors the historical updatedAt (terminal time), + # not createdAt. + assert payload['statusUpdated'] == 1600000003600 + + class TestExcepthook: """Test the sys.excepthook integration for FAILED status detection.""" @@ -315,6 +353,61 @@ def test_excepthook_then_finish_sets_failed(self): assert STATUS[settings._op_status] == 'FAILED' +class TestTerminalStatusResilience: + """finish() must not silently swallow a failed terminal status update — it + records it on the op so callers (the migration loader) can tell the run was + not actually finalized, and a post-status teardown error must not re-flip a + confirmed run to FAILED. + """ + + def test_finish_records_unconfirmed_status_without_crashing(self): + from unittest.mock import MagicMock + + from pluto.iface import PlutoRequestError + from pluto.op import Op + from pluto.sets import Settings + + settings = Settings() + settings.mode = 'noop' + op = Op(config={}, settings=settings) + op.start() + + # Terminal status POST fails after its own retries (raises up to finish). + failing_iface = MagicMock() + failing_iface.update_status.side_effect = PlutoRequestError( + 'peer reset', status_code=None + ) + op._iface = failing_iface + + op.finish() # must NOT raise + + failing_iface.update_status.assert_called_once() + assert isinstance(op._status_update_error, PlutoRequestError) + # The intended terminal code stands — the failed POST is not re-reported + # as FAILED (status_confirmed guard); the run's code stays COMPLETED (0). + assert settings._op_status == 0 + + def test_finish_confirmed_status_clears_error_flag(self): + from unittest.mock import MagicMock + + from pluto.op import Op + from pluto.sets import Settings + + settings = Settings() + settings.mode = 'noop' + op = Op(config={}, settings=settings) + op.start() + + ok_iface = MagicMock() # update_status() succeeds + op._iface = ok_iface + + op.finish() + + ok_iface.update_status.assert_called_once() + assert op._status_update_error is None + assert settings._op_status == 0 + + class TestExcepthookSubprocess: """ End-to-end test: run a script that raises an unhandled exception diff --git a/tests/test_shutdown.py b/tests/test_shutdown.py index fc09c53..58451c1 100644 --- a/tests/test_shutdown.py +++ b/tests/test_shutdown.py @@ -1361,6 +1361,9 @@ def _make_sigterm_op(status=-1, op_id=123): _op_status=status, url_stop='http://server/api/runs/status/update', x_sigterm_status_timeout_seconds=5.0, + # make_compat_status_v1 reads compat['updatedAt'] (historical statusUpdated + # for backfilled runs); empty for a normal run. + compat={}, # Read when the handler builds its fresh httpx client. insecure_disable_ssl=False, http_proxy=None, diff --git a/tests/test_sweep.py b/tests/test_sweep.py new file mode 100644 index 0000000..6650b08 --- /dev/null +++ b/tests/test_sweep.py @@ -0,0 +1,309 @@ +"""Unit tests for pluto.sweep (native hyperparameter sweeps). + +Covers the pure combo generators, config validation, sweep() storage, and the +agent() loop's context handling — without touching the network (pluto.init is +not called; the agent's function only reads the active sweep context). +""" + +from __future__ import annotations + +import sys + +import pytest + +import pluto + +# The pluto.sweep *function* shadows the pluto.sweep *module* attribute, so +# `import pluto.sweep as sw` would bind the function. Grab the real module. +sw = sys.modules['pluto.sweep'] + + +class TestComboGeneration: + def test_grid_is_cartesian_product(self): + combos = list( + sw._grid_combos({'lr': {'values': [0.1, 0.01]}, 'bs': {'values': [16, 32]}}) + ) + assert combos == [ + {'lr': 0.1, 'bs': 16}, + {'lr': 0.1, 'bs': 32}, + {'lr': 0.01, 'bs': 16}, + {'lr': 0.01, 'bs': 32}, + ] + + def test_grid_constant_value(self): + combos = list( + sw._grid_combos({'opt': {'value': 'adam'}, 'lr': {'values': [1, 2]}}) + ) + assert combos == [{'opt': 'adam', 'lr': 1}, {'opt': 'adam', 'lr': 2}] + + def test_grid_rejects_range_param(self): + with pytest.raises(ValueError, match="use method='random'"): + list(sw._grid_combos({'lr': {'min': 0.0, 'max': 1.0}})) + + def test_random_respects_each_spec(self): + params = { + 'const': {'value': 'x'}, + 'choice': {'values': ['a', 'b']}, + 'frange': {'min': 0.0, 'max': 1.0}, + 'irange': {'min': 1, 'max': 3}, # both ints -> int sample + } + for _ in range(50): + c = sw._random_combo(params) + assert c['const'] == 'x' + assert c['choice'] in ('a', 'b') + assert 0.0 <= c['frange'] <= 1.0 + assert isinstance(c['irange'], int) and 1 <= c['irange'] <= 3 + + def test_random_log_uniform(self): + for _ in range(50): + v = sw._sample_param( + 'lr', {'min': 1e-4, 'max': 1e-1, 'distribution': 'log_uniform_values'} + ) + assert 1e-4 <= v <= 1e-1 + + +class TestValidation: + def test_bad_method_rejected(self): + with pytest.raises(ValueError, match='method must be one of'): + sw._validate_config( + {'method': 'nope', 'parameters': {'a': {'values': [1]}}} + ) + + def test_missing_parameters_rejected(self): + with pytest.raises(ValueError, match="non-empty 'parameters'"): + sw._validate_config({'method': 'grid'}) + + def test_bayes_requires_metric(self): + with pytest.raises(ValueError, match='bayes sweep needs metric'): + pluto.sweep({'method': 'bayes', 'parameters': {'a': {'values': [1]}}}) + + def test_bayes_with_metric_is_accepted(self): + sid = pluto.sweep( + { + 'method': 'bayes', + 'metric': {'name': 'loss', 'goal': 'minimize'}, + 'parameters': {'a': {'min': 0.0, 'max': 1.0}}, + } + ) + assert isinstance(sid, str) + + +class TestSweepStorage: + def test_sweep_returns_id_and_round_trips(self, tmp_path, monkeypatch): + monkeypatch.setenv('PLUTO_DIR', str(tmp_path)) + sid = pluto.sweep( + {'method': 'grid', 'parameters': {'a': {'values': [1, 2]}}}, + project='demo', + ) + assert isinstance(sid, str) and len(sid) == 8 + loaded = sw._load_sweep(sid) + assert loaded['method'] == 'grid' + assert loaded['_project'] == 'demo' + # survives losing the in-memory registry (reads the on-disk copy) + sw._SWEEP_REGISTRY.pop(sid) + assert sw._load_sweep(sid)['parameters'] == {'a': {'values': [1, 2]}} + + def test_load_unknown_sweep_raises(self): + with pytest.raises(ValueError, match='unknown sweep id'): + sw._load_sweep('doesnotexist') + + +class TestAgent: + def test_grid_agent_sets_context_per_combo(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + seen = [] + + def fake_fn(): + # the agent must have installed the sampled combo before calling us + assert sw._active_sweep is not None + seen.append(dict(sw._active_sweep['config'])) + + sid = pluto.sweep( + { + 'method': 'grid', + 'parameters': {'a': {'values': [1, 2]}, 'b': {'values': [9]}}, + } + ) + pluto.agent(sid, fake_fn) + assert seen == [{'a': 1, 'b': 9}, {'a': 2, 'b': 9}] + assert sw._active_sweep is None # cleared afterwards + + def test_grid_agent_count_caps_runs(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + n = [] + sid = pluto.sweep( + {'method': 'grid', 'parameters': {'a': {'values': [1, 2, 3, 4]}}} + ) + pluto.agent(sid, lambda: n.append(1), count=2) + assert len(n) == 2 + + def test_random_agent_requires_count(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + sid = pluto.sweep( + {'method': 'random', 'parameters': {'a': {'min': 0.0, 'max': 1.0}}} + ) + with pytest.raises(ValueError, match='needs count'): + pluto.agent(sid, lambda: None) + + def test_agent_clears_context_even_if_function_raises(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + + def boom(): + raise RuntimeError('training blew up') + + sid = pluto.sweep({'method': 'grid', 'parameters': {'a': {'values': [1]}}}) + pluto.agent(sid, boom) # must not propagate; context cleared + assert sw._active_sweep is None + + def test_agent_exposes_declared_spec(self, monkeypatch): + # The declared spec (method/metric/search-space) is available during the + # run so init() can stamp it onto config.sweep; cleared afterwards. + monkeypatch.setattr(pluto, 'ops', []) + captured = {} + + def fn(): + captured['declared'] = dict(sw._active_declared) + + sid = pluto.sweep( + { + 'method': 'grid', + 'metric': {'name': 'loss', 'goal': 'minimize'}, + 'parameters': {'a': {'values': [1]}}, + } + ) + pluto.agent(sid, fn) + d = captured['declared'] + assert d['id'] == sid + assert d['method'] == 'grid' + assert d['metric'] == {'name': 'loss', 'goal': 'minimize'} + assert d['parameters'] == {'a': {'values': [1]}} + assert sw._active_declared is None # cleared after the agent finishes + + +class TestResume: + def test_grid_resume_skips_completed_combos(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + import pluto.query as pq + + # two combos (lr=0.1) already COMPLETED in the sweep + monkeypatch.setattr( + pq, + 'list_runs', + lambda project, tags=None, limit=200, offset=0: ( + [ + {'displayId': 'A', 'status': 'COMPLETED'}, + {'displayId': 'B', 'status': 'COMPLETED'}, + ] + if offset == 0 + else [] + ), + ) + configs = {'A': {'lr': 0.1, 'bs': 16}, 'B': {'lr': 0.1, 'bs': 32}} + monkeypatch.setattr( + pq, 'get_run', lambda project, did: {'config': configs[did]} + ) + + seen = [] + sid = pluto.sweep( + { + 'method': 'grid', + 'parameters': { + 'lr': {'values': [0.1, 0.01]}, + 'bs': {'values': [16, 32]}, + }, + }, + project='p', + ) + pluto.agent( + sid, + lambda: seen.append( + (sw._active_sweep['config']['lr'], sw._active_sweep['config']['bs']) + ), + ) + assert seen == [(0.01, 16), (0.01, 32)] # only the not-done combos + + def test_random_resume_runs_remaining_count(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + import pluto.query as pq + + # 3 of a target 5 already done -> run 2 more + monkeypatch.setattr( + pq, + 'list_runs', + lambda project, tags=None, limit=200, offset=0: ( + [{'displayId': str(i), 'status': 'COMPLETED'} for i in range(3)] + if offset == 0 + else [] + ), + ) + n = [] + sid = pluto.sweep( + {'method': 'random', 'parameters': {'a': {'min': 0.0, 'max': 1.0}}}, + project='p', + ) + pluto.agent(sid, lambda: n.append(1), count=5) + assert len(n) == 2 + + def test_resume_query_failure_runs_everything(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + import pluto.query as pq + + def boom(*a, **k): + raise RuntimeError('backend down') + + monkeypatch.setattr(pq, 'list_runs', boom) + n = [] + sid = pluto.sweep( + {'method': 'grid', 'parameters': {'a': {'values': [1, 2, 3]}}}, project='p' + ) + pluto.agent(sid, lambda: n.append(1)) # query fails -> run all 3 + assert len(n) == 3 + + +class TestBayes: + def test_optuna_suggest_honors_spec(self): + import optuna + + trial = optuna.create_study().ask() + assert sw._optuna_suggest(trial, 'c', {'value': 5}) == 5 + assert sw._optuna_suggest(trial, 'ch', {'values': ['a', 'b']}) in ('a', 'b') + iv = sw._optuna_suggest(trial, 'i', {'min': 1, 'max': 3}) + assert isinstance(iv, int) and 1 <= iv <= 3 + fv = sw._optuna_suggest(trial, 'f', {'min': 0.0, 'max': 1.0}) + assert isinstance(fv, float) and 0.0 <= fv <= 1.0 + + def test_bayes_runs_count_and_optimizes(self, monkeypatch): + # Drive _run_bayes without the network by faking the run: objective is + # (x-0.7)^2, so optuna should home in near x=0.7. + monkeypatch.setattr(pluto, 'ops', []) + tried = [] + + def fake_run_combo(sweep_id, combo, project, function, metric_name, i, total): + x = combo['x'] + tried.append(x) + return (x - 0.7) ** 2 + + monkeypatch.setattr(sw, '_run_combo', fake_run_combo) + sid = pluto.sweep( + { + 'method': 'bayes', + 'metric': {'name': 'loss', 'goal': 'minimize'}, + 'parameters': {'x': {'min': 0.0, 'max': 1.0}}, + } + ) + pluto.agent(sid, lambda: None, count=25) + assert len(tried) == 25 + best = min(tried, key=lambda x: (x - 0.7) ** 2) + assert abs(best - 0.7) < 0.15 # optuna found the neighborhood + + def test_bayes_requires_count(self, monkeypatch): + monkeypatch.setattr(pluto, 'ops', []) + sid = pluto.sweep( + { + 'method': 'bayes', + 'metric': {'name': 'loss', 'goal': 'minimize'}, + 'parameters': {'x': {'min': 0.0, 'max': 1.0}}, + } + ) + with pytest.raises(ValueError, match='needs count'): + pluto.agent(sid, lambda: None) diff --git a/tests/test_sync_process.py b/tests/test_sync_process.py index 62014be..e92c33b 100644 --- a/tests/test_sync_process.py +++ b/tests/test_sync_process.py @@ -423,6 +423,39 @@ def test_graceful_shutdown_waits_for_pending(self): elapsed = time.time() - start assert elapsed < 60, f'Shutdown took too long: {elapsed}s' + def test_finish_terminates_sync_subprocess(self): + """finish()/stop(wait=True) must TERMINATE the sync subprocess, not just + flush its data and leave it running. + + The sync process is a persistent daemon (its loop only breaks on SIGTERM + or parent death). If stop() only flushed and returned, one subprocess per + run would survive for the life of the caller — a bulk single-invocation + migrate load of N runs would then hold N live subprocesses and exhaust + memory. See SyncProcessManager.stop() / _terminate_process(). + """ + run = pluto.init( + project=TESTING_PROJECT_NAME, + name=get_task_name(), + config={}, + sync_process_enabled=True, + ) + run.log({'metric': 1}) + + proc = run._sync_manager._process + assert proc is not None, 'expected a spawned sync subprocess' + assert proc.poll() is None, 'sync subprocess should be alive before finish' + + run.finish() + + # stop() terminates + reaps synchronously; give a short grace margin. + deadline = time.time() + 10 + while proc.poll() is None and time.time() < deadline: + time.sleep(0.1) + assert proc.poll() is not None, ( + 'sync subprocess is still running after finish() — it would ' + 'accumulate one process per run in a bulk load' + ) + def test_sync_manager_pending_count(self): """Test that pending count tracks metrics and files.""" run = pluto.init(