wandb 0.21.2__py3-none-win_amd64.whl → 0.21.4__py3-none-win_amd64.whl

wandb/sdk/artifacts/artifact_file_cache.py

@@ -87,7 +87,7 @@ class ArtifactFileCache:
     ) -> tuple[FilePathStr, bool, Opener]:
         opener = self._opener(path, size, skip_cache=skip_cache)
         hit = path.is_file() and path.stat().st_size == size
-        return FilePathStr(str(path)), hit, opener
+        return FilePathStr(path), hit, opener
 
     def cleanup(
         self,

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -11,8 +11,8 @@ from typing import TYPE_CHECKING
 from urllib.parse import urlparse
 
 from wandb.proto.wandb_deprecated import Deprecated
-from wandb.sdk.lib import filesystem
 from wandb.sdk.lib.deprecate import deprecate
+from wandb.sdk.lib.filesystem import copy_or_overwrite_changed
 from wandb.sdk.lib.hashutil import (
     B64MD5,
     ETag,
@@ -184,13 +184,11 @@ class ArtifactManifestEntry:
                 executor=executor,
                 multipart=multipart,
             )
-
-        if skip_cache:
-            return FilePathStr(dest_path)
-        else:
-            return FilePathStr(
-                str(filesystem.copy_or_overwrite_changed(cache_path, dest_path))
-            )
+        return FilePathStr(
+            dest_path
+            if skip_cache
+            else copy_or_overwrite_changed(cache_path, dest_path)
+        )
 
     def ref_target(self) -> FilePathStr | URIStr:
         """Get the reference URL that is targeted by this artifact entry.

wandb/sdk/artifacts/exceptions.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, TypeVar
 
 from wandb import errors
+from wandb._strutils import nameof
 
 if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact import Artifact
@@ -39,7 +40,7 @@ class ArtifactNotLoggedError(ArtifactStatusError):
         *_, name = fullname.split(".")
         msg = (
             f"{fullname!r} used prior to logging artifact or while in offline mode. "
-            f"Call {type(obj).wait.__qualname__}() before accessing logged artifact properties."
+            f"Call {nameof(obj.wait)}() before accessing logged artifact properties."
         )
         super().__init__(msg=msg, name=name, obj=obj)
 

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -198,7 +198,7 @@ class GCSHandler(StorageHandler):
             posix_ref = posix_path / relpath
         return ArtifactManifestEntry(
             path=posix_name,
-            ref=URIStr(f"{self._scheme}://{str(posix_ref)}"),
+            ref=URIStr(f"{self._scheme}://{posix_ref}"),
             digest=obj.etag,
             size=obj.size,
             extra={"versionID": obj.generation},

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Sequence
 from urllib.parse import parse_qsl, urlparse
 
 from wandb import util
+from wandb._strutils import ensureprefix
 from wandb.errors import CommError
 from wandb.errors.term import termlog
 from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache
@@ -328,7 +329,7 @@ class S3Handler(StorageHandler):
             return True
 
         # Enforce HTTPS otherwise
-        https_url = url if url.startswith("https://") else f"https://{url}"
+        https_url = ensureprefix(url, "https://")
         netloc = urlparse(https_url).netloc
         return bool(
             # Match for https://cwobject.com

wandb/sdk/launch/inputs/internal.py

@@ -171,6 +171,29 @@ def _replace_refs_and_allofs(schema: dict, defs: Optional[dict]) -> dict:
     return ret
 
 
+def _prepare_schema(schema: Any) -> dict:
+    """Prepare a schema for validation.
+
+    This function prepares a schema for validation by:
+    1. Converting a Pydantic model instance or class to a dict
+    2. Replacing $ref with their associated definition in defs
+    3. Removing any "allOf" lists that only have one item, "lifting" the item up
+
+    We support both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
+    or the BaseModel class itself (e.g. schema=MySchema)
+    """
+    if hasattr(schema, "model_json_schema") and callable(
+        schema.model_json_schema  # type: ignore
+    ):
+        schema = schema.model_json_schema()
+    if not isinstance(schema, dict):
+        raise LaunchError(
+            "schema must be a dict, Pydantic model instance, or Pydantic model class."
+        )
+    defs = schema.pop("$defs", None)
+    return _replace_refs_and_allofs(schema, defs)
+
+
 def _validate_schema(schema: dict) -> None:
     jsonschema = get_module(
         "jsonschema",
@@ -210,18 +233,7 @@ def handle_config_file_input(
         dest,
     )
     if schema:
-        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
-        # or the BaseModel class itself (e.g. schema=MySchema)
-        if hasattr(schema, "model_json_schema") and callable(
-            schema.model_json_schema  # type: ignore
-        ):
-            schema = schema.model_json_schema()
-        if not isinstance(schema, dict):
-            raise LaunchError(
-                "schema must be a dict, Pydantic model instance, or Pydantic model class."
-            )
-        defs = schema.pop("$defs", None)
-        schema = _replace_refs_and_allofs(schema, defs)
+        schema = _prepare_schema(schema)
         _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,
@@ -251,18 +263,7 @@ def handle_run_config_input(
     when a run is created.
     """
     if schema:
-        # This supports both an instance of a pydantic BaseModel class (e.g. schema=MySchema(...))
-        # or the BaseModel class itself (e.g. schema=MySchema)
-        if hasattr(schema, "model_json_schema") and callable(
-            schema.model_json_schema  # type: ignore
-        ):
-            schema = schema.model_json_schema()
-        if not isinstance(schema, dict):
-            raise LaunchError(
-                "schema must be a dict, Pydantic model instance, or Pydantic model class."
-            )
-        defs = schema.pop("$defs", None)
-        schema = _replace_refs_and_allofs(schema, defs)
+        schema = _prepare_schema(schema)
         _validate_schema(schema)
     arguments = JobInputArguments(
         include=include,

wandb/sdk/launch/inputs/schema.py

@@ -3,7 +3,7 @@ META_SCHEMA = {
     "properties": {
         "type": {
             "type": "string",
-            "enum": ["boolean", "integer", "number", "string", "object"],
+            "enum": ["boolean", "integer", "number", "string", "object", "array"],
         },
         "title": {"type": "string"},
         "description": {"type": "string"},
@@ -11,6 +11,11 @@ META_SCHEMA = {
         "enum": {"type": "array", "items": {"type": ["integer", "number", "string"]}},
         "properties": {"type": "object", "patternProperties": {".*": {"$ref": "#"}}},
         "allOf": {"type": "array", "items": {"$ref": "#"}},
+        # Array-specific properties
+        "items": {"$ref": "#"},
+        "uniqueItems": {"type": "boolean"},
+        "minItems": {"type": "integer", "minimum": 0},
+        "maxItems": {"type": "integer", "minimum": 0},
     },
     "allOf": [
         {
@@ -35,6 +40,31 @@ META_SCHEMA = {
                 }
             },
         },
+        {
+            "if": {"properties": {"type": {"const": "array"}}},
+            "then": {
+                "required": ["items"],
+                "properties": {
+                    "items": {
+                        "properties": {
+                            "type": {"enum": ["integer", "number", "string"]},
+                            "enum": {
+                                "type": "array",
+                                "items": {"type": ["integer", "number", "string"]},
+                            },
+                            "title": {"type": "string"},
+                            "description": {"type": "string"},
+                            "format": {"type": "string"},
+                        },
+                        "required": ["type", "enum"],
+                        "unevaluatedProperties": False,
+                    },
+                    "uniqueItems": {"type": "boolean"},
+                    "minItems": {"type": "integer", "minimum": 0},
+                    "maxItems": {"type": "integer", "minimum": 0},
+                },
+            },
+        },
     ],
     "unevaluatedProperties": False,
 }

wandb/sdk/lib/asyncio_compat.py

@@ -7,7 +7,7 @@ import concurrent
 import concurrent.futures
 import contextlib
 import threading
-from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, TypeVar
+from typing import Any, AsyncIterator, Callable, Coroutine, TypeVar
 
 _T = TypeVar("_T")
 
@@ -143,34 +143,71 @@ class TaskGroup:
         """
         self._tasks.append(asyncio.create_task(coro))
 
-    async def _wait_all(self) -> None:
-        """Block until all tasks complete.
+    async def _wait_all(self, *, race: bool, timeout: float | None) -> None:
+        """Block until tasks complete.
+
+        Args:
+            race: If true, blocks until the first task completes and then
+                cancels the rest. Otherwise, waits for all tasks or until
+                the first exception.
+            timeout: How long to wait.
 
         Raises:
+            TimeoutError: If the timeout expires.
             Exception: If one or more tasks raises an exception, one of these
                 is raised arbitrarily.
         """
-        done, _ = await asyncio.wait(
+        if not self._tasks:
+            return
+
+        if race:
+            return_when = asyncio.FIRST_COMPLETED
+        else:
+            return_when = asyncio.FIRST_EXCEPTION
+
+        done, pending = await asyncio.wait(
             self._tasks,
-            # NOTE: Cancelling a task counts as a normal exit,
-            #   not an exception.
-            return_when=concurrent.futures.FIRST_EXCEPTION,
+            timeout=timeout,
+            return_when=return_when,
         )
 
+        if not done:
+            raise TimeoutError(f"Timed out after {timeout} seconds.")
+
+        # If any of the finished tasks raised an exception, pick the first one.
         for task in done:
-            with contextlib.suppress(asyncio.CancelledError):
-                if exc := task.exception():
-                    raise exc
+            if exc := task.exception():
+                raise exc
 
-    def _cancel_all(self) -> None:
-        """Cancel all tasks."""
+        # Wait for remaining tasks to clean up, then re-raise any exceptions
+        # that arise. Note that pending is only non-empty when race=True.
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        for task in pending:
+            if task.cancelled():
+                continue
+            if exc := task.exception():
+                raise exc
+
+    async def _cancel_all(self) -> None:
+        """Cancel all tasks.
+
+        Blocks until cancelled tasks complete to allow them to clean up.
+        Ignores exceptions.
+        """
         for task in self._tasks:
             # NOTE: It is safe to cancel tasks that have already completed.
             task.cancel()
+        await asyncio.gather(*self._tasks, return_exceptions=True)
 
 
 @contextlib.asynccontextmanager
-async def open_task_group() -> AsyncIterator[TaskGroup]:
+async def open_task_group(
+    *,
+    exit_timeout: float | None = None,
+    race: bool = False,
+) -> AsyncIterator[TaskGroup]:
     """Create a task group.
 
     `asyncio` gained task groups in Python 3.11.
@@ -184,30 +221,58 @@ async def open_task_group() -> AsyncIterator[TaskGroup]:
     NOTE: Subtask exceptions do not propagate until the context manager exits.
     This means that the task group cannot cancel code running inside the
     `async with` block .
+
+    Args:
+        exit_timeout: An optional timeout in seconds. When exiting the
+            context manager, if tasks don't complete in this time,
+            they are cancelled and a TimeoutError is raised.
+        race: If true, all pending tasks are cancelled once any task
+            in the group completes. Prefer to use the race() function instead.
+
+    Raises:
+        TimeoutError: if exit_timeout is specified and tasks don't finish
+            in time.
     """
     task_group = TaskGroup()
 
     try:
         yield task_group
-        await task_group._wait_all()
+        await task_group._wait_all(race=race, timeout=exit_timeout)
     finally:
-        task_group._cancel_all()
+        await task_group._cancel_all()
 
 
-@contextlib.contextmanager
-def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> Iterator[None]:
+@contextlib.asynccontextmanager
+async def cancel_on_exit(coro: Coroutine[Any, Any, Any]) -> AsyncIterator[None]:
     """Schedule a task, cancelling it when exiting the context manager.
 
     If the context manager exits successfully but the given coroutine raises
     an exception, that exception is reraised. The exception is suppressed
     if the context manager raises an exception.
     """
-    task = asyncio.create_task(coro)
 
-    try:
+    async def stop_immediately():
+        pass
+
+    async with open_task_group(race=True) as group:
+        group.start_soon(stop_immediately())
+        group.start_soon(coro)
         yield
 
-        if task.done() and (exception := task.exception()):
-            raise exception
-    finally:
-        task.cancel()
+
+async def race(*coros: Coroutine[Any, Any, Any]) -> None:
+    """Wait until the first completed task.
+
+    After any coroutine completes, all others are cancelled.
+    If the current task is cancelled, all coroutines are cancelled too.
+
+    If coroutines complete simultaneously and any one of them raises
+    an exception, an arbitrary one is propagated. Similarly, if any coroutines
+    raise exceptions during cancellation, one of them propagates.
+
+    Args:
+        coros: Coroutines to race.
+    """
+    async with open_task_group(race=True) as tg:
+        for coro in coros:
+            tg.start_soon(coro)

wandb/sdk/lib/gql_request.py

@@ -4,7 +4,9 @@ Note: This was originally wandb/vendor/gql-0.2.0/wandb_gql/transport/requests.py
 The only substantial change is to reuse a requests.Session object.
 """
 
-from typing import Any, Callable, Dict, Optional, Tuple, Union
+from __future__ import annotations
+
+from typing import Any, Callable
 
 import requests
 from wandb_gql.transport.http import HTTPTransport
@@ -12,15 +14,17 @@ from wandb_graphql.execution import ExecutionResult
 from wandb_graphql.language import ast
 from wandb_graphql.language.printer import print_ast
 
+from wandb._analytics import tracked_func
+
 
 class GraphQLSession(HTTPTransport):
     def __init__(
         self,
         url: str,
-        auth: Optional[Union[Tuple[str, str], Callable]] = None,
+        auth: tuple[str, str] | Callable | None = None,
         use_json: bool = False,
-        timeout: Optional[Union[int, float]] = None,
-        proxies: Optional[Dict[str, str]] = None,
+        timeout: int | float | None = None,
+        proxies: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         """Setup a session for sending GraphQL queries and mutations.
@@ -42,15 +46,22 @@ class GraphQLSession(HTTPTransport):
     def execute(
         self,
         document: ast.Node,
-        variable_values: Optional[Dict] = None,
-        timeout: Optional[Union[int, float]] = None,
+        variable_values: dict[str, Any] | None = None,
+        timeout: int | float | None = None,
     ) -> ExecutionResult:
         query_str = print_ast(document)
         payload = {"query": query_str, "variables": variable_values or {}}
 
         data_key = "json" if self.use_json else "data"
+
+        headers = self.headers.copy() if self.headers else {}
+
+        # If we're tracking a calling python function, include it in the headers
+        if func_info := tracked_func():
+            headers.update(func_info.to_headers())
+
         post_args = {
-            "headers": self.headers,
+            "headers": headers or None,
             "cookies": self.cookies,
             "timeout": timeout or self.default_timeout,
             data_key: payload,

wandb/sdk/lib/paths.py

@@ -1,18 +1,20 @@
+from __future__ import annotations
+
 import os
 import platform
 from functools import wraps
 from pathlib import PurePath, PurePosixPath
-from typing import Any, NewType, Union
+from typing import Any, Union
+
+from typing_extensions import TypeAlias
 
 # Path _inputs_ should generally accept any kind of path. This is named the same and
 # modeled after the hint defined in the Python standard library's `typeshed`:
 # https://github.com/python/typeshed/blob/0b1cd5989669544866213807afa833a88f649ee7/stdlib/_typeshed/__init__.pyi#L56-L65
-StrPath = Union[str, "os.PathLike[str]"]
-
-# A native path to a file on a local filesystem.
-FilePathStr = NewType("FilePathStr", str)
+StrPath: TypeAlias = Union[str, "os.PathLike[str]"]
 
-URIStr = NewType("URIStr", str)
+FilePathStr: TypeAlias = str  #: A native path to a file on a local filesystem.
+URIStr: TypeAlias = str
 
 
 class LogicalPath(str):
@@ -54,7 +56,7 @@ class LogicalPath(str):
     # will result in different outputs on different platforms; however, it doesn't alter
     # absolute paths or check for prohibited characters etc.
 
-    def __new__(cls, path: StrPath) -> "LogicalPath":
+    def __new__(cls, path: StrPath) -> LogicalPath:
         if isinstance(path, LogicalPath):
             return super().__new__(cls, path)
         if hasattr(path, "as_posix"):
@@ -77,30 +79,30 @@ class LogicalPath(str):
         """Convert this path to a PurePosixPath."""
         return PurePosixPath(self)
 
-    def __getattr__(self, attr: str) -> Any:
+    def __getattr__(self, name: str) -> Any:
         """Act like a subclass of PurePosixPath for all methods not defined on str."""
         try:
-            result = getattr(self.to_path(), attr)
-        except AttributeError as e:
-            raise AttributeError(f"LogicalPath has no attribute {attr!r}") from e
+            attr = getattr(self.to_path(), name)
+        except AttributeError:
+            classname = type(self).__qualname__
+            raise AttributeError(f"{classname!r} has no attribute {name!r}") from None
 
-        if isinstance(result, PurePosixPath):
-            return LogicalPath(result)
+        if isinstance(attr, PurePosixPath):
+            return LogicalPath(attr)
 
         # If the result is a callable (a method), wrap it so that it has the same
         # behavior: if the call result returns a PurePosixPath, return a LogicalPath.
-        if callable(result):
+        if callable(fn := attr):
 
-            @wraps(result)
+            @wraps(fn)
             def wrapper(*args: Any, **kwargs: Any) -> Any:
-                inner_result = result(*args, **kwargs)
-                if isinstance(inner_result, PurePosixPath):
-                    return LogicalPath(inner_result)
-                return inner_result
+                if isinstance(res := fn(*args, **kwargs), PurePosixPath):
+                    return LogicalPath(res)
+                return res
 
             return wrapper
-        return result
+        return attr
 
-    def __truediv__(self, other: StrPath) -> "LogicalPath":
+    def __truediv__(self, other: StrPath) -> LogicalPath:
         """Act like a PurePosixPath for the / operator, but return a LogicalPath."""
         return LogicalPath(self.to_path() / LogicalPath(other))

wandb/sdk/lib/printer.py

@@ -156,14 +156,10 @@ class Printer(abc.ABC):
         """
 
     @abc.abstractmethod
-    def progress_close(self, text: str | None = None) -> None:
+    def progress_close(self) -> None:
         """Close the progress indicator.
 
         After this, `progress_update` should not be used.
-
-        Args:
-            text: The final text to set on the progress indicator.
-                Ignored in Jupyter notebooks.
         """
 
     @staticmethod
@@ -342,13 +338,10 @@ class _PrinterTerm(Printer):
         wandb.termlog(f"{next(self._progress)} {text}", newline=False)
 
     @override
-    def progress_close(self, text: str | None = None) -> None:
+    def progress_close(self) -> None:
         if self._settings and self._settings.silent:
             return
 
-        text = text or " " * 79
-        wandb.termlog(text)
-
     @property
     @override
     def supports_html(self) -> bool:
@@ -462,11 +455,14 @@ class _PrinterJupyter(Printer):
             display_id=True,
         )
 
-        if handle:
+        if not handle:
+            yield None
+            return
+
+        try:
             yield _DynamicJupyterText(handle)
+        finally:
             handle.update(self._ipython_display.HTML(""))
-        else:
-            yield None
 
     @override
     def display(
@@ -539,7 +535,7 @@ class _PrinterJupyter(Printer):
         self._progress.update(percent_done, text)
 
     @override
-    def progress_close(self, text: str | None = None) -> None:
+    def progress_close(self) -> None:
         if self._progress:
             self._progress.close()
 

wandb/sdk/lib/progress.py

@@ -70,12 +70,14 @@ def progress_printer(
         default_text: The text to show if no information is available.
     """
     with printer.dynamic_text() as text_area:
-        yield ProgressPrinter(
-            printer,
-            text_area,
-            default_text=default_text,
-        )
-        printer.progress_close()
+        try:
+            yield ProgressPrinter(
+                printer,
+                text_area,
+                default_text=default_text,
+            )
+        finally:
+            printer.progress_close()
 
 
 class ProgressPrinter:

wandb/sdk/lib/service/service_connection.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import atexit
+import pathlib
 from typing import Callable
 
 from wandb.proto import wandb_server_pb2 as spb
-from wandb.proto import wandb_settings_pb2
+from wandb.proto import wandb_settings_pb2, wandb_sync_pb2
 from wandb.sdk import wandb_settings
 from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.interface.interface_sock import InterfaceSock
@@ -12,6 +13,7 @@ from wandb.sdk.lib import asyncio_manager
 from wandb.sdk.lib.exit_hooks import ExitHooks
 from wandb.sdk.lib.service.service_client import ServiceClient
 from wandb.sdk.mailbox import HandleAbandonedError, MailboxClosedError
+from wandb.sdk.mailbox.mailbox_handle import MailboxHandle
 
 from . import service_process, service_token
 
@@ -96,6 +98,45 @@ class ServiceConnection:
         """Returns an interface for communicating with the service."""
         return InterfaceSock(self._client, stream_id=stream_id)
 
+    def init_sync(
+        self,
+        paths: set[pathlib.Path],
+        settings: wandb_settings.Settings,
+    ) -> MailboxHandle[wandb_sync_pb2.ServerInitSyncResponse]:
+        """Send a ServerInitSyncRequest."""
+        init_sync = wandb_sync_pb2.ServerInitSyncRequest(
+            path=(str(path) for path in paths),
+            settings=settings.to_proto(),
+        )
+        request = spb.ServerRequest(init_sync=init_sync)
+
+        handle = self._client.deliver(request)
+        return handle.map(lambda r: r.init_sync_response)
+
+    def sync(
+        self,
+        id: str,
+        *,
+        parallelism: int,
+    ) -> MailboxHandle[wandb_sync_pb2.ServerSyncResponse]:
+        """Send a ServerSyncRequest."""
+        sync = wandb_sync_pb2.ServerSyncRequest(id=id, parallelism=parallelism)
+        request = spb.ServerRequest(sync=sync)
+
+        handle = self._client.deliver(request)
+        return handle.map(lambda r: r.sync_response)
+
+    def sync_status(
+        self,
+        id: str,
+    ) -> MailboxHandle[wandb_sync_pb2.ServerSyncStatusResponse]:
+        """Send a ServerSyncStatusRequest."""
+        sync_status = wandb_sync_pb2.ServerSyncStatusRequest(id=id)
+        request = spb.ServerRequest(sync_status=sync_status)
+
+        handle = self._client.deliver(request)
+        return handle.map(lambda r: r.sync_status_response)
+
     def inform_init(
         self,
         settings: wandb_settings_pb2.Settings,
@@ -143,17 +184,6 @@ class ServiceConnection:
         else:
             return response.inform_attach_response.settings
 
-    def inform_start(
-        self,
-        settings: wandb_settings_pb2.Settings,
-        run_id: str,
-    ) -> None:
-        """Send a start request to the service."""
-        request = spb.ServerInformStartRequest()
-        request.settings.CopyFrom(settings)
-        request._info.stream_id = run_id
-        self._client.publish(spb.ServerRequest(inform_start=request))
-
     def teardown(self, exit_code: int) -> int | None:
         """Close the connection.