wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/sdk/internal/file_pusher.py

@@ -3,6 +3,7 @@ import logging
 import os
 import queue
 import tempfile
+import threading
 import time
 from typing import TYPE_CHECKING, Optional, Tuple
 
@@ -12,8 +13,8 @@ from wandb.filesync import stats, step_checksum, step_upload
 from wandb.sdk.lib.paths import LogicalPath
 
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts import artifact_saver
     from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+    from wandb.sdk.artifacts.artifact_saver import SaveFn, SaveFnAsync
     from wandb.sdk.internal import file_stream, internal_api
     from wandb.sdk.internal.settings_static import SettingsStatic
 
@@ -72,6 +73,21 @@ class FilePusher:
         )
         self._step_upload.start()
 
+        self._stats_thread_stop = threading.Event()
+        if os.environ.get("WANDB_DEBUG"):
+            # debug thread to monitor and report file pusher stats
+            self._stats_thread = threading.Thread(
+                target=self._file_pusher_stats,
+                daemon=True,
+                name="FPStatsThread",
+            )
+            self._stats_thread.start()
+
+    def _file_pusher_stats(self) -> None:
+        while not self._stats_thread_stop.is_set():
+            logger.info(f"FilePusher stats: {self._stats._stats}")
+            time.sleep(1)
+
     def get_status(self) -> Tuple[bool, stats.Summary]:
         running = self.is_alive()
         summary = self._stats.summary()
@@ -134,8 +150,8 @@ class FilePusher:
         self,
         manifest: "ArtifactManifest",
         artifact_id: str,
