wandb 0.21.3__py3-none-win32.whl → 0.22.0__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
@@ -95,7 +96,7 @@ class S3Handler(StorageHandler):
         path, hit, cache_open = self._cache.check_etag_obj_path(
             URIStr(manifest_entry.ref),
             ETag(manifest_entry.digest),
-            manifest_entry.size if manifest_entry.size is not None else 0,
+            manifest_entry.size or 0,
         )
         if hit:
             return path
@@ -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/artifacts/storage_policies/_factories.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Final
+
+from requests import Response, Session
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from ..storage_handler import StorageHandler
+from ..storage_handlers.azure_handler import AzureHandler
+from ..storage_handlers.gcs_handler import GCSHandler
+from ..storage_handlers.http_handler import HTTPHandler
+from ..storage_handlers.local_file_handler import LocalFileHandler
+from ..storage_handlers.s3_handler import S3Handler
+from ..storage_handlers.wb_artifact_handler import WBArtifactHandler
+from ..storage_handlers.wb_local_artifact_handler import WBLocalArtifactHandler
+
+# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120
+# seconds, i.e. a total of 20min 6s.
+HTTP_RETRY_STRATEGY: Final[Retry] = Retry(
+    backoff_factor=1,
+    total=16,
+    status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504),
+)
+HTTP_POOL_CONNECTIONS: Final[int] = 64
+HTTP_POOL_MAXSIZE: Final[int] = 64
+
+
+def raise_for_status(response: Response, *_, **__) -> None:
+    """A `requests.Session` hook to raise for status on all requests."""
+    response.raise_for_status()
+
+
+def make_http_session() -> Session:
+    """A factory that returns a `requests.Session` for use with artifact storage handlers."""
+    session = Session()
+
+    # Explicitly configure the retry strategy for http/https adapters.
+    adapter = HTTPAdapter(
+        max_retries=HTTP_RETRY_STRATEGY,
+        pool_connections=HTTP_POOL_CONNECTIONS,
+        pool_maxsize=HTTP_POOL_MAXSIZE,
+    )
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    # Always raise on HTTP status errors.
+    session.hooks["response"].append(raise_for_status)
+    return session
+
+
+def make_storage_handlers(session: Session) -> list[StorageHandler]:
+    """A factory that returns the default artifact storage handlers."""
+    return [
+        S3Handler(),  # s3
+        GCSHandler(),  # gcs
+        AzureHandler(),  # azure
+        HTTPHandler(session, scheme="http"),  # http
+        HTTPHandler(session, scheme="https"),  # https
+        WBArtifactHandler(),  # artifact
+        WBLocalArtifactHandler(),  # local_artifact
+        LocalFileHandler(),  # file_handler
+    ]

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -16,7 +16,6 @@ from typing import IO, TYPE_CHECKING, Any, NamedTuple, Sequence
 from urllib.parse import quote
 
 import requests
-import urllib3
 
 from wandb import env
 from wandb.errors.term import termwarn
@@ -27,40 +26,24 @@ from wandb.sdk.artifacts.artifact_file_cache import (
     get_artifact_file_cache,
 )
 from wandb.sdk.artifacts.staging import get_staging_dir
-from wandb.sdk.artifacts.storage_handlers.azure_handler import AzureHandler
-from wandb.sdk.artifacts.storage_handlers.gcs_handler import GCSHandler
-from wandb.sdk.artifacts.storage_handlers.http_handler import HTTPHandler
-from wandb.sdk.artifacts.storage_handlers.local_file_handler import LocalFileHandler
 from wandb.sdk.artifacts.storage_handlers.multi_handler import MultiHandler
-from wandb.sdk.artifacts.storage_handlers.s3_handler import S3Handler
 from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandler
-from wandb.sdk.artifacts.storage_handlers.wb_artifact_handler import WBArtifactHandler
-from wandb.sdk.artifacts.storage_handlers.wb_local_artifact_handler import (
-    WBLocalArtifactHandler,
-)
 from wandb.sdk.artifacts.storage_layout import StorageLayout
 from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
 from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings
-from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id
+from wandb.sdk.lib.hashutil import b64_to_hex_id, hex_to_b64_id
 from wandb.sdk.lib.paths import FilePathStr, URIStr
 
+from ._factories import make_http_session, make_storage_handlers
+
 if TYPE_CHECKING:
     from wandb.filesync.step_prepare import StepPrepare
     from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal import progress
 
-# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120
-# seconds, i.e. a total of 20min 6s.
-_REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry(
-    backoff_factor=1,
-    total=16,
-    status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504),
-)
-_REQUEST_POOL_CONNECTIONS = 64
-_REQUEST_POOL_MAXSIZE = 64
 
 # AWS S3 max upload parts without having to make additional requests for extra parts
 S3_MAX_PART_NUMBERS = 1000
@@ -96,48 +79,23 @@ class WandbStoragePolicy(StoragePolicy):
 
     @classmethod
     def from_config(
-        cls, config: dict, api: InternalApi | None = None
+        cls, config: dict[str, Any], api: InternalApi | None = None
     ) -> WandbStoragePolicy:
         return cls(config=config, api=api)
 
     def __init__(
         self,
-        config: dict | None = None,
+        config: dict[str, Any] | None = None,
         cache: ArtifactFileCache | None = None,
         api: InternalApi | None = None,
+        session: requests.Session | None = None,
     ) -> None:
-        self._cache = cache or get_artifact_file_cache()
         self._config = config or {}
-        self._session = requests.Session()
-        adapter = requests.adapters.HTTPAdapter(
-            max_retries=_REQUEST_RETRY_STRATEGY,
-            pool_connections=_REQUEST_POOL_CONNECTIONS,
-            pool_maxsize=_REQUEST_POOL_MAXSIZE,
-        )
-        self._session.mount("http://", adapter)
-        self._session.mount("https://", adapter)
-
-        s3 = S3Handler()
-        gcs = GCSHandler()
-        azure = AzureHandler()
-        http = HTTPHandler(self._session)
-        https = HTTPHandler(self._session, scheme="https")
-        artifact = WBArtifactHandler()
-        local_artifact = WBLocalArtifactHandler()
-        file_handler = LocalFileHandler()
-
+        self._cache = cache or get_artifact_file_cache()
+        self._session = session or make_http_session()
         self._api = api or InternalApi()
         self._handler = MultiHandler(
-            handlers=[
-                s3,
-                gcs,
-                azure,
-                http,
-                https,
-                artifact,
-                local_artifact,
-                file_handler,
-            ],
+            handlers=make_storage_handlers(self._session),
             default_handler=TrackingHandler(),
         )
 
@@ -167,54 +125,52 @@ class WandbStoragePolicy(StoragePolicy):
             self._cache._override_cache_path = dest_path
 
         path, hit, cache_open = self._cache.check_md5_obj_path(
-            B64MD5(manifest_entry.digest),
-            manifest_entry.size if manifest_entry.size is not None else 0,
+            manifest_entry.digest,
+            size=manifest_entry.size or 0,
         )
         if hit:
             return path
 
-        if manifest_entry._download_url is not None:
+        if (url := manifest_entry._download_url) is not None:
             # Use multipart parallel download for large file
             if (
-                executor is not None
-                and manifest_entry.size is not None
-                and self._should_multipart_download(manifest_entry.size, multipart)
+                executor
+                and (size := manifest_entry.size)
+                and self._should_multipart_download(size, multipart)
             ):
-                self._multipart_file_download(
-                    executor,
-                    manifest_entry._download_url,
-                    manifest_entry.size,
-                    cache_open,
-                )
+                self._multipart_file_download(executor, url, size, cache_open)
                 return path
+
             # Serial download
-            response = self._session.get(manifest_entry._download_url, stream=True)
             try:
-                response.raise_for_status()
-            except Exception:
+                response = self._session.get(url, stream=True)
+            except requests.HTTPError:
                 # Signed URL might have expired, fall back to fetching it one by one.
                 manifest_entry._download_url = None
+
         if manifest_entry._download_url is None:
             auth = None
-            http_headers = _thread_local_api_settings.headers or {}
-            if self._api.access_token is not None:
-                http_headers["Authorization"] = f"Bearer {self._api.access_token}"
-            elif _thread_local_api_settings.cookies is None:
+            headers = _thread_local_api_settings.headers
+            cookies = _thread_local_api_settings.cookies
+
+            # For auth, prefer using (in order): auth header, cookies, HTTP Basic Auth
+            if token := self._api.access_token:
+                headers = {**(headers or {}), "Authorization": f"Bearer {token}"}
+            elif cookies is not None:
+                pass
+            else:
                 auth = ("api", self._api.api_key or "")
+
+            file_url = self._file_url(
+                self._api,
+                artifact.entity,
+                artifact.project,
+                artifact.name.split(":")[0],
+                manifest_entry,
+            )
             response = self._session.get(
-                self._file_url(
-                    self._api,
-                    artifact.entity,
-                    artifact.project,
-                    artifact.name.split(":")[0],
-                    manifest_entry,
-                ),
-                auth=auth,
-                cookies=_thread_local_api_settings.cookies,
-                headers=http_headers,
-                stream=True,
+                file_url, auth=auth, cookies=cookies, headers=headers, stream=True
             )
-            response.raise_for_status()
 
         with cache_open(mode="wb") as file:
             for data in response.iter_content(chunk_size=16 * 1024):
@@ -269,12 +225,7 @@ class WandbStoragePolicy(StoragePolicy):
         # Other threads has error, no need to start
         if download_has_error.is_set():
             return
-        response = self._session.get(
-            url=download_url,
-            headers=headers,
-            stream=True,
-        )
-        response.raise_for_status()
+        response = self._session.get(url=download_url, headers=headers, stream=True)
 
         file_offset = start
         for content in response.iter_content(chunk_size=_HTTP_RES_CHUNK_SIZE_BYTES):
@@ -376,43 +327,27 @@ class WandbStoragePolicy(StoragePolicy):
         entity_name: str,
         project_name: str,
         artifact_name: str,
-        manifest_entry: ArtifactManifestEntry,
+        entry: ArtifactManifestEntry,
     ) -> str:
-        storage_layout = self._config.get("storageLayout", StorageLayout.V1)
-        storage_region = self._config.get("storageRegion", "default")
-        md5_hex = b64_to_hex_id(B64MD5(manifest_entry.digest))
+        layout = self._config.get("storageLayout", StorageLayout.V1)
+        region = self._config.get("storageRegion", "default")
+        md5_hex = b64_to_hex_id(entry.digest)
 
-        if storage_layout == StorageLayout.V1:
-            return "{}/artifacts/{}/{}".format(
-                api.settings("base_url"), entity_name, md5_hex
-            )
-        elif storage_layout == StorageLayout.V2:
+        base_url: str = api.settings("base_url")
+
+        if layout == StorageLayout.V1:
+            return f"{base_url}/artifacts/{entity_name}/{md5_hex}"
+
+        if layout == StorageLayout.V2:
+            birth_artifact_id = entry.birth_artifact_id or ""
             if api._server_supports(
-                ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER  # type: ignore
+                ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER
             ):
-                return "{}/artifactsV2/{}/{}/{}/{}/{}/{}/{}".format(
-                    api.settings("base_url"),
-                    storage_region,
-                    quote(entity_name),
-                    quote(project_name),
-                    quote(artifact_name),
-                    quote(manifest_entry.birth_artifact_id or ""),
-                    md5_hex,
-                    manifest_entry.path.name,
-                )
-            return "{}/artifactsV2/{}/{}/{}/{}".format(
-                api.settings("base_url"),
-                storage_region,
-                entity_name,
-                quote(
-                    manifest_entry.birth_artifact_id
-                    if manifest_entry.birth_artifact_id is not None
-                    else ""
-                ),
-                md5_hex,
-            )
-        else:
-            raise Exception(f"unrecognized storage layout: {storage_layout}")
+                return f"{base_url}/artifactsV2/{region}/{quote(entity_name)}/{quote(project_name)}/{quote(artifact_name)}/{quote(birth_artifact_id)}/{md5_hex}/{entry.path.name}"
+
+            return f"{base_url}/artifactsV2/{region}/{entity_name}/{quote(birth_artifact_id)}/{md5_hex}"
+
+        raise ValueError(f"unrecognized storage layout: {layout!r}")
 
     def s3_multipart_file_upload(
         self,
@@ -486,7 +421,7 @@ class WandbStoragePolicy(StoragePolicy):
             True if the file was a duplicate (did not need to be uploaded),
             False if it needed to be uploaded or was a reference (nothing to dedupe).
         """
-        file_size = entry.size if entry.size is not None else 0
+        file_size = entry.size or 0
         chunk_size = self.calc_chunk_size(file_size)
         upload_parts = []
         hex_digests = {}
@@ -562,8 +497,8 @@ class WandbStoragePolicy(StoragePolicy):
 
         # Cache upon successful upload.
         _, hit, cache_open = self._cache.check_md5_obj_path(
-            B64MD5(entry.digest),
-            entry.size if entry.size is not None else 0,
+            entry.digest,
+            size=entry.size or 0,
         )
 
         staging_dir = get_staging_dir()

wandb/sdk/interface/interface.py

@@ -883,6 +883,16 @@ class InterfaceBase:
     ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
+    def publish_probe_system_info(self) -> None:
+        probe_system_info = pb.ProbeSystemInfoRequest()
+        return self._publish_probe_system_info(probe_system_info)
+
+    @abstractmethod
+    def _publish_probe_system_info(
+        self, probe_system_info: pb.ProbeSystemInfoRequest
+    ) -> None:
+        raise NotImplementedError
+
     def join(self) -> None:
         # Drop indicates that the internal process has already been shutdown
         if self._drop:

wandb/sdk/interface/interface_shared.py

@@ -112,6 +112,7 @@ class InterfaceShared(InterfaceBase):
         python_packages: Optional[pb.PythonPackagesRequest] = None,
         job_input: Optional[pb.JobInputRequest] = None,
         run_finish_without_exit: Optional[pb.RunFinishWithoutExitRequest] = None,
+        probe_system_info: Optional[pb.ProbeSystemInfoRequest] = None,
     ) -> pb.Record:
         request = pb.Request()
         if get_summary:
@@ -178,6 +179,8 @@ class InterfaceShared(InterfaceBase):
             request.job_input.CopyFrom(job_input)
         elif run_finish_without_exit:
             request.run_finish_without_exit.CopyFrom(run_finish_without_exit)
+        elif probe_system_info:
+            request.probe_system_info.CopyFrom(probe_system_info)
         else:
             raise Exception("Invalid request")
         record = self._make_record(request=request)
@@ -330,6 +333,12 @@ class InterfaceShared(InterfaceBase):
         rec = self._make_record(use_artifact=use_artifact)
         self._publish(rec)
 
+    def _publish_probe_system_info(
+        self, probe_system_info: pb.ProbeSystemInfoRequest
+    ) -> None:
+        record = self._make_request(probe_system_info=probe_system_info)
+        self._publish(record)
+
     def _deliver_artifact(
         self,
         log_artifact: pb.LogArtifactRequest,

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,