wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

wandb/errors/__init__.py

@@ -2,30 +2,22 @@ __all__ = [
     "Error",
     "UsageError",
     "CommError",
-    "LogError",
-    "DockerError",
-    "LogMultiprocessError",
-    "MultiprocessError",
-    "RequireError",
-    "ExecutionError",
-    "LaunchError",
-    "SweepError",
+    "UnsupportedError",
     "WaitTimeoutError",
-    "ContextCancelledError",
-    "ServiceStartProcessError",
-    "ServiceStartTimeoutError",
-    "ServiceStartPortError",
 ]
 
-from typing import List, Optional
+from typing import Optional
 
 
 class Error(Exception):
     """Base W&B Error"""
 
-    def __init__(self, message) -> None:
+    def __init__(self, message, context: Optional[dict] = None) -> None:
         super().__init__(message)
         self.message = message
+        # sentry context capture
+        if context:
+            self.context = context
 
     # For python 2 support
     def encode(self, encoding):
@@ -47,111 +39,11 @@ class UsageError(Error):
     pass
 
 
-class LogError(Error):
-    """Raised when wandb.log() fails"""
-
-    pass
-
-
-class LogMultiprocessError(LogError):
-    """Raised when wandb.log() fails because of multiprocessing"""
-
-    pass
-
-
-class MultiprocessError(Error):
-    """Raised when fails because of multiprocessing"""
-
-    pass
-
-
-class RequireError(Error):
-    """Raised when wandb.require() fails"""
-
-    pass
-
-
-class ExecutionError(Error):
-    """Generic execution exception"""
-
-    pass
-
-
-class DockerError(Error):
-    """Raised when attempting to execute a docker command"""
-
-    def __init__(
-        self,
-        command_launched: List[str],
-        return_code: int,
-        stdout: Optional[bytes] = None,
-        stderr: Optional[bytes] = None,
-    ) -> None:
-        command_launched_str = " ".join(command_launched)
-        error_msg = (
-            f"The docker command executed was `{command_launched_str}`.\n"
-            f"It returned with code {return_code}\n"
-        )
-        if stdout is not None:
-            error_msg += f"The content of stdout is '{stdout.decode()}'\n"
-        else:
-            error_msg += (
-                "The content of stdout can be found above the "
-                "stacktrace (it wasn't captured).\n"
-            )
-        if stderr is not None:
-            error_msg += f"The content of stderr is '{stderr.decode()}'\n"
-        else:
-            error_msg += (
-                "The content of stderr can be found above the "
-                "stacktrace (it wasn't captured)."
-            )
-        super().__init__(error_msg)
-
-
-class LaunchError(Error):
-    """Raised when a known error occurs in wandb launch"""
-
-    pass
-
-
-class SweepError(Error):
-    """Raised when a known error occurs with wandb sweeps"""
-
-    pass
+class UnsupportedError(UsageError):
+    """Raised when trying to use a feature that is not supported"""
 
 
 class WaitTimeoutError(Error):
     """Raised when wait() timeout occurs before process is finished"""
 
     pass
-
-
-class MailboxError(Error):
-    """Generic Mailbox Exception"""
-
-    pass
-
-
-class ContextCancelledError(Error):
-    """Context cancelled Exception"""
-
-    pass
-
-
-class ServiceStartProcessError(Error):
-    """Raised when a known error occurs when launching wandb service"""
-
-    pass
-
-
-class ServiceStartTimeoutError(Error):
-    """Raised when service start times out"""
-
-    pass
-
-
-class ServiceStartPortError(Error):
-    """Raised when service start fails to find a port"""
-
-    pass

wandb/errors/term.py

@@ -17,7 +17,7 @@ _show_errors = True
 _logger = None
 
 
-def termsetup(settings, logger):
+def termsetup(settings, logger) -> None:
     global _silent, _show_info, _show_warnings, _show_errors, _logger
     _silent = settings.silent
     _show_info = settings.show_info

wandb/fastai/__init__.py

@@ -1,5 +1,4 @@
-"""
-Compatibility fastai module.
+"""Compatibility fastai module.
 
 In the future use:
     from wandb.integration.fastai import WandbCallback

wandb/filesync/dir_watcher.py

@@ -71,7 +71,7 @@ class FileEventHandler(abc.ABC):
 
 
 class PolicyNow(FileEventHandler):
-    """This policy only uploads files now"""
+    """This policy only uploads files now."""
 
     def on_modified(self, force: bool = False) -> None:
         # only upload if we've never uploaded or when .save is called