-        save_fn: "artifact_saver.SaveFn",
-        save_fn_async: "artifact_saver.SaveFnAsync",
+        save_fn: "SaveFn",
+        save_fn_async: "SaveFnAsync",
     ) -> None:
         event = step_checksum.RequestStoreManifestFiles(
             manifest, artifact_id, save_fn, save_fn_async
@@ -158,6 +174,7 @@ class FilePusher:
     def finish(self, callback: Optional[step_upload.OnRequestFinishFn] = None):
         logger.info("shutting down file pusher")
         self._incoming_queue.put(step_checksum.RequestFinish(callback))
+        self._stats_thread_stop.set()
 
     def join(self) -> None:
         # NOTE: must have called finish before join

wandb/sdk/internal/file_stream.py

@@ -335,11 +335,13 @@ class FileStreamApi:
         self._run_id = run_id
         self._start_time = start_time
         self._client = requests.Session()
+        timeout = timeout or 0
         if timeout > 0:
             self._client.post = functools.partial(self._client.post, timeout=timeout)  # type: ignore[method-assign]
         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._client.proxies.update(api.client.transport.session.proxies or {})
         self._file_policies: Dict[str, DefaultFilePolicy] = {}
         self._dropped_chunks: int = 0
         self._queue: queue.Queue = queue.Queue()
@@ -499,7 +501,7 @@ class FileStreamApi:
             wandb.termerror(
                 "Dropped streaming file chunk (see wandb/debug-internal.log)"
             )
-            logging.exception("dropped chunk %s" % response)
+            logger.exception("dropped chunk %s" % response)
             self._dropped_chunks += 1
         else:
             parsed: Optional[dict] = None

wandb/sdk/internal/handler.py

@@ -23,6 +23,7 @@ from typing import (
 
 from wandb.proto.wandb_internal_pb2 import (
     HistoryRecord,
+    InternalMessages,
     MetricRecord,
     Record,
     Result,
@@ -34,13 +35,13 @@ from wandb.proto.wandb_internal_pb2 import (
 )
 
 from ..interface.interface_queue import InterfaceQueue
-from ..lib import handler_util, proto_util, tracelog
+from ..lib import handler_util, proto_util, tracelog, wburls
 from . import context, sample, tb_watcher
 from .settings_static import SettingsStatic
 from .system.system_monitor import SystemMonitor
 
 if TYPE_CHECKING:
-    from wandb.proto.wandb_internal_pb2 import ArtifactDoneRequest, MetricSummary
+    from wandb.proto.wandb_internal_pb2 import MetricSummary
 
 
 SummaryDict = Dict[str, Any]
@@ -80,7 +81,6 @@ class HandleManager:
     _metric_copy: Dict[Tuple[str, ...], Any]
     _track_time: Optional[float]
     _accumulate_time: float
-    _artifact_xid_done: Dict[str, "ArtifactDoneRequest"]
     _run_start_time: Optional[float]
     _context_keeper: context.ContextKeeper
 
@@ -119,9 +119,9 @@ class HandleManager:
         self._metric_globs = defaultdict(MetricRecord)
         self._metric_track = dict()
         self._metric_copy = dict()
+        self._internal_messages = InternalMessages()
 
-        # TODO: implement release protocol to clean this up
-        self._artifact_xid_done = dict()
+        self._dropped_history = False
 
     def __len__(self) -> int:
         return self._record_q.qsize()
@@ -235,7 +235,7 @@ class HandleManager:
             record = Record(summary=summary)
             self._dispatch_record(record)
         elif not self._settings._offline:
-            # Send this summary update as a request since we arent persisting every update
+            # Send this summary update as a request since we aren't persisting every update
             summary_record = SummaryRecordRequest(summary=summary)
             request_record = self._interface._make_request(
                 summary_record=summary_record
@@ -266,7 +266,7 @@ class HandleManager:
         if s.none:
             return False
         if s.copy:
-            # non key list copy already done in _update_summary
+            # non-key list copy already done in _update_summary
             if len(kl) > 1:
                 _dict_nested_set(self._consolidated_summary, kl, v)
                 return True
@@ -290,7 +290,7 @@ class HandleManager:
                 if best_key:
                     _dict_nested_set(self._consolidated_summary, best_key, v)
                     updated = True
-        # defaulting to minimize if goal is not supecified
+        # defaulting to minimize if goal is not specified
         if s.min or best_key and not goal_max:
             min_key = tuple(kl + ["min"])
             old_min = self._metric_track.get(min_key)
@@ -358,7 +358,7 @@ class HandleManager:
     ) -> bool:
         metric_key = ".".join([k.replace(".", "\\.") for k in kl])
         d = self._metric_defines.get(metric_key, d)
-        # if the dict has _type key, its a wandb table object
+        # if the dict has _type key, it's a wandb table object
         if isinstance(v, dict) and not handler_util.metric_is_wandb_dict(v):
             updated = False
             for nk, nv in v.items():
@@ -374,7 +374,7 @@ class HandleManager:
         return updated
 
     def _update_summary_media_objects(self, v: Dict[str, Any]) -> Dict[str, Any]:
-        # For now, non recursive - just top level
+        # For now, non-recursive - just top level
         for nk, nv in v.items():
             if (
                 isinstance(nv, dict)
@@ -387,17 +387,17 @@ class HandleManager:
                 v[nk] = nv
         return v
 
-    def _update_summary(self, history_dict: Dict[str, Any]) -> bool:
+    def _update_summary(self, history_dict: Dict[str, Any]) -> List[str]:
         # keep old behavior fast path if no define metrics have been used
         if not self._metric_defines:
             history_dict = self._update_summary_media_objects(history_dict)
             self._consolidated_summary.update(history_dict)
-            return True
-        updated = False
+            return list(history_dict.keys())
+        updated_keys = []
         for k, v in history_dict.items():
             if self._update_summary_list(kl=[k], v=v):
-                updated = True
-        return updated
+                updated_keys.append(k)
+        return updated_keys
 
     def _history_assign_step(
         self,
@@ -511,24 +511,28 @@ class HandleManager:
         self._history_update(record.history, history_dict)
         self._dispatch_record(record)
         self._save_history(record.history)
-        updated = self._update_summary(history_dict)
-        if updated:
-            self._save_summary(self._consolidated_summary)
+        # update summary from history
+        updated_keys = self._update_summary(history_dict)
+        if updated_keys:
+            updated_items = {k: self._consolidated_summary[k] for k in updated_keys}
+            self._save_summary(updated_items)
 
     def _flush_partial_history(
         self,
         step: Optional[int] = None,
     ) -> None:
-        if self._partial_history:
-            history = HistoryRecord()
-            for k, v in self._partial_history.items():
-                item = history.item.add()
-                item.key = k
-                item.value_json = json.dumps(v)
-            if step is not None:
-                history.step.num = step
-            self.handle_history(Record(history=history))
-            self._partial_history = {}
+        if not self._partial_history:
+            return
+
+        history = HistoryRecord()
+        for k, v in self._partial_history.items():
+            item = history.item.add()
+            item.key = k
+            item.value_json = json.dumps(v)
+        if step is not None:
+            history.step.num = step
+        self.handle_history(Record(history=history))
+        self._partial_history = {}
 
     def handle_request_sender_mark_report(self, record: Record) -> None:
         self._dispatch_record(record, always_send=True)
@@ -550,9 +554,18 @@ class HandleManager:
         history_dict = proto_util.dict_from_proto_list(partial_history.item)
         if step is not None:
             if step < self._step:
-                logger.warning(
-                    f"Step {step} < {self._step}. Dropping entry: {history_dict}."
+                if not self._dropped_history:
+                    message = (
+                        "Step only supports monotonically increasing values, use define_metric to set a custom x "
+                        f"axis. For details see: {wburls.wburls.get('wandb_define_metric')}"
+                    )
+                    self._internal_messages.warning.append(message)
+                    self._dropped_history = True
+                message = (
+                    f"(User provided step: {step} is less than current step: {self._step}. "
+                    f"Dropping entry: {history_dict})."
                 )
+                self._internal_messages.warning.append(message)
                 return
             elif step > self._step:
                 self._flush_partial_history()
@@ -639,44 +652,6 @@ class HandleManager:
     def handle_request_log_artifact(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_request_artifact_send(self, record: Record) -> None:
-        assert record.control.req_resp
-        result = proto_util._result_from_record(record)
-
-        self._dispatch_record(record)
-
-        # send response immediately, the request will be polled for result
-        xid = record.uuid
-        result.response.artifact_send_response.xid = xid
-        self._respond_result(result)
-
-    def handle_request_artifact_poll(self, record: Record) -> None:
-        assert record.control.req_resp
-        xid = record.request.artifact_poll.xid
-        assert xid
-
-        result = proto_util._result_from_record(record)
-        done_req = self._artifact_xid_done.get(xid)
-        if done_req:
-            result.response.artifact_poll_response.artifact_id = done_req.artifact_id
-            result.response.artifact_poll_response.error_message = (
-                done_req.error_message
-            )
-            result.response.artifact_poll_response.ready = True
-        self._respond_result(result)
-
-    def handle_request_artifact_done(self, record: Record) -> None:
-        assert not record.control.req_resp
-        done_req = record.request.artifact_done
-        xid = done_req.xid
-        assert xid
-
-        self._artifact_xid_done[xid] = done_req
-
-    # def handle_request_artifact_release(self, record: Record) -> None:
-    #     assert record.control.req_resp
-    #     # TODO: implement release protocol to clean up _artifact_xid_done dict
-
     def handle_telemetry(self, record: Record) -> None:
         self._dispatch_record(record)
 
@@ -716,7 +691,7 @@ class HandleManager:
 
     def handle_request_resume(self, record: Record) -> None:
         if self._system_monitor is not None:
-            logger.info("starting system metrics thread or process")
+            logger.info("starting system metrics thread")
             self._system_monitor.start()
 
         if self._track_time is not None:
@@ -725,7 +700,7 @@ class HandleManager:
 
     def handle_request_pause(self, record: Record) -> None:
         if self._system_monitor is not None:
-            logger.info("stopping system metrics thread or process")
+            logger.info("stopping system metrics thread")
             self._system_monitor.finish()
         if self._track_time is not None:
             self._accumulate_time += time.time() - self._track_time
@@ -740,6 +715,14 @@ class HandleManager:
     def handle_request_network_status(self, record: Record) -> None:
         self._dispatch_record(record)
 
+    def handle_request_internal_messages(self, record: Record) -> None:
+        result = proto_util._result_from_record(record)
+        result.response.internal_messages_response.messages.CopyFrom(
+            self._internal_messages
+        )
+        self._internal_messages.Clear()
+        self._respond_result(result)
+
     def handle_request_status(self, record: Record) -> None:
         # TODO(mempressure): do something better?
         assert record.control.req_resp