wandb 0.16.5__py3-none-any.whl → 0.17.0__py3-none-any.whl

wandb/filesync/step_upload.py

@@ -1,6 +1,5 @@
 """Batching file prepare requests to our API."""
 
-import asyncio
 import concurrent.futures
 import logging
 import queue
@@ -8,7 +7,6 @@ import sys
 import threading
 from typing import (
     TYPE_CHECKING,
-    Awaitable,
     Callable,
     MutableMapping,
     MutableSequence,
@@ -43,7 +41,6 @@ if TYPE_CHECKING:
 PreCommitFn = Callable[[], None]
 OnRequestFinishFn = Callable[[], None]
 SaveFn = Callable[["progress.ProgressFn"], bool]
-SaveFnAsync = Callable[["progress.ProgressFn"], Awaitable[bool]]
 
 logger = logging.getLogger(__name__)
 
@@ -55,7 +52,6 @@ class RequestUpload(NamedTuple):
     md5: Optional[str]
     copied: bool
     save_fn: Optional[SaveFn]
-    save_fn_async: Optional[SaveFnAsync]
     digest: Optional[str]
 
 
@@ -78,47 +74,6 @@ class EventJobDone(NamedTuple):
 Event = Union[RequestUpload, RequestCommitArtifact, RequestFinish, EventJobDone]
 
 
-class AsyncExecutor:
-    """Runs async file uploads in a background thread."""
-
-    def __init__(
-        self,
-        pool: concurrent.futures.ThreadPoolExecutor,
-        concurrency_limit: Optional[int],
-    ) -> None:
-        self.loop = asyncio.new_event_loop()
-        self.loop.set_default_executor(pool)
-        self.loop_thread = threading.Thread(
-            target=self.loop.run_forever,
-            daemon=True,
-            name="wandb-upload-async",
-        )
-
-        self.concurrency_limiter = asyncio.Semaphore(
-            value=concurrency_limit or 128,
-            # Before Python 3.10: if we don't set `loop=loop`,
-            #   then the Semaphore will bind to the wrong event loop,
-            #   causing errors when a coroutine tries to wait for it;
-            #   see https://pastebin.com/XcrS9suX .
-            # After 3.10: the `loop` argument doesn't exist.
-            # So we need to only conditionally pass in `loop`.
-            **({} if sys.version_info >= (3, 10) else {"loop": self.loop}),
-        )
-
-    def start(self) -> None:
-        self.loop_thread.start()
-
-    def stop(self) -> None:
-        self.loop.call_soon_threadsafe(self.loop.stop)
-
-    def submit(self, coro: Awaitable[None]) -> None:
-        async def run_with_limiter() -> None:
-            async with self.concurrency_limiter:
-                await coro
-
-        asyncio.run_coroutine_threadsafe(run_with_limiter(), self.loop)
-
-
 class StepUpload:
     def __init__(
         self,
@@ -142,15 +97,6 @@ class StepUpload:
             max_workers=max_threads,
         )
 
-        self._async_executor = (
-            AsyncExecutor(
-                pool=self._pool,
-                concurrency_limit=settings._async_upload_concurrency_limit,
-            )
-            if settings is not None and settings._async_upload_concurrency_limit
-            else None
-        )
-
         # Indexed by files' `save_name`'s, which are their ID's in the Run.
         self._running_jobs: MutableMapping[LogicalPath, RequestUpload] = {}
         self._pending_jobs: MutableSequence[RequestUpload] = []
@@ -186,8 +132,6 @@ class StepUpload:
             elif not self._running_jobs:
                 # Queue was empty and no jobs left.
                 self._pool.shutdown(wait=False)
-                if self._async_executor:
-                    self._async_executor.stop()
                 if finish_callback:
                     finish_callback()
                 break
@@ -245,18 +189,9 @@ class StepUpload:
             self._pending_jobs.append(event)
             return
 
-        if self._async_executor and event.save_fn_async is not None:
-            # (The `and save_fn_async is not None` is because the async code path
-            # doesn't support all uploads yet: even if the user has requested async,
-            # we sometimes need to use the sync method instead.)
-            self._spawn_upload_async(
-                event,
-                async_executor=self._async_executor,
-            )
-        else:
-            self._spawn_upload_sync(event)
+        self._spawn_upload(event)
 
-    def _spawn_upload_sync(self, event: RequestUpload) -> None:
+    def _spawn_upload(self, event: RequestUpload) -> None:
         """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`,
@@ -285,35 +220,13 @@ class StepUpload:
 
         def run_and_notify() -> None:
             try:
-                self._do_upload_sync(event)
+                self._do_upload(event)
             finally:
                 self._event_queue.put(EventJobDone(event, exc=sys.exc_info()[1]))
 
         self._pool.submit(run_and_notify)
 
-    def _spawn_upload_async(
-        self,
-        event: RequestUpload,
-        async_executor: AsyncExecutor,
-    ) -> None:
-        """Equivalent to _spawn_upload_sync, but uses the async event loop instead of a thread, and requires `event.save_fn_async`.
-
-        Raises:
-            AssertionError: if `event.save_fn_async` is None.
-        """
-        assert event.save_fn_async is not None
-
-        self._running_jobs[event.save_name] = event
-
-        async def run_and_notify() -> None:
-            try:
-                await self._do_upload_async(event)
-            finally:
-                self._event_queue.put(EventJobDone(event, exc=sys.exc_info()[1]))
-
-        async_executor.submit(run_and_notify())
-
-    def _do_upload_sync(self, event: RequestUpload) -> None:
+    def _do_upload(self, event: RequestUpload) -> None:
         job = upload_job.UploadJob(
             self._stats,
             self._api,
@@ -329,19 +242,6 @@ class StepUpload:
         )
         job.run()
 
-    async def _do_upload_async(self, event: RequestUpload) -> None:
-        """Upload a file and returns when it's done. Requires `event.save_fn_async`."""
-        assert event.save_fn_async is not None
-        job = upload_job.UploadJobAsync(
-            stats=self._stats,
-            api=self._api,
-            file_stream=self._file_stream,
-            silent=self.silent,
-            request=event,
-            save_fn_async=event.save_fn_async,
-        )
-        await job.run()
-
     def _init_artifact(self, artifact_id: str) -> None:
         self._artifacts[artifact_id] = {
             "finalize": False,
@@ -385,8 +285,6 @@ class StepUpload:
 
     def start(self) -> None:
         self._thread.start()
-        if self._async_executor:
-            self._async_executor.start()
 
     def is_alive(self) -> bool:
         return self._thread.is_alive()

wandb/filesync/upload_job.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import os
 from typing import TYPE_CHECKING, Optional
@@ -141,78 +140,3 @@ class UploadJob:
 
     def progress(self, total_bytes: int) -> None:
         self._stats.update_uploaded_file(self.save_name, total_bytes)
-
-
-class UploadJobAsync:
-    """Roughly an async equivalent of UploadJob.
-
-    Important differences:
-    - `run` is a coroutine
-    - If `run()` fails, it falls back to the synchronous UploadJob
-    """
-
-    def __init__(
-        self,
-        stats: "stats.Stats",
-        api: "internal_api.Api",
-        file_stream: "file_stream.FileStreamApi",
-        silent: bool,
-        request: "step_upload.RequestUpload",
-        save_fn_async: "step_upload.SaveFnAsync",
-    ) -> None:
-        self._stats = stats
-        self._api = api
-        self._file_stream = file_stream
-        self.silent = silent
-        self._request = request
-        self._save_fn_async = save_fn_async
-
-    async def run(self) -> None:
-        try:
-            deduped = await self._save_fn_async(
-                lambda _, t: self._stats.update_uploaded_file(self._request.path, t)
-            )
-        except Exception as e:
-            # Async uploads aren't yet (2023-01) battle-tested.
-            # Fall back to the "normal" synchronous upload.
-            loop = asyncio.get_event_loop()
-            logger.exception("async upload failed", exc_info=e)
-            loop.run_in_executor(None, wandb._sentry.exception, e)
-            wandb.termwarn(
-                "Async file upload failed; falling back to sync", repeat=False
-            )
-            sync_job = UploadJob(
-                self._stats,
-                self._api,
-                self._file_stream,
-                self.silent,
-                self._request.save_name,
-                self._request.path,
-                self._request.artifact_id,
-                self._request.md5,
-                self._request.copied,
-                self._request.save_fn,
-                self._request.digest,
-            )
-
-            await loop.run_in_executor(None, sync_job.run)
-        else:
-            self._file_stream.push_success(
-                self._request.artifact_id,  # type: ignore
-                self._request.save_name,
-            )
-
-            if deduped:
-                logger.info("Skipped uploading %s", self._request.path)
-                self._stats.set_file_deduped(self._request.path)
-            else:
-                logger.info("Uploaded file %s", self._request.path)
-        finally:
-            # If we fell back to the sync impl, the file will have already been deleted.
-            # Doesn't matter, we only try to delete it if it exists.
-            if self._request.copied:
-                try:
-                    os.remove(self._request.path)
-                except OSError:
-                    # The file has already been deleted, we don't have permissions, or something else we can't fix.
-                    pass

wandb/integration/catboost/catboost.py

@@ -81,7 +81,7 @@ def _checkpoint_artifact(
 
 
 def _log_feature_importance(
-    model: Union[CatBoostClassifier, CatBoostRegressor]
+    model: Union[CatBoostClassifier, CatBoostRegressor],
 ) -> None:
     """Log feature importance with default settings."""
     if wandb.run is None:

wandb/integration/fastai/__init__.py

@@ -35,6 +35,7 @@ Examples:
         learn.fit(..., callbacks=WandbCallback(learn, ...))
     ```
 """
+
 import random
 import sys
 from pathlib import Path

wandb/integration/huggingface/resolver.py

@@ -91,8 +91,8 @@ class HuggingFacePipelineRequestResponseResolver:
 
     # TODO: This should have a dependency on PreTrainedModel. i.e. isinstance(PreTrainedModel)
     # from transformers.modeling_utils import PreTrainedModel
-    # We do not want this dependency explicity in our codebase so we make a very general assumption about
-    # the structure of the pipeline which may have unintended consequences
+    # We do not want this dependency explicitly in our codebase so we make a very general
+    # assumption about the structure of the pipeline which may have unintended consequences
     def _get_model(self, pipe) -> Optional[Any]:
         """Extracts model from the pipeline.
 

wandb/integration/keras/__init__.py

@@ -2,6 +2,7 @@
 
 Keras is a deep learning API for [`TensorFlow`](https://www.tensorflow.org/).
 """
+
 __all__ = (
     "WandbCallback",
     "WandbMetricsLogger",

wandb/integration/keras/callbacks/metrics_logger.py

@@ -43,7 +43,7 @@ class WandbMetricsLogger(callbacks.Callback):
             at the end of each epoch. If "batch", logs metrics at the end
             of each batch. If an integer, logs metrics at the end of that
             many batches. Defaults to "epoch".
-        initial_global_step: (int) Use this argument to correcly log the
+        initial_global_step: (int) Use this argument to correctly log the
             learning rate when you resume training from some `initial_epoch`,
             and a learning rate scheduler is used. This can be computed as
             `step_size * initial_step`. Defaults to 0.

wandb/integration/keras/keras.py

@@ -576,7 +576,7 @@ class WandbCallback(tf.keras.callbacks.Callback):
                     )
                     self._model_trained_since_last_eval = False
             except Exception as e:
-                wandb.termwarn("Error durring prediction logging for epoch: " + str(e))
+                wandb.termwarn("Error during prediction logging for epoch: " + str(e))
 
     def on_epoch_end(self, epoch, logs=None):
         if logs is None:
@@ -616,9 +616,9 @@ class WandbCallback(tf.keras.callbacks.Callback):
         self.current = logs.get(self.monitor)
         if self.current and self.monitor_op(self.current, self.best):
             if self.log_best_prefix:
-                wandb.run.summary[
-                    f"{self.log_best_prefix}{self.monitor}"
-                ] = self.current
+                wandb.run.summary[f"{self.log_best_prefix}{self.monitor}"] = (
+                    self.current
+                )
                 wandb.run.summary["{}{}".format(self.log_best_prefix, "epoch")] = epoch
                 if self.verbose and not self.save_model:
                     print(
@@ -937,9 +937,9 @@ class WandbCallback(tf.keras.callbacks.Callback):
         grads = self._grad_accumulator_callback.grads
         metrics = {}
         for weight, grad in zip(weights, grads):
-            metrics[
-                "gradients/" + weight.name.split(":")[0] + ".gradient"
-            ] = wandb.Histogram(grad)
+            metrics["gradients/" + weight.name.split(":")[0] + ".gradient"] = (
+                wandb.Histogram(grad)
+            )
         return metrics
 
     def _log_dataframe(self):

wandb/integration/langchain/wandb_tracer.py

@@ -14,6 +14,7 @@ integration will not break user code. The one exception to the rule is at import
 LangChain is not installed, or the symbols are not in the same place, the appropriate error
 will be raised when importing this module.
 """
+
 from packaging import version
 
 import wandb.util

wandb/integration/lightning/fabric/logger.py

@@ -401,9 +401,7 @@ class WandbLogger(Logger):
                         "*", step_metric="trainer/global_step", step_sync=True
                     )
 
-        self._experiment._label(
-            repo="lightning_fabric_logger"
-        )  # pylint: disable=protected-access
+        self._experiment._label(repo="lightning_fabric_logger")  # pylint: disable=protected-access
         with telemetry.context(run=self._experiment) as tel:
             tel.feature.lightning_fabric_logger = True
         return self._experiment

wandb/integration/metaflow/metaflow.py

@@ -36,7 +36,15 @@ try:
     import pandas as pd
 
     @typedispatch  # noqa: F811
-    def _wandb_use(name: str, data: pd.DataFrame, datasets=False, run=None, testing=False, *args, **kwargs):  # type: ignore
+    def _wandb_use(
+        name: str,
+        data: pd.DataFrame,
+        datasets=False,
+        run=None,
+        testing=False,
+        *args,
+        **kwargs,
+    ):  # type: ignore
         if testing:
             return "datasets" if datasets else None
 
@@ -74,7 +82,15 @@ try:
     import torch.nn as nn
 
     @typedispatch  # noqa: F811
-    def _wandb_use(name: str, data: nn.Module, models=False, run=None, testing=False, *args, **kwargs):  # type: ignore
+    def _wandb_use(
+        name: str,
+        data: nn.Module,
+        models=False,
+        run=None,
+        testing=False,
+        *args,
+        **kwargs,
+    ):  # type: ignore
         if testing:
             return "models" if models else None
 
@@ -111,7 +127,15 @@ try:
     from sklearn.base import BaseEstimator
 
     @typedispatch  # noqa: F811
-    def _wandb_use(name: str, data: BaseEstimator, models=False, run=None, testing=False, *args, **kwargs):  # type: ignore
+    def _wandb_use(
+        name: str,
+        data: BaseEstimator,
+        models=False,
+        run=None,
+        testing=False,
+        *args,
+        **kwargs,
+    ):  # type: ignore
         if testing:
             return "models" if models else None
 
@@ -169,7 +193,14 @@ class ArtifactProxy:
 
 
 @typedispatch  # noqa: F811
-def wandb_track(name: str, data: (dict, list, set, str, int, float, bool), run=None, testing=False, *args, **kwargs):  # type: ignore
+def wandb_track(
+    name: str,
+    data: (dict, list, set, str, int, float, bool),
+    run=None,
+    testing=False,
+    *args,
+    **kwargs,
+):  # type: ignore
     if testing:
         return "scalar"
 
@@ -222,12 +253,16 @@ def wandb_use(name: str, data, *args, **kwargs):
 
 
 @typedispatch  # noqa: F811
-def wandb_use(name: str, data: (dict, list, set, str, int, float, bool), *args, **kwargs):  # type: ignore
+def wandb_use(
+    name: str, data: (dict, list, set, str, int, float, bool), *args, **kwargs
+):  # type: ignore
     pass  # do nothing for these types
 
 
 @typedispatch  # noqa: F811
-def _wandb_use(name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs):  # type: ignore
+def _wandb_use(
+    name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs
+):  # type: ignore
     if testing:
         return "datasets" if datasets else None