@@ -88,7 +88,7 @@ class PolicyNow(FileEventHandler):
 
 
 class PolicyEnd(FileEventHandler):
-    """This policy only updates at the end of the run"""
+    """This policy only updates at the end of the run."""
 
     def on_modified(self, force: bool = False) -> None:
         pass
@@ -106,8 +106,11 @@ class PolicyEnd(FileEventHandler):
 
 
 class PolicyLive(FileEventHandler):
-    """This policy will upload files every RATE_LIMIT_SECONDS as it
-    changes throttling as the size increases"""
+    """Event handler that uploads respecting throttling.
+
+    Uploads files every RATE_LIMIT_SECONDS, which changes as the size increases to deal
+    with throttling.
+    """
 
     RATE_LIMIT_SECONDS = 15
     unit_dict = dict(util.POW_10_BYTES)
@@ -250,7 +253,7 @@ class DirWatcher:
             feh.on_modified(force=True)
 
     def _per_file_event_handler(self) -> "wd_events.FileSystemEventHandler":
-        """Create a Watchdog file event handler that does different things for every file"""
+        """Create a Watchdog file event handler that does different things for every file."""
         file_event_handler = wd_events.PatternMatchingEventHandler()
         file_event_handler.on_created = self._on_file_created
         file_event_handler.on_modified = self._on_file_modified

wandb/filesync/step_prepare.py

@@ -1,11 +1,11 @@
 """Batching file prepare requests to our API."""
 
 import queue
