wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/integration_utils/auto_logging.py

@@ -0,0 +1,215 @@
+import functools
+import logging
+import sys
+from typing import Any, Dict, Optional, Sequence, TypeVar
+
+import wandb.sdk
+import wandb.util
+from wandb.sdk.lib import telemetry as wb_telemetry
+from wandb.sdk.lib.timer import Timer
+
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+
+logger = logging.getLogger(__name__)
+
+
+AutologInitArgs = Optional[Dict[str, Any]]
+
+
+K = TypeVar("K", bound=str)
+V = TypeVar("V")
+
+
+class Response(Protocol[K, V]):
+    def __getitem__(self, key: K) -> V:
+        ...  # pragma: no cover
+
+    def get(self, key: K, default: Optional[V] = None) -> Optional[V]:
+        ...  # pragma: no cover
+
+
+class ArgumentResponseResolver(Protocol):
+    def __call__(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        response: Response,
+        start_time: float,
+        time_elapsed: float,
+    ) -> Optional[Dict[str, Any]]:
+        ...  # pragma: no cover
+
+
+class PatchAPI:
+    def __init__(
+        self,
+        name: str,
+        symbols: Sequence[str],
+        resolver: ArgumentResponseResolver,
+    ) -> None:
+        """Patches the API to log wandb Media or metrics."""
+        # name of the LLM provider, e.g. "Cohere" or "OpenAI" or package name like "Transformers"
+        self.name = name
+        # api library name, e.g. "cohere" or "openai" or "transformers"
+        self._api = None
+        # dictionary of original methods
+        self.original_methods: Dict[str, Any] = {}
+        # list of symbols to patch, e.g. ["Client.generate", "Edit.create"] or ["Pipeline.__call__"]
+        self.symbols = symbols
+        # resolver callable to convert args/response into a dictionary of wandb media objects or metrics
+        self.resolver = resolver
+
+    @property
+    def set_api(self) -> Any:
+        """Returns the API module."""
+        lib_name = self.name.lower()
+        if self._api is None:
+            self._api = wandb.util.get_module(
+                name=lib_name,
+                required=f"To use the W&B {self.name} Autolog, "
+                f"you need to have the `{lib_name}` python "
+                f"package installed. Please install it with `pip install {lib_name}`.",
+                lazy=False,
+            )
+        return self._api
+
+    def patch(self, run: "wandb.sdk.wandb_run.Run") -> None:
+        """Patches the API to log media or metrics to W&B."""
+        for symbol in self.symbols:
+            # split on dots, e.g. "Client.generate" -> ["Client", "generate"]
+            symbol_parts = symbol.split(".")
+            # and get the attribute from the module
+            original = functools.reduce(getattr, symbol_parts, self.set_api)
+
+            def method_factory(original_method: Any):
+                @functools.wraps(original_method)
+                def method(*args, **kwargs):
+                    with Timer() as timer:
+                        result = original_method(*args, **kwargs)
+                    try:
+                        loggable_dict = self.resolver(
+                            args, kwargs, result, timer.start_time, timer.elapsed
+                        )
+                        if loggable_dict is not None:
+                            run.log(loggable_dict)
+                    except Exception as e:
+                        logger.warning(e)
+                    return result
+
+                return method
+
+            # save original method
+            self.original_methods[symbol] = original
+            # monkey patch the method
+            if len(symbol_parts) == 1:
+                setattr(self.set_api, symbol_parts[0], method_factory(original))
+            else:
+                setattr(
+                    functools.reduce(getattr, symbol_parts[:-1], self.set_api),
+                    symbol_parts[-1],
+                    method_factory(original),
+                )
+
+    def unpatch(self) -> None:
+        """Unpatches the API."""
+        for symbol, original in self.original_methods.items():
+            # split on dots, e.g. "Client.generate" -> ["Client", "generate"]
+            symbol_parts = symbol.split(".")
+            # unpatch the method
+            if len(symbol_parts) == 1:
+                setattr(self.set_api, symbol_parts[0], original)
+            else:
+                setattr(
+                    functools.reduce(getattr, symbol_parts[:-1], self.set_api),
+                    symbol_parts[-1],
+                    original,
+                )
+
+
+class AutologAPI:
+    def __init__(
+        self,
+        name: str,
+        symbols: Sequence[str],
+        resolver: ArgumentResponseResolver,
+        telemetry_feature: Optional[str] = None,
+    ) -> None:
+        """Autolog API calls to W&B."""
+        self._telemetry_feature = telemetry_feature
+        self._patch_api = PatchAPI(
+            name=name,
+            symbols=symbols,
+            resolver=resolver,
+        )
+        self._name = self._patch_api.name
+        self._run: Optional["wandb.sdk.wandb_run.Run"] = None
+        self.__run_created_by_autolog: bool = False
+
+    @property
+    def _is_enabled(self) -> bool:
+        """Returns whether autologging is enabled."""
+        return self._run is not None
+
+    def __call__(self, init: AutologInitArgs = None) -> None:
+        """Enable autologging."""
+        self.enable(init=init)
+
+    def _run_init(self, init: AutologInitArgs = None) -> None:
+        """Handle wandb run initialization."""
+        # - autolog(init: dict = {...}) calls wandb.init(**{...})
+        #   regardless of whether there is a wandb.run or not,
+        #   we only track if the run was created by autolog
+        #    - todo: autolog(init: dict | run = run) would use the user-provided run
+        # - autolog() uses the wandb.run if there is one, otherwise it calls wandb.init()
+        if init:
+            _wandb_run = wandb.run
+            # we delegate dealing with the init dict to wandb.init()
+            self._run = wandb.init(**init)
+            if _wandb_run != self._run:
+                self.__run_created_by_autolog = True
+        elif wandb.run is None:
+            self._run = wandb.init()
+            self.__run_created_by_autolog = True
+        else:
+            self._run = wandb.run
+
+    def enable(self, init: AutologInitArgs = None) -> None:
+        """Enable autologging.
+
+        Args:
+            init: Optional dictionary of arguments to pass to wandb.init().
+
+        """
+        if self._is_enabled:
+            logger.info(
+                f"{self._name} autologging is already enabled, disabling and re-enabling."
+            )
+            self.disable()
+
+        logger.info(f"Enabling {self._name} autologging.")
+        self._run_init(init=init)
+
+        self._patch_api.patch(self._run)
+
+        if self._telemetry_feature:
+            with wb_telemetry.context(self._run) as tel:
+                setattr(tel.feature, self._telemetry_feature, True)
+
+    def disable(self) -> None:
+        """Disable autologging."""
+        if self._run is None:
+            return
+
+        logger.info(f"Disabling {self._name} autologging.")
+
+        if self.__run_created_by_autolog:
+            self._run.finish()
+            self.__run_created_by_autolog = False
+
+        self._run = None
+
+        self._patch_api.unpatch()

