wandb 0.21.3__py3-none-musllinux_1_2_aarch64.whl → 0.21.4__py3-none-musllinux_1_2_aarch64.whl

wandb/apis/public/registries/registry.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
 from wandb_gql import gql
 
 import wandb
+from wandb._analytics import tracked
 from wandb.proto.wandb_internal_pb2 import ServerFeature
 from wandb.sdk.artifacts._validators import REGISTRY_PREFIX, validate_project_name
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -177,6 +178,7 @@ class Registry:
         """
         self._visibility = value
 
+    @tracked
     def collections(self, filter: Optional[Dict[str, Any]] = None) -> Collections:
         """Returns the collections belonging to the registry."""
         registry_filter = {
@@ -184,6 +186,7 @@ class Registry:
         }
         return Collections(self.client, self.organization, registry_filter, filter)
 
+    @tracked
     def versions(self, filter: Optional[Dict[str, Any]] = None) -> Versions:
         """Returns the versions belonging to the registry."""
         registry_filter = {
@@ -192,6 +195,7 @@ class Registry:
         return Versions(self.client, self.organization, registry_filter, None, filter)
 
     @classmethod
+    @tracked
     def create(
         cls,
         client: "Client",
@@ -256,6 +260,7 @@ class Registry:
             response["upsertModel"]["project"],
         )
 
+    @tracked
     def delete(self) -> None:
         """Delete the registry. This is irreversible."""
         try:
@@ -272,6 +277,7 @@ class Registry:
                 f"Failed to delete registry: {self.name!r} in organization: {self.organization!r}"
             )
 
+    @tracked
     def load(self) -> None:
         """Load the registry attributes from the backend to reflect the latest saved state."""
         load_failure_message = (
@@ -295,6 +301,7 @@ class Registry:
             raise ValueError(load_failure_message)
         self._update_attributes(self.attrs)
 
+    @tracked
     def save(self) -> None:
         """Save registry attributes to the backend."""
         if not InternalApi()._server_supports(

wandb/automations/_filters/expressions.py

@@ -9,6 +9,7 @@ from pydantic import ConfigDict, model_serializer
 from typing_extensions import Self, TypeAlias, get_args
 
 from wandb._pydantic import CompatBaseModel, model_validator
+from wandb._strutils import nameof
 
 from .operators import (
     Contains,
@@ -59,7 +60,7 @@ class FilterableField:
         return self._name
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({self._name!r})"
+        return f"{nameof(type(self))}({self._name!r})"
 
     # Methods to define filter expressions through chaining
     def matches_regex(self, pattern: str) -> FilterExpr:
@@ -145,7 +146,7 @@ class FilterExpr(CompatBaseModel, SupportsLogicalOpSyntax):
     op: Op
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({self.field!s}: {self.op!r})"
+        return f"{nameof(type(self))}({self.field!s}: {self.op!r})"
 
     def __rich_repr__(self) -> RichReprResult:
         # https://rich.readthedocs.io/en/stable/pretty.html

wandb/automations/_filters/operators.py

@@ -8,6 +8,7 @@ from pydantic import ConfigDict, Field, StrictBool, StrictFloat, StrictInt, Stri
 from typing_extensions import TypeAlias, get_args
 
 from wandb._pydantic import GQLBase
+from wandb._strutils import nameof
 
 # for type annotations
 Scalar = Union[StrictStr, StrictInt, StrictFloat, StrictBool]
@@ -62,7 +63,7 @@ class BaseOp(GQLBase, SupportsLogicalOpSyntax):
     def __repr__(self) -> str:
         # Display operand as a positional arg
         values_repr = ", ".join(map(repr, self.model_dump().values()))
-        return f"{type(self).__name__}({values_repr})"
+        return f"{nameof(type(self))}({values_repr})"
 
     def __rich_repr__(self) -> RichReprResult:
         # Display field values as positional args:

wandb/automations/_validators.py

@@ -5,13 +5,33 @@ from functools import singledispatch
 from itertools import chain
 from typing import Any, TypeVar
 
+from pydantic import BeforeValidator, Json, PlainSerializer
 from pydantic_core import PydanticUseDefault
+from typing_extensions import Annotated
+
+from wandb._pydantic import to_json
 
 from ._filters import And, FilterExpr, In, Nor, Not, NotIn, Op, Or
 
 T = TypeVar("T")
 
 
+def ensure_json(v: Any) -> Any:
+    """In case the incoming value isn't serialized JSON, reserialize it.
+
+    This lets us use `Json[...]` fields with values that are already deserialized.
+    """
+    # NOTE: Assumes that the deserialized type is not itself a string.
+    # Revisit this if we need to support deserialized types that are str/bytes.
+    return v if isinstance(v, (str, bytes)) else to_json(v)
+
+
+# Allow lenient instantiation/validation: incoming data may already be deserialized.
+SerializedToJson = Annotated[
+    Json[T], BeforeValidator(ensure_json), PlainSerializer(to_json)
+]
+
+
 class LenientStrEnum(str, Enum):
     """A string enum allowing for case-insensitive lookups by value.
 