-import sys
 import threading
 import time
 from typing import (
     TYPE_CHECKING,
+    Callable,
     List,
     Mapping,
     NamedTuple,
@@ -16,31 +16,16 @@ from typing import (
 )
 
 if TYPE_CHECKING:
-    from wandb.sdk.internal import internal_api
-
-    if sys.version_info >= (3, 8):
-        from typing import Protocol
-    else:
-        from typing_extensions import Protocol
-
-    class DoPrepareFn(Protocol):
-        def __call__(self) -> "internal_api.CreateArtifactFileSpecInput":
-            pass
-
-    class OnPrepareFn(Protocol):
-        def __call__(
-            self,
-            upload_url: Optional[str],  # GraphQL type File.uploadUrl
-            upload_headers: Sequence[str],  # GraphQL type File.uploadHeaders
-            artifact_id: str,  # GraphQL type File.artifact.id
-        ) -> None:
-            pass
+    from wandb.sdk.internal.internal_api import (
+        Api,
+        CreateArtifactFileSpecInput,
+        CreateArtifactFilesResponseFile,
+    )
 
 
 # Request for a file to be prepared.
 class RequestPrepare(NamedTuple):
-    prepare_fn: "DoPrepareFn"
-    on_prepare: Optional["OnPrepareFn"]
+    file_spec: "CreateArtifactFileSpecInput"
     response_queue: "queue.Queue[ResponsePrepare]"
 
 
@@ -54,7 +39,49 @@ class ResponsePrepare(NamedTuple):
     birth_artifact_id: str
 
 
-Event = Union[RequestPrepare, RequestFinish, ResponsePrepare]
+Request = Union[RequestPrepare, RequestFinish]
+
+
+def _clamp(x: float, low: float, high: float) -> float:
+    return max(low, min(x, high))
+
+
+def gather_batch(
+    request_queue: "queue.Queue[Request]",
+    batch_time: float,
+    inter_event_time: float,
+    max_batch_size: int,
+    clock: Callable[[], float] = time.monotonic,
+) -> Tuple[bool, Sequence[RequestPrepare]]:
+
+    batch_start_time = clock()
+    remaining_time = batch_time
+
+    first_request = request_queue.get()
+    if isinstance(first_request, RequestFinish):
+        return True, []
+
+    batch: List[RequestPrepare] = [first_request]
+
+    while remaining_time > 0 and len(batch) < max_batch_size:
+        try:
+            request = request_queue.get(
+                timeout=_clamp(
+                    x=inter_event_time,
+                    low=1e-12,  # 0 = "block forever", so just use something tiny
+                    high=remaining_time,
+                ),
+            )
+            if isinstance(request, RequestFinish):
+                return True, batch
+
+            batch.append(request)
+            remaining_time = batch_time - (clock() - batch_start_time)
+
+        except queue.Empty:
+            break
+
+    return False, batch
 
 
 class StepPrepare:
@@ -66,68 +93,46 @@ class StepPrepare:
 
     def __init__(
         self,
-        api: "internal_api.Api",
+        api: "Api",
         batch_time: float,
         inter_event_time: float,
         max_batch_size: int,
+        request_queue: Optional["queue.Queue[Request]"] = None,
     ) -> None:
         self._api = api
         self._inter_event_time = inter_event_time
         self._batch_time = batch_time
         self._max_batch_size = max_batch_size
-        self._request_queue: "queue.Queue[RequestPrepare | RequestFinish]" = (
-            queue.Queue()
-        )
+        self._request_queue: "queue.Queue[Request]" = request_queue or queue.Queue()
         self._thread = threading.Thread(target=self._thread_body)
         self._thread.daemon = True
 
     def _thread_body(self) -> None:
         while True:
-            request = self._request_queue.get()
-            if isinstance(request, RequestFinish):
-                break
-            finish, batch = self._gather_batch(request)
-            prepare_response = self._prepare_batch(batch)
-            # send responses
-            for prepare_request in batch:
-                name = prepare_request.prepare_fn()["name"]
-                response_file = prepare_response[name]
-                upload_url = response_file["uploadUrl"]
-                upload_headers = response_file["uploadHeaders"]
-                birth_artifact_id = response_file["artifact"]["id"]
-                if prepare_request.on_prepare:
-                    prepare_request.on_prepare(
-                        upload_url, upload_headers, birth_artifact_id
+            finish, batch = gather_batch(
+                request_queue=self._request_queue,
+                batch_time=self._batch_time,
+                inter_event_time=self._inter_event_time,
+                max_batch_size=self._max_batch_size,
+            )
+            if batch:
+                prepare_response = self._prepare_batch(batch)
+                # send responses
+                for prepare_request in batch:
+                    name = prepare_request.file_spec["name"]
+                    response_file = prepare_response[name]
+                    upload_url = response_file["uploadUrl"]
+                    upload_headers = response_file["uploadHeaders"]
+                    birth_artifact_id = response_file["artifact"]["id"]
+                    prepare_request.response_queue.put(
+                        ResponsePrepare(upload_url, upload_headers, birth_artifact_id)
                     )
-                prepare_request.response_queue.put(
-                    ResponsePrepare(upload_url, upload_headers, birth_artifact_id)
-                )
             if finish:
                 break
 
-    def _gather_batch(
-        self, first_request: RequestPrepare
-    ) -> Tuple[bool, Sequence[RequestPrepare]]:
-        batch_start_time = time.time()
-        batch: List[RequestPrepare] = [first_request]
-        while True:
-            try:
-                request = self._request_queue.get(
-                    block=True, timeout=self._inter_event_time
-                )
-                if isinstance(request, RequestFinish):
-                    return True, batch
-                batch.append(request)
-                remaining_time = self._batch_time - (time.time() - batch_start_time)
-                if remaining_time < 0 or len(batch) >= self._max_batch_size:
-                    break
-            except queue.Empty:
-                break
-        return False, batch
-
     def _prepare_batch(
         self, batch: Sequence[RequestPrepare]
-    ) -> Mapping[str, "internal_api.CreateArtifactFilesResponseFile"]:
+    ) -> Mapping[str, "CreateArtifactFilesResponseFile"]:
         """Execute the prepareFiles API call.
 
         Arguments:
@@ -137,14 +142,10 @@ class StepPrepare:
                 an uploadUrl key. The value of the uploadUrl key is None if the file
                 already exists, or a url string if the file should be uploaded.
         """
-        file_specs: List["internal_api.CreateArtifactFileSpecInput"] = []
-        for prepare_request in batch:
-            file_spec = prepare_request.prepare_fn()
-            file_specs.append(file_spec)
-        return self._api.create_artifact_files(file_specs)
+        return self._api.create_artifact_files([req.file_spec for req in batch])
 
     def prepare_async(
-        self, prepare_fn: "DoPrepareFn", on_prepare: Optional["OnPrepareFn"] = None
+        self, file_spec: "CreateArtifactFileSpecInput"
     ) -> "queue.Queue[ResponsePrepare]":
         """Request the backend to prepare a file for upload.
 
@@ -153,11 +154,11 @@ class StepPrepare:
                 either a file upload url, or None if the file doesn't need to be uploaded.
         """
         response_queue: "queue.Queue[ResponsePrepare]" = queue.Queue()
-        self._request_queue.put(RequestPrepare(prepare_fn, on_prepare, response_queue))
+        self._request_queue.put(RequestPrepare(file_spec, response_queue))
         return response_queue
 
-    def prepare(self, prepare_fn: "DoPrepareFn") -> ResponsePrepare:
-        return self.prepare_async(prepare_fn).get()
+    def prepare(self, file_spec: "CreateArtifactFileSpecInput") -> ResponsePrepare:
+        return self.prepare_async(file_spec).get()
 
     def start(self) -> None:
         self._thread.start()

wandb/filesync/step_upload.py

@@ -203,7 +203,7 @@ class StepUpload:
         self._spawn_upload(job)
 
     def _spawn_upload(self, job: upload_job.UploadJob) -> None:
-        """Spawns an upload job, and handles the bookkeeping of `self._running_jobs`.
+        """Spawn an upload job, and handles the bookkeeping of `self._running_jobs`.
 
         Context: it's important that, whenever we add an entry to `self._running_jobs`,
         we ensure that a corresponding `EventJobDone` message will eventually get handled;
@@ -214,7 +214,6 @@ class StepUpload:
         to `self._running_jobs` is textually right next to the code that eventually enqueues
         the `EventJobDone` message. This should help keep them in sync.
         """
-
         # Adding the entry to `self._running_jobs` MUST happen in the main thread,
         # NOT in the job that gets submitted to the thread-pool, to guard against
         # this sequence of events:

wandb/integration/catboost/__init__.py

@@ -1,6 +1,4 @@
-"""
-W&B callback for CatBoost
-"""
+"""W&B callback for CatBoost."""
 
 from .catboost import WandbCallback, log_summary
 

wandb/integration/catboost/catboost.py

@@ -1,6 +1,4 @@
-"""
-catboost init
-"""
+"""catboost init."""
 
 from pathlib import Path
 from types import SimpleNamespace
@@ -65,9 +63,7 @@ class WandbCallback:
 def _checkpoint_artifact(
     model: Union[CatBoostClassifier, CatBoostRegressor], aliases: List[str]
 ) -> None:
-    """
-    Upload model checkpoint as W&B artifact
-    """
+    """Upload model checkpoint as W&B artifact."""
     if wandb.run is None:
         raise wandb.Error(
             "You must call `wandb.init()` before `_checkpoint_artifact()`"
@@ -87,9 +83,7 @@ def _checkpoint_artifact(
 def _log_feature_importance(
     model: Union[CatBoostClassifier, CatBoostRegressor]
 ) -> None:
-    """
-    Log feature importance with default settings.
-    """
+    """Log feature importance with default settings."""
     if wandb.run is None:
         raise wandb.Error(
             "You must call `wandb.init()` before `_checkpoint_artifact()`"
@@ -119,7 +113,7 @@ def log_summary(
     save_model_checkpoint: bool = False,
     log_feature_importance: bool = True,
 ) -> None:
-    """`log_summary` logs useful metrics about catboost model after training is done
+    """`log_summary` logs useful metrics about catboost model after training is done.
 
     Arguments:
         model: it can be CatBoostClassifier or CatBoostRegressor.
@@ -136,13 +130,13 @@ def log_summary(
 
     Example:
         ```python
-        train_pool = Pool(train[features], label=train['label'], cat_features=cat_features)
-        test_pool = Pool(test[features], label=test['label'], cat_features=cat_features)
+        train_pool = Pool(train[features], label=train["label"], cat_features=cat_features)
+        test_pool = Pool(test[features], label=test["label"], cat_features=cat_features)
 
         model = CatBoostRegressor(
             iterations=100,
-            loss_function='Cox',
-            eval_metric='Cox',
+            loss_function="Cox",
+            eval_metric="Cox",
         )
 
         model.fit(

wandb/integration/fastai/__init__.py

@@ -1,5 +1,5 @@
-"""
-This module hooks fast.ai v1 Learners to Weights & Biases through a callback.
+"""Hooks that add fast.ai v1 Learners to Weights & Biases through a callback.
+
 Requested logged data can be configured through the callback constructor.
 
 Examples:
@@ -61,8 +61,8 @@ except ImportError:
 
 
 class WandbCallback(TrackerCallback):
-    """
-    Automatically saves model topology, losses & metrics.
+    """Callback for saving model topology, losses & metrics.
+
     Optionally logs weights, gradients, sample predictions and best trained model.
 
     Arguments:
@@ -92,7 +92,6 @@ class WandbCallback(TrackerCallback):
         predictions: int = 36,
         seed: int = 12345,
     ) -> None:
-
         # Check if wandb.init has been called
         if wandb.run is None:
             raise ValueError("You must call wandb.init() before WandbCallback()")
@@ -119,8 +118,7 @@ class WandbCallback(TrackerCallback):
             self.validation_data = [learn.data.valid_ds[i] for i in indices]
 
     def on_train_begin(self, **kwargs: Any) -> None:
-        """Call watch method to log model topology, gradients & weights"""
-
+        """Call watch method to log model topology, gradients & weights."""
         # Set self.best, method inherited from "TrackerCallback" by "SaveModelCallback"
         super().on_train_begin()
 
@@ -134,8 +132,7 @@ class WandbCallback(TrackerCallback):
     def on_epoch_end(
         self, epoch: int, smooth_loss: float, last_metrics: list, **kwargs: Any
     ) -> None:
-        """Logs training loss, validation loss and custom metrics & log prediction samples & save model"""
-
+        """Log training loss, validation loss and custom metrics & log prediction samples & save model."""
         if self.save_model:
             # Adapted from fast.ai "SaveModelCallback"
             current = self.get_monitor_value()
@@ -174,7 +171,6 @@ class WandbCallback(TrackerCallback):
 
     def on_train_end(self, **kwargs: Any) -> None:
         """Load the best model."""
-
         if self.save_model:
             # Adapted from fast.ai "SaveModelCallback"
             if self.model_path.is_file():
@@ -183,8 +179,7 @@ class WandbCallback(TrackerCallback):
                     print(f"Loaded best saved model from {self.model_path}")
 
     def _wandb_log_predictions(self) -> None:
-        """Log prediction samples"""
-
+        """Log prediction samples."""
         pred_log = []
 
         if self.validation_data is None:
@@ -234,7 +229,6 @@ class WandbCallback(TrackerCallback):
             elif hasattr(y, "shape") and (
                 (len(y.shape) == 2) or (len(y.shape) == 3 and y.shape[0] in [1, 3, 4])
             ):
-
                 pred_log.extend(
                     [
                         wandb.Image(x.data, caption="Input data", grouping=3),

wandb/integration/gym/__init__.py

@@ -1,21 +1,52 @@
 import re
+import sys
 from typing import Optional
 
 import wandb
+import wandb.util
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 
 _gym_version_lt_0_26: Optional[bool] = None
+_required_error_msg = (
+    "Couldn't import the gymnasium python package, "
+    "install with `pip install gymnasium`"
+)
+GymLib = Literal["gym", "gymnasium"]
 
 
 def monitor():
+    """Monitor a gym environment.
+
+    Supports both gym and gymnasium.
+    """
+    gym_lib: Optional[GymLib] = None
+
+    # gym is not maintained anymore, gymnasium is the drop-in replacement - prefer it
+    if wandb.util.get_module("gymnasium") is not None:
+        gym_lib = "gymnasium"
+    elif wandb.util.get_module("gym") is not None:
+        gym_lib = "gym"
+
+    if gym_lib is None:
+        raise wandb.Error(_required_error_msg)
+
     vcr = wandb.util.get_module(
-        "gym.wrappers.monitoring.video_recorder",
-        required="Couldn't import the gym python package, install with `pip install gym`",
+        f"{gym_lib}.wrappers.monitoring.video_recorder",
+        required=_required_error_msg,
     )
 
     global _gym_version_lt_0_26
 
     if _gym_version_lt_0_26 is None:
-        import gym  # type: ignore
+        if gym_lib == "gym":
+            import gym
+        else:
+            import gymnasium as gym  # type: ignore
         from pkg_resources import parse_version
 
         if parse_version(gym.__version__) < parse_version("0.26.0"):
@@ -47,7 +78,7 @@ def monitor():
     recorder.close = close
     wandb.patched["gym"].append(
         [
-            f"gym.wrappers.monitoring.video_recorder.{vcr_recorder_attribute}",
+            f"{gym_lib}.wrappers.monitoring.video_recorder.{vcr_recorder_attribute}",
             "close",
         ]
     )

wandb/integration/keras/__init__.py

@@ -1,6 +1,6 @@
-"""
-Tools for integrating `wandb` with [`Keras`](https://keras.io/),
-a deep learning API for [`TensorFlow`](https://www.tensorflow.org/).
+"""Tools for integrating `wandb` with [`Keras`](https://keras.io/).
+
+Keras is a deep learning API for [`TensorFlow`](https://www.tensorflow.org/).
 """
 __all__ = (
     "WandbCallback",