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

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

wandb/sdk/mailbox/wait_with_progress.py

@@ -62,7 +62,7 @@ def wait_all_with_progress(
     start_time = time.monotonic()
 
     async def progress_loop_with_timeout() -> list[_T]:
-        with asyncio_compat.cancel_on_exit(display_progress()):
+        async with asyncio_compat.cancel_on_exit(display_progress()):
             if timeout is not None:
                 elapsed_time = time.monotonic() - start_time
                 remaining_timeout = timeout - elapsed_time

wandb/sdk/wandb_init.py

@@ -1009,14 +1009,6 @@ class _WandbInit:
         run._set_run_obj(result.run_result.run)
 
         self._logger.info("starting run threads in backend")
-        # initiate run (stats and metadata probing)
-
-        if service:
-            assert settings.run_id
-            service.inform_start(
-                settings=settings.to_proto(),
-                run_id=settings.run_id,
-            )
 
         assert backend.interface
 

wandb/sdk/wandb_run.py

@@ -892,7 +892,19 @@ class Run:
     def tags(self, tags: Sequence) -> None:
         with telemetry.context(run=self) as tel:
             tel.feature.set_run_tags = True
-        self._settings.run_tags = tuple(tags)
+
+        try:
+            self._settings.run_tags = tuple(tags)
+        except ValueError as e:
+            # For runtime tag setting, warn instead of crash
+            # Extract the core error message without the pydantic wrapper
+            error_msg = str(e)
+            if "Value error," in error_msg:
+                # Extract the actual error message after "Value error, "
+                error_msg = error_msg.split("Value error, ")[1].split(" [type=")[0]
+            wandb.termwarn(f"Invalid tag detected: {error_msg} Tags not updated.")
+            return
+
         if self._backend and self._backend.interface:
             self._backend.interface.publish_run(self)
 

wandb/sdk/wandb_settings.py

@@ -1435,6 +1435,61 @@ class Settings(BaseModel, validate_assignment=True):
             raise UsageError("Sweep ID cannot contain only whitespace")
         return value
 
+    @field_validator("run_tags", mode="before")
+    @classmethod
+    def validate_run_tags(cls, value):
+        """Validate run tags.
+
+        Validates that each tag:
+        - Is between 1 and 64 characters in length (inclusive)
+        - Converts single string values to tuple format
+        - Preserves None values
+
+        Args:
+            value: A string, list, tuple, or None representing tags
+
+        Returns:
+            tuple: A tuple of validated tags, or None
+
+        Raises:
+            ValueError: If any tag is empty or exceeds 64 characters
+
+        <!-- lazydoc-ignore-classmethod: internal -->
+        """
+        if value is None:
+            return None
+
+        # Convert to tuple if needed
+        if isinstance(value, str):
+            tags = (value,)
+        else:
+            tags = tuple(value)
+
+        # Validate each tag and accumulate errors
+        errors = []
+        for i, tag in enumerate(tags):
+            tag_str = str(tag)
+            if len(tag_str) == 0:
+                errors.append(
+                    f"Tag at index {i} is empty. Tags must be between 1 and 64 characters"
+                )
+            elif len(tag_str) > 64:
+                # Truncate long tags for display
+                display_tag = (
+                    f"{tag_str[:20]}...{tag_str[-20:]}"
+                    if len(tag_str) > 43
+                    else tag_str
+                )
+                errors.append(
+                    f"Tag '{display_tag}' is {len(tag_str)} characters. Tags must be between 1 and 64 characters"
+                )
+
+        # Raise combined error if any validation issues were found
+        if errors:
+            raise ValueError("; ".join(errors))
+
+        return tags
+
     @field_validator("sweep_param_path", mode="before")
     @classmethod
     def validate_sweep_param_path(cls, value):

{wandb-0.21.3.dist-info → wandb-0.21.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: wandb
-Version: 0.21.3
+Version: 0.21.4
 Summary: A CLI and library for interacting with the Weights & Biases API.
 Project-URL: Source, https://github.com/wandb/wandb
 Project-URL: Bug Reports, https://github.com/wandb/wandb/issues