wandb/sdk/interface/interface.py

@@ -17,9 +17,9 @@ import time
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Iterable, NewType, Optional, Tuple, Union
 
-from wandb.apis.public import Artifact as PublicArtifact
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
+from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
 from wandb.util import (
     WandBJSONEncoderOld,
     get_h5_typename,
@@ -32,14 +32,14 @@ from wandb.util import (
 
 from ..data_types.utils import history_dict_to_json, val_to_json
 from ..lib.mailbox import MailboxHandle
-from ..wandb_artifacts import Artifact
 from . import summary_record as sr
-from .artifacts import ArtifactManifest
 from .message_future import MessageFuture
 
 GlobStr = NewType("GlobStr", str)
 
 if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
     from ..wandb_run import Run
 
     if sys.version_info >= (3, 8):
@@ -180,8 +180,9 @@ class InterfaceBase:
             proto_run.telemetry.MergeFrom(run._telemetry_obj)
         return proto_run
 
-    def publish_run(self, run: "pb.RunRecord") -> None:
-        self._publish_run(run)
+    def publish_run(self, run: "Run") -> None:
+        run_record = self._make_run(run)
+        self._publish_run(run_record)
 
     @abstractmethod
     def _publish_run(self, run: pb.RunRecord) -> None:
@@ -372,7 +373,7 @@ class InterfaceBase:
     def _publish_files(self, files: pb.FilesRecord) -> None:
         raise NotImplementedError
 
-    def _make_artifact(self, artifact: Artifact) -> pb.ArtifactRecord:
+    def _make_artifact(self, artifact: "Artifact") -> pb.ArtifactRecord:
         proto_artifact = pb.ArtifactRecord()
         proto_artifact.type = artifact.type
         proto_artifact.name = artifact.name
@@ -424,14 +425,14 @@ class InterfaceBase:
     def publish_link_artifact(
         self,
         run: "Run",
-        artifact: Union[Artifact, PublicArtifact],
+        artifact: "Artifact",
         portfolio_name: str,
         aliases: Iterable[str],
         entity: Optional[str] = None,
         project: Optional[str] = None,
     ) -> None:
         link_artifact = pb.LinkArtifactRecord()
-        if isinstance(artifact, Artifact):
+        if artifact.is_draft():
             link_artifact.client_id = artifact._client_id
         else:
             link_artifact.server_id = artifact.id if artifact.id else ""
@@ -448,10 +449,8 @@ class InterfaceBase:
 
     def publish_use_artifact(
         self,
-        artifact: Artifact,
+        artifact: "Artifact",
     ) -> None:
-        # use_artifact is either a public.Artifact or a wandb.Artifact that has been
-        # waited on and has an id
         assert artifact.id is not None, "Artifact must have an id"
         use_artifact = pb.UseArtifactRecord(
             id=artifact.id, type=artifact.type, name=artifact.name
@@ -466,7 +465,7 @@ class InterfaceBase:
     def communicate_artifact(
         self,
         run: "Run",
-        artifact: Artifact,
+        artifact: "Artifact",
         aliases: Iterable[str],
         history_step: Optional[int] = None,
         is_user_created: bool = False,
@@ -516,7 +515,7 @@ class InterfaceBase:
     def publish_artifact(
         self,
         run: "Run",
-        artifact: Artifact,
+        artifact: "Artifact",
         aliases: Iterable[str],
         is_user_created: bool = False,
         use_after_commit: bool = False,
@@ -744,8 +743,9 @@ class InterfaceBase:
     def _communicate_shutdown(self) -> None:
         raise NotImplementedError
 
-    def deliver_run(self, run: "pb.RunRecord") -> MailboxHandle:
-        return self._deliver_run(run)
+    def deliver_run(self, run: "Run") -> MailboxHandle:
+        run_record = self._make_run(run)
+        return self._deliver_run(run_record)
 
     @abstractmethod
     def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle:

wandb/sdk/internal/file_pusher.py

@@ -8,11 +8,13 @@ from typing import TYPE_CHECKING, Optional, Tuple
 
 import wandb
 import wandb.util
-from wandb.filesync import dir_watcher, stats, step_checksum, step_upload
+from wandb.filesync import stats, step_checksum, step_upload
+from wandb.sdk.lib.paths import LogicalPath
 
 if TYPE_CHECKING:
-    from wandb.sdk.interface import artifacts
-    from wandb.sdk.internal import artifact_saver, file_stream, internal_api
+    from wandb.sdk.artifacts import artifact_saver
+    from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+    from wandb.sdk.internal import file_stream, internal_api
     from wandb.sdk.internal.settings_static import SettingsStatic
 
 
@@ -111,12 +113,7 @@ class FilePusher:
     def file_counts_by_category(self) -> stats.FileCountsByCategory:
         return self._stats.file_counts_by_category()
 
-    def file_changed(
-        self,
-        save_name: dir_watcher.SaveName,
-        path: str,
-        copy: bool = True,
-    ):
+    def file_changed(self, save_name: LogicalPath, path: str, copy: bool = True):
         """Tell the file pusher that a file's changed and should be uploaded.
 
         Arguments:
@@ -130,17 +127,12 @@ class FilePusher:
         if os.path.getsize(path) == 0:
             return
 
-        save_name = dir_watcher.SaveName(wandb.util.to_forward_slash_path(save_name))
-        event = step_checksum.RequestUpload(
-            path,
-            dir_watcher.SaveName(save_name),
-            copy,
-        )
+        event = step_checksum.RequestUpload(path, save_name, copy)
         self._incoming_queue.put(event)
 
     def store_manifest_files(
         self,
-        manifest: "artifacts.ArtifactManifest",
+        manifest: "ArtifactManifest",
         artifact_id: str,
         save_fn: "artifact_saver.SaveFn",
         save_fn_async: "artifact_saver.SaveFnAsync",

wandb/sdk/internal/file_stream.py

@@ -336,14 +336,9 @@ class FileStreamApi:
         self._client = requests.Session()
         # todo: actually use the timeout once more thorough error injection in testing covers it
         # self._client.post = functools.partial(self._client.post, timeout=self.HTTP_TIMEOUT)
-        self._client.auth = ("api", api.api_key or "")
-        self._client.headers.update(
-            {
-                "User-Agent": api.user_agent,
-                "X-WANDB-USERNAME": env.get_username() or "",
-                "X-WANDB-USER-EMAIL": env.get_user_email() or "",
-            }
-        )
+        self._client.auth = api.client.transport.session.auth
+        self._client.headers.update(api.client.transport.headers or {})
+        self._client.cookies.update(api.client.transport.cookies or {})  # type: ignore[no-untyped-call]
         self._file_policies: Dict[str, "DefaultFilePolicy"] = {}
         self._dropped_chunks: int = 0
         self._queue: queue.Queue = queue.Queue()
@@ -658,9 +653,8 @@ def request_with_retry(
                 e.response is not None and e.response.status_code == 429
             ):
                 err_str = (
-                    "Filestream rate limit exceeded, retrying in {} seconds".format(
-                        delay
-                    )
+                    "Filestream rate limit exceeded, "
+                    f"retrying in {delay:.1f} seconds. "
                 )
                 if retry_callback:
                     retry_callback(e.response.status_code, err_str)

wandb/sdk/internal/handler.py

@@ -191,6 +191,11 @@ class HandleManager:
         self._dispatch_record(record)
 
     def handle_run(self, record: Record) -> None:
+        if self._settings._offline:
+            self._run_proto = record.run
+            result = proto_util._result_from_record(record)
+            result.run_result.run.CopyFrom(record.run)
+            self._respond_result(result)
         self._dispatch_record(record)
 
     def handle_stats(self, record: Record) -> None:
@@ -624,7 +629,12 @@ class HandleManager:
         self._dispatch_record(record)
 
     def handle_request_attach(self, record: Record) -> None:
-        self._dispatch_record(record)
+        result = proto_util._result_from_record(record)
+        attach_id = record.request.attach.attach_id
+        assert attach_id
+        assert self._run_proto
+        result.response.attach_response.run.CopyFrom(self._run_proto)
+        self._respond_result(result)
 
     def handle_request_log_artifact(self, record: Record) -> None:
         self._dispatch_record(record)
@@ -675,6 +685,8 @@ class HandleManager:
         assert run_start
         assert run_start.run
 
+        self._run_proto = run_start.run
+
         self._run_start_time = run_start.run.start_time.ToMicroseconds() / 1e6
 
         self._track_time = time.time()