feat(migrate): wandb → Pluto historical data migration

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
 ```
 

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": [

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.
 
 ---
 

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'

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)

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()
 

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()

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: <HxW class-id array>}`` (native — encode a PNG
+        with the class id in the red channel) or ``{path: <mask.png>}``
+        (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:

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 (