wandb/automations/actions.py

@@ -7,7 +7,8 @@ from typing import Any, Literal, Optional, Union
 from pydantic import BeforeValidator, Field
 from typing_extensions import Annotated, Self, get_args
 
-from wandb._pydantic import GQLBase, GQLId, SerializedToJson, Typename
+from wandb._pydantic import GQLBase, GQLId, Typename
+from wandb._strutils import nameof
 
 from ._generated import (
     AlertSeverity,
@@ -21,6 +22,7 @@ from ._generated import (
 )
 from ._validators import (
     LenientStrEnum,
+    SerializedToJson,
     default_if_none,
     to_input_action,
     to_saved_action,
@@ -214,5 +216,5 @@ InputActionTypes: tuple[type, ...] = get_args(InputAction.__origin__)  # type: i
 
 __all__ = [
     "ActionType",
-    *(cls.__name__ for cls in InputActionTypes),
+    *(nameof(cls) for cls in InputActionTypes),
 ]

wandb/automations/events.py

@@ -9,18 +9,17 @@ from typing_extensions import Annotated, Self, get_args
 
 from wandb._pydantic import (
     GQLBase,
-    SerializedToJson,
-    ensure_json,
     field_validator,
     model_validator,
     pydantic_isinstance,
 )
+from wandb._strutils import nameof
 
 from ._filters import And, MongoLikeFilter, Or
 from ._filters.expressions import FilterableField
 from ._filters.run_metrics import MetricChangeFilter, MetricThresholdFilter, MetricVal
 from ._generated import FilterEventFields
-from ._validators import LenientStrEnum, simplify_op
+from ._validators import LenientStrEnum, SerializedToJson, ensure_json, simplify_op
 from .actions import InputAction, InputActionTypes, SavedActionTypes
 from .scopes import ArtifactCollectionScope, AutomationScope, ProjectScope
 
@@ -156,7 +155,7 @@ class _BaseEventInput(GQLBase):
         if isinstance(action, (InputActionTypes, SavedActionTypes)):
             return NewAutomation(event=self, action=action)
 
-        raise TypeError(f"Expected a valid action, got: {type(action).__qualname__!r}")
+        raise TypeError(f"Expected a valid action, got: {nameof(type(action))!r}")
 
     def __rshift__(self, other: InputAction) -> NewAutomation:
         """Implements `event >> action` to define an Automation with this event and action."""
@@ -277,7 +276,7 @@ OnRunMetric.model_rebuild()
 
 __all__ = [
     "EventType",
-    *(cls.__name__ for cls in InputEventTypes),
+    *(nameof(cls) for cls in InputEventTypes),
     "RunEvent",
     "ArtifactEvent",
     "MetricThresholdFilter",

wandb/cli/beta.py

@@ -6,13 +6,10 @@ These commands are experimental and may change or be removed in future versions.
 from __future__ import annotations
 
 import pathlib
-import sys
 
 import click
 
-import wandb
 from wandb.errors import WandbCoreNotAvailableError
-from wandb.sdk.wandb_sync import _sync
 from wandb.util import get_core_path
 
 
@@ -35,141 +32,62 @@ def beta():
         )
 
 
-@beta.command(
-    name="sync",
-    context_settings={"default_map": {}},
-    help="Upload a training run to W&B",
-)
-@click.pass_context
-@click.argument("wandb_dir", nargs=1, type=click.Path(exists=True))
-@click.option("--id", "run_id", help="The run you want to upload to.")
-@click.option("--project", "-p", help="The project you want to upload to.")
-@click.option("--entity", "-e", help="The entity to scope to.")
-@click.option("--skip-console", is_flag=True, default=False, help="Skip console logs")
-@click.option("--append", is_flag=True, default=False, help="Append run")
-@click.option(
-    "--include",
-    "-i",
-    help="Glob to include. Can be used multiple times.",
-    multiple=True,
-)
+@beta.command()
+@click.argument("paths", type=click.Path(exists=True), nargs=-1)
 @click.option(
-    "--exclude",
-    "-e",
-    help="Glob to exclude. Can be used multiple times.",
-    multiple=True,
+    "--skip-synced/--no-skip-synced",
+    is_flag=True,
+    default=True,
+    help="Skip runs that have already been synced with this command.",
 )
 @click.option(
-    "--mark-synced/--no-mark-synced",
+    "--dry-run",
     is_flag=True,
-    default=True,
-    help="Mark runs as synced",
+    default=False,
+    help="Print what would happen without uploading anything.",
 )
 @click.option(
-    "--skip-synced/--no-skip-synced",
+    "-v",
+    "--verbose",
     is_flag=True,
-    default=True,
-    help="Skip synced runs",
+    default=False,
+    help="Print more information.",
 )
 @click.option(
-    "--dry-run", is_flag=True, help="Perform a dry run without uploading anything."
+    "-n",
+    default=5,
+    help="Max number of runs to sync at a time.",
 )
-def sync_beta(  # noqa: C901
-    ctx,
-    wandb_dir=None,
-    run_id: str | None = None,
-    project: str | None = None,
-    entity: str | None = None,
-    skip_console: bool = False,
-    append: bool = False,
-    include: str | None = None,
-    exclude: str | None = None,
-    skip_synced: bool = True,
-    mark_synced: bool = True,
-    dry_run: bool = False,
+def sync(
+    paths: tuple[str, ...],
+    skip_synced: bool,
+    dry_run: bool,
+    verbose: bool,
+    n: int,
 ) -> None:
-    import concurrent.futures
-    from multiprocessing import cpu_count
-
-    paths = set()
-
-    # TODO: test file discovery logic
-    # include and exclude globs are evaluated relative to the provided base_path
-    if include:
-        for pattern in include:
-            matching_dirs = list(pathlib.Path(wandb_dir).glob(pattern))
-            for d in matching_dirs:
-                if not d.is_dir():
-                    continue
-                wandb_files = [p for p in d.glob("*.wandb") if p.is_file()]
-                if len(wandb_files) > 1:
-                    wandb.termwarn(
-                        f"Multiple wandb files found in directory {d}, skipping"
-                    )
-                elif len(wandb_files) == 1:
-                    paths.add(d)
-    else:
-        paths.update({p.parent for p in pathlib.Path(wandb_dir).glob("**/*.wandb")})
-
-    for pattern in exclude:
-        matching_dirs = list(pathlib.Path(wandb_dir).glob(pattern))
-        for d in matching_dirs:
-            if not d.is_dir():
-                continue
-            if d in paths:
-                paths.remove(d)
-
-    # remove paths that are already synced, if requested
-    if skip_synced:
-        synced_paths = set()
-        for path in paths:
-            wandb_synced_files = [p for p in path.glob("*.wandb.synced") if p.is_file()]
-            if len(wandb_synced_files) > 1:
-                wandb.termwarn(
-                    f"Multiple wandb.synced files found in directory {path}, skipping"
-                )
-            elif len(wandb_synced_files) == 1:
-                synced_paths.add(path)
-        paths -= synced_paths
-
-    if run_id and len(paths) > 1:
-        # TODO: handle this more gracefully
-        click.echo("id can only be set for a single run.", err=True)
-        sys.exit(1)
-
-    if not paths:
-        click.echo("No runs to sync.")
-        return
-
-    click.echo("Found runs:")
-    for path in paths:
-        click.echo(f"  {path}")
-
-    if dry_run:
-        return
-
-    wandb.setup()
-
-    # TODO: make it thread-safe in the Rust code
-    with concurrent.futures.ProcessPoolExecutor(
-        max_workers=min(len(paths), cpu_count())
-    ) as executor:
-        futures = []
-        for path in paths:
-            # we already know there is only one wandb file in the directory
-            wandb_file = [p for p in path.glob("*.wandb") if p.is_file()][0]
-            future = executor.submit(
-                _sync,
-                wandb_file,
-                run_id=run_id,
-                project=project,
-                entity=entity,
-                skip_console=skip_console,
-                append=append,
-                mark_synced=mark_synced,
-            )
-            futures.append(future)
-
-        # Wait for tasks to complete
-        for _ in concurrent.futures.as_completed(futures):
-            pass
+    """Upload .wandb files specified by PATHS.
+
+    PATHS can include .wandb files, run directories containing .wandb files,
+    and "wandb" directories containing run directories.
+
+    For example, to sync all runs in a directory:
+
+        wandb beta sync ./wandb
+
+    To sync a specific run:
+
+        wandb beta sync ./wandb/run-20250813_124246-n67z9ude
+
+    Or equivalently:
+
+        wandb beta sync ./wandb/run-20250813_124246-n67z9ude/run-n67z9ude.wandb
+    """
+    from . import beta_sync
+
+    beta_sync.sync(
+        [pathlib.Path(path) for path in paths],
+        dry_run=dry_run,
+        skip_synced=skip_synced,
+        verbose=verbose,
+        parallelism=n,
+    )

wandb/cli/beta_sync.py

@@ -0,0 +1,226 @@
+"""Implements `wandb sync` using wandb-core."""
+
+from __future__ import annotations
+
+import asyncio
+import pathlib
+import time
+from itertools import filterfalse
+from typing import Iterable, Iterator
+
+import click
+
+import wandb
+from wandb.proto.wandb_sync_pb2 import ServerSyncResponse
+from wandb.sdk import wandb_setup
+from wandb.sdk.lib import asyncio_compat
+from wandb.sdk.lib.printer import ERROR, Printer, new_printer
+from wandb.sdk.lib.progress import progress_printer
+from wandb.sdk.lib.service.service_connection import ServiceConnection
+from wandb.sdk.mailbox.mailbox_handle import MailboxHandle
+
+_MAX_LIST_LINES = 20
+_POLL_WAIT_SECONDS = 0.1
+_SLEEP = asyncio.sleep  # patched in tests
+
+
+def sync(
+    paths: list[pathlib.Path],
+    *,
+    dry_run: bool,
+    skip_synced: bool,
+    verbose: bool,
+    parallelism: int,
+) -> None:
+    """Replay one or more .wandb files.
+
+    Args:
+        paths: One or more .wandb files, run directories containing
+            .wandb files, and wandb directories containing run directories.
+        dry_run: If true, just prints what it would do and exits.
+        skip_synced: If true, skips files that have already been synced
+            as indicated by a .wandb.synced marker file in the same directory.
+        verbose: Verbose mode for printing more info.
+        parallelism: Max number of runs to sync at a time.
+    """
+    wandb_files: set[pathlib.Path] = set()
+    for path in paths:
+        for wandb_file in _find_wandb_files(path, skip_synced=skip_synced):
+            wandb_files.add(wandb_file.resolve())
+
+    if not wandb_files:
+        click.echo("No files to sync.")
+        return
+
+    if dry_run:
+        click.echo(f"Would sync {len(wandb_files)} file(s):")
+        _print_sorted_paths(wandb_files, verbose=verbose)
+        return
+
+    click.echo(f"Syncing {len(wandb_files)} file(s):")
+    _print_sorted_paths(wandb_files, verbose=verbose)
+
+    singleton = wandb_setup.singleton()
+    service = singleton.ensure_service()
+    printer = new_printer()
+    singleton.asyncer.run(
+        lambda: _do_sync(
+            wandb_files,
+            service=service,
+            settings=singleton.settings,
+            printer=printer,
+            parallelism=parallelism,
+        )
+    )
+
+
+async def _do_sync(
+    wandb_files: set[pathlib.Path],
+    *,
+    service: ServiceConnection,
+    settings: wandb.Settings,
+    printer: Printer,
+    parallelism: int,
+) -> None:
+    """Sync the specified files.
+
+    This is factored out to make the progress animation testable.
+    """
+    init_result = await service.init_sync(
+        wandb_files,
+        settings,
+    ).wait_async(timeout=5)
+
+    sync_handle = service.sync(init_result.id, parallelism=parallelism)
+
+    await _SyncStatusLoop(
+        init_result.id,
+        service,
+        printer,
+    ).wait_with_progress(sync_handle)
+
+
+class _SyncStatusLoop:
+    """Displays a sync operation's status until it completes."""
+
+    def __init__(
+        self,
+        id: str,
+        service: ServiceConnection,
+        printer: Printer,
+    ) -> None:
+        self._id = id
+        self._service = service
+        self._printer = printer
+
+        self._rate_limit_last_time: float | None = None
+        self._done = asyncio.Event()
+
+    async def wait_with_progress(
+        self,
+        handle: MailboxHandle[ServerSyncResponse],
+    ) -> None:
+        """Display status updates until the handle completes."""
+        async with asyncio_compat.open_task_group() as group:
+            group.start_soon(self._wait_then_mark_done(handle))
+            group.start_soon(self._show_progress_until_done())
+
+    async def _wait_then_mark_done(
+        self,
+        handle: MailboxHandle[ServerSyncResponse],
+    ) -> None:
+        response = await handle.wait_async(timeout=None)
+        if messages := list(response.errors):
+            self._printer.display(messages, level=ERROR)
+        self._done.set()
+
+    async def _show_progress_until_done(self) -> None:
+        """Show rate-limited status updates until _done is set."""
+        with progress_printer(self._printer, "Syncing...") as progress:
+            while not await self._rate_limit_check_done():
+                handle = self._service.sync_status(self._id)
+                response = await handle.wait_async(timeout=None)
+
+                if messages := list(response.new_errors):
+                    self._printer.display(messages, level=ERROR)
+                progress.update(response.stats)
+
+    async def _rate_limit_check_done(self) -> bool:
+        """Wait for rate limit and return whether _done is set."""
+        now = time.monotonic()
+        last_time = self._rate_limit_last_time
+        self._rate_limit_last_time = now
+
+        if last_time and (time_since_last := now - last_time) < _POLL_WAIT_SECONDS:
+            await asyncio_compat.race(
+                _SLEEP(_POLL_WAIT_SECONDS - time_since_last),
+                self._done.wait(),
+            )
+
+        return self._done.is_set()
+
+
+def _find_wandb_files(
+    path: pathlib.Path,
+    *,
+    skip_synced: bool,
+) -> Iterator[pathlib.Path]:
+    """Returns paths to the .wandb files to sync."""
+    if skip_synced:
+        yield from filterfalse(_is_synced, _expand_wandb_files(path))
+    else:
+        yield from _expand_wandb_files(path)
+
+
+def _expand_wandb_files(
+    path: pathlib.Path,
+) -> Iterator[pathlib.Path]:
+    """Iterate over .wandb files selected by the path."""
+    if path.suffix == ".wandb":
+        yield path
+        return
+
+    files_in_run_directory = path.glob("*.wandb")
+    try:
+        first_file = next(files_in_run_directory)
+    except StopIteration:
+        pass
+    else:
+        yield first_file
+        yield from files_in_run_directory
+        return
+
+    yield from path.glob("*/*.wandb")
+
+
+def _is_synced(path: pathlib.Path) -> bool:
+    """Returns whether the .wandb file is synced."""
+    return path.with_suffix(".wandb.synced").exists()
+
+
+def _print_sorted_paths(paths: Iterable[pathlib.Path], verbose: bool) -> None:
+    """Print file paths, sorting them and truncating the list if needed.
+
+    Args:
+        paths: Paths to print. Must be absolute with symlinks resolved.
+        verbose: If true, doesn't truncate paths.
+    """
+    # Prefer to print paths relative to the current working directory.
+    cwd = pathlib.Path(".").resolve()
+    formatted_paths: list[str] = []
+    for path in paths:
+        try:
+            formatted_path = str(path.relative_to(cwd))
+        except ValueError:
+            formatted_path = str(path)
+        formatted_paths.append(formatted_path)
+
+    sorted_paths = sorted(formatted_paths)
+    max_lines = len(sorted_paths) if verbose else _MAX_LIST_LINES
+
+    for i in range(min(len(sorted_paths), max_lines)):
+        click.echo(f"  {sorted_paths[i]}")
+
+    if len(sorted_paths) > max_lines:
+        remaining = len(sorted_paths) - max_lines
+        click.echo(f"  +{remaining:,d} more (pass --verbose to see all)")

wandb/integration/dspy/__init__.py

@@ -0,0 +1,5 @@
+"""W&B DSPy integration package."""
+
+from .dspy import WandbDSPyCallback
+
+__all__ = ["WandbDSPyCallback"]