wandb 0.20.1rc20250604__py3-none-any.whl → 0.21.0__py3-none-any.whl

wandb/sdk/internal/handler.py

@@ -33,16 +33,12 @@ from wandb.proto.wandb_internal_pb2 import (
     SummaryItem,
     SummaryRecord,
     SummaryRecordRequest,
-    SystemMetricSample,
-    SystemMetricsBuffer,
 )
 
 from ..interface.interface_queue import InterfaceQueue
 from ..lib import handler_util, proto_util
-from ..wandb_metadata import Metadata
 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 MetricSummary
@@ -89,7 +85,6 @@ class HandleManager:
     _stopped: Event
     _writer_q: "Queue[Record]"
     _interface: InterfaceQueue
-    _system_monitor: Optional[SystemMonitor]
     _tb_watcher: Optional[tb_watcher.TBWatcher]
     _metric_defines: Dict[str, MetricRecord]
     _metric_globs: Dict[str, MetricRecord]
@@ -119,8 +114,6 @@ class HandleManager:
         self._context_keeper = context_keeper
 
         self._tb_watcher = None
-        self._system_monitor = None
-        self._metadata: Optional[Metadata] = None
         self._step = 0
 
         self._track_time = None
@@ -178,21 +171,12 @@ class HandleManager:
     def handle_request_cancel(self, record: Record) -> None:
         self._dispatch_record(record)
 
-    def handle_request_metadata(self, record: Record) -> None:
-        logger.warning("Metadata updates are ignored when using the legacy service.")
-
     def handle_request_defer(self, record: Record) -> None:
         defer = record.request.defer
         state = defer.state
 
         logger.info(f"handle defer: {state}")
-        # only handle flush tb (sender handles the rest)
-        if state == defer.FLUSH_STATS:
-            # TODO(jhr): this could block so we dont really want to call shutdown
-            # from handler thread
-            if self._system_monitor is not None:
-                self._system_monitor.finish()
-        elif state == defer.FLUSH_TB:
+        if state == defer.FLUSH_TB:
             if self._tb_watcher:
                 # shutdown tensorboard workers so we get all metrics flushed
                 self._tb_watcher.finish()
@@ -660,6 +644,9 @@ class HandleManager:
     def handle_footer(self, record: Record) -> None:
         self._dispatch_record(record)
 
+    def handle_metadata(self, record: Record) -> None:
+        self._dispatch_record(record)
+
     def handle_request_attach(self, record: Record) -> None:
         result = proto_util._result_from_record(record)
         attach_id = record.request.attach.attach_id
@@ -689,24 +676,6 @@ class HandleManager:
         else:
             self._accumulate_time = 0
 
-        # system monitor
-        self._system_monitor = SystemMonitor(
-            self._settings,
-            self._interface,
-        )
-        if not (
-            self._settings.x_disable_stats or self._settings.x_disable_machine_info
-        ):
-            self._system_monitor.start()
-        if (
-            not (self._settings.x_disable_meta or self._settings.x_disable_machine_info)
-            and not run_start.run.resumed
-        ):
-            try:
-                self._metadata = Metadata(**self._system_monitor.probe(publish=True))
-            except Exception:
-                logger.exception("Error probing system metadata.")
-
         self._tb_watcher = tb_watcher.TBWatcher(
             self._settings, interface=self._interface, run_proto=run_start.run
         )
@@ -717,18 +686,11 @@ class HandleManager:
         self._respond_result(result)
 
     def handle_request_resume(self, record: Record) -> None:
-        if self._system_monitor is not None:
-            logger.info("starting system metrics thread")
-            self._system_monitor.start()
-
         if self._track_time is not None:
             self._accumulate_time += time.time() - self._track_time
         self._track_time = time.time()
 
     def handle_request_pause(self, record: Record) -> None:
-        if self._system_monitor is not None:
-            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
             self._track_time = None
@@ -763,36 +725,6 @@ class HandleManager:
             result.response.get_summary_response.item.append(item)
         self._respond_result(result)
 
-    def handle_request_get_system_metrics(self, record: Record) -> None:
-        result = proto_util._result_from_record(record)
-        if self._system_monitor is None:
-            return
-
-        buffer = self._system_monitor.buffer
-        for key, samples in buffer.items():
-            buff = []
-            for s in samples:
-                sms = SystemMetricSample()
-                sms.timestamp.FromMicroseconds(int(s[0] * 1e6))
-                sms.value = s[1]
-                buff.append(sms)
-
-            result.response.get_system_metrics_response.system_metrics[key].CopyFrom(
-                SystemMetricsBuffer(record=buff)
-            )
-
-        self._respond_result(result)
-
-    def handle_request_get_system_metadata(self, record: Record) -> None:
-        result = proto_util._result_from_record(record)
-        if self._system_monitor is None or self._metadata is None:
-            return
-
-        result.response.get_system_metadata_response.metadata.CopyFrom(
-            self._metadata.to_proto()
-        )
-        self._respond_result(result)
-
     def handle_tbrecord(self, record: Record) -> None:
         logger.info("handling tbrecord: %s", record)
         if self._tb_watcher:
@@ -895,8 +827,6 @@ class HandleManager:
 
     def finish(self) -> None:
         logger.info("shutting down handler")
-        if self._system_monitor is not None:
-            self._system_monitor.finish()
         if self._tb_watcher:
             self._tb_watcher.finish()
         # self._context_keeper._debug_print_orphans()

wandb/sdk/internal/internal_api.py

@@ -362,6 +362,7 @@ class Api:
         self.server_create_run_queue_supports_priority: Optional[bool] = None
         self.server_supports_template_variables: Optional[bool] = None
         self.server_push_to_run_queue_supports_priority: Optional[bool] = None
+
         self._server_features_cache: Optional[Dict[str, bool]] = None
 
     def gql(self, *args: Any, **kwargs: Any) -> Any:
@@ -4661,3 +4662,56 @@ class Api:
         success: bool = response["stopRun"].get("success")
 
         return success
+
+    @normalize_exceptions
+    def create_custom_chart(
+        self,
+        entity: str,
+        name: str,
+        display_name: str,
+        spec_type: str,
+        access: str,
+        spec: Union[str, Mapping[str, Any]],
+    ) -> Optional[Dict[str, Any]]:
+        if not isinstance(spec, str):
+            spec = json.dumps(spec)
+
+        mutation = gql(
+            """
+            mutation CreateCustomChart(
+                $entity: String!
+                $name: String!
+                $displayName: String!
+                $type: String!
+                $access: String!
+                $spec: JSONString!
+            ) {
+                createCustomChart(
+                    input: {
+                        entity: $entity
+                        name: $name
+                        displayName: $displayName
+                        type: $type
+                        access: $access
+                        spec: $spec
+                    }
+                ) {
+                    chart { id }
+                }
+            }
+            """
+        )
+
+        variable_values = {
+            "entity": entity,
+            "name": name,
+            "displayName": display_name,
+            "type": spec_type,
+            "access": access,
+            "spec": spec,
+        }
+
+        result: Optional[Dict[str, Any]] = self.gql(mutation, variable_values)[
+            "createCustomChart"
+        ]
+        return result

wandb/sdk/internal/sender.py

@@ -63,6 +63,7 @@ if TYPE_CHECKING:
         ArtifactManifest,
         ArtifactManifestEntry,
         ArtifactRecord,
+        EnvironmentRecord,
         HttpResponse,
         LocalInfo,
         Record,
@@ -212,6 +213,7 @@ class SendManager:
     _context_keeper: context.ContextKeeper
 
     _telemetry_obj: telemetry.TelemetryRecord
+    _environment_obj: "EnvironmentRecord"
     _fs: Optional["file_stream.FileStreamApi"]
     _run: Optional["RunRecord"]
     _entity: Optional[str]
@@ -268,6 +270,7 @@ class SendManager:
 
         self._start_time: int = 0
         self._telemetry_obj = telemetry.TelemetryRecord()
+        self._environment_obj = wandb_internal_pb2.EnvironmentRecord()
         self._config_metric_pbdict_list: List[Dict[int, Any]] = []
         self._metadata_summary: Dict[str, Any] = defaultdict()
         self._cached_summary: Dict[str, Any] = dict()
@@ -790,12 +793,12 @@ class SendManager:
 
     def _config_backend_dict(self) -> sender_config.BackendConfigDict:
         config = self._consolidated_config or sender_config.ConfigState()
-
         return config.to_backend_dict(
             telemetry_record=self._telemetry_obj,
             framework=self._telemetry_get_framework(),
             start_time_millis=self._start_time,
             metric_pbdicts=self._config_metric_pbdict_list,
+            environment_record=self._environment_obj,
         )
 
     def _config_save(
@@ -1379,11 +1382,11 @@ class SendManager:
             next_idx = len(self._config_metric_pbdict_list)
             self._config_metric_pbdict_list.append(md)
             self._config_metric_index_dict[metric.name] = next_idx
-        self._update_config()
+        self._debounce_config()
 
     def _update_telemetry_record(self, telemetry: telemetry.TelemetryRecord) -> None:
         self._telemetry_obj.MergeFrom(telemetry)
-        self._update_config()
+        self._debounce_config()
 
     def send_telemetry(self, record: "Record") -> None:
         self._update_telemetry_record(record.telemetry)
@@ -1417,6 +1420,23 @@ class SendManager:
         # tbrecord watching threads are handled by handler.py
         pass
 
+    def _update_environment_record(self, environment: "EnvironmentRecord") -> None:
+        self._environment_obj.MergeFrom(environment)
+        self._debounce_config()
+
+    def send_environment(self, record: "Record") -> None:
+        """Inject environment info into config and upload as a JSON file."""
+        self._update_environment_record(record.environment)
+
+        environment_json = json.dumps(proto_util.message_to_dict(self._environment_obj))
+
+        with open(
+            os.path.join(self._settings.files_dir, filenames.METADATA_FNAME), "w"
+        ) as f:
+            f.write(environment_json)
+
+        self._save_file(interface.GlobStr(filenames.METADATA_FNAME), policy="now")
+
     def send_request_link_artifact(self, record: "Record") -> None:
         if not (record.control.req_resp or record.control.mailbox_slot):
             raise ValueError(

wandb/sdk/internal/sender_config.py

@@ -79,6 +79,7 @@ class ConfigState:
         framework: Optional[str],
         start_time_millis: int,
         metric_pbdicts: Sequence[Dict[int, Any]],
+        environment_record: wandb_internal_pb2.EnvironmentRecord,
     ) -> BackendConfigDict:
         """Returns a dictionary representation expected by the backend.
 
@@ -125,6 +126,14 @@ class ConfigState:
         if metric_pbdicts:
             wandb_internal["m"] = metric_pbdicts
 
+        ###################################################
+        # Environment
+        ###################################################
+        writer_id = environment_record.writer_id
+        if writer_id:
+            environment_dict = proto_util.message_to_dict(environment_record)
+            wandb_internal["e"] = {writer_id: environment_dict}
+
         return BackendConfigDict(
             {
                 key: {

wandb/sdk/launch/_project_spec.py

@@ -370,9 +370,9 @@ class LaunchProject:
 
     def set_job_entry_point(self, command: List[str]) -> "EntryPoint":
         """Set job entrypoint for the project."""
-        assert (
-            self._entry_point is None
-        ), "Cannot set entry point twice. Use LaunchProject.override_entrypoint"
+        assert self._entry_point is None, (
+            "Cannot set entry point twice. Use LaunchProject.override_entrypoint"
+        )
         new_entrypoint = EntryPoint(name=command[-1], command=command)
         self._entry_point = new_entrypoint
         return new_entrypoint

wandb/sdk/launch/agent/agent.py

@@ -76,9 +76,9 @@ class JobSpecAndQueue:
 def _convert_access(access: str) -> str:
     """Convert access string to a value accepted by wandb."""
     access = access.upper()
-    assert (
-        access == "PROJECT" or access == "USER"
-    ), "Queue access must be either project or user"
+    assert access == "PROJECT" or access == "USER", (
+        "Queue access must be either project or user"
+    )
     return access
 
 

wandb/sdk/launch/agent/job_status_tracker.py

@@ -44,7 +44,9 @@ class JobAndRunStatusTracker:
             self.run_id is not None
             and self.project is not None
             and self.entity is not None
-        ), "Job tracker does not contain run info. Update with run info before checking if run stopped"
+        ), (
+            "Job tracker does not contain run info. Update with run info before checking if run stopped"
+        )
         check_stop = event_loop_thread_exec(api.api.check_stop_requested)
         try:
             return bool(await check_stop(self.project, self.entity, self.run_id))

wandb/sdk/launch/utils.py

@@ -380,9 +380,9 @@ def diff_pip_requirements(req_1: List[str], req_2: List[str]) -> Dict[str, str]:
             else:
                 raise ValueError(f"Unable to parse pip requirements file line: {line}")
             if _name is not None:
-                assert re.match(
-                    _VALID_PIP_PACKAGE_REGEX, _name
-                ), f"Invalid pip package name {_name}"
+                assert re.match(_VALID_PIP_PACKAGE_REGEX, _name), (
+                    f"Invalid pip package name {_name}"
+                )
                 d[_name] = _version
         return d
 

wandb/sdk/lib/console_capture.py

@@ -25,17 +25,38 @@ In particular, it does not work with some combinations of pytest's
 
 from __future__ import annotations
 
+import logging
 import sys
 import threading
 from typing import IO, AnyStr, Callable, Protocol
 
+from . import wb_logging
+
+_logger = logging.getLogger(__name__)
+
 
 class CannotCaptureConsoleError(Exception):
     """The module failed to patch stdout or stderr."""
 
 
 class _WriteCallback(Protocol):
-    """A callback that receives intercepted bytes or string data."""
+    """A callback that receives intercepted bytes or string data.
+
+    This may be called from any thread, but is only called from one thread
+    at a time.
+
+    Note on errors: Any error raised during the callback will clear all
+    callbacks. This means that if a user presses Ctrl-C at an unlucky time
+    during a run, we will stop uploading console output---but it's not
+    likely to be a problem unless something catches the KeyboardInterrupt.
+
+    Regular Exceptions are caught and logged instead of bubbling up to the
+    user's print() statements; other exceptions like KeyboardInterrupt are
+    re-raised.
+
+    Callbacks should handle all exceptions---a callback that raises any
+    Exception is considered buggy.
+    """
 
     def __call__(
         self,
@@ -45,6 +66,8 @@ class _WriteCallback(Protocol):
     ) -> None:
         """Intercept data passed to `write()`.
 
+        See the protocol docstring for information about exceptions.
+
         Args:
             data: The object passed to stderr's or stdout's `write()`.
             written: The number of bytes or characters written.
@@ -52,7 +75,9 @@ class _WriteCallback(Protocol):
         """
 
 
-_module_lock = threading.Lock()
+# A reentrant lock is used to catch callbacks that write to stderr/stdout.
+_module_rlock = threading.RLock()
+_is_writing = False
 
 _patch_exception: CannotCaptureConsoleError | None = None
 
@@ -67,9 +92,6 @@ def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
 
     Args:
         callback: A callback to invoke after running `sys.stdout.write`.
-            This may be called from any thread, so it must be thread-safe.
-            Exceptions are propagated to the caller of `write`.
-            See `_WriteCallback` for the exact protocol.
 
     Returns:
         A function to uninstall the callback.
@@ -77,7 +99,7 @@ def capture_stdout(callback: _WriteCallback) -> Callable[[], None]:
     Raises:
         CannotCaptureConsoleError: If patching failed on import.
     """
-    with _module_lock:
+    with _module_rlock:
         if _patch_exception:
             raise _patch_exception
 
@@ -92,9 +114,6 @@ def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
 
     Args:
         callback: A callback to invoke after running `sys.stderr.write`.
-            This may be called from any thread, so it must be thread-safe.
-            Exceptions are propagated to the caller of `write`.
-            See `_WriteCallback` for the exact protocol.
 
     Returns:
         A function to uninstall the callback.
@@ -102,7 +121,7 @@ def capture_stderr(callback: _WriteCallback) -> Callable[[], None]:
     Raises:
         CannotCaptureConsoleError: If patching failed on import.
     """
-    with _module_lock:
+    with _module_rlock:
         if _patch_exception:
             raise _patch_exception
 
@@ -125,11 +144,11 @@ def _insert_disposably(
     def dispose() -> None:
         nonlocal disposed
 
-        with _module_lock:
+        with _module_rlock:
             if disposed:
                 return
 
-            del callback_dict[id]
+            callback_dict.pop(id, None)
 
             disposed = True
 
@@ -143,16 +162,44 @@ def _patch(
 ) -> None:
     orig_write: Callable[[AnyStr], int]
 
+    @wb_logging.log_to_all_runs()
     def write_with_callbacks(s: AnyStr, /) -> int:
+        global _is_writing
         n = orig_write(s)
 
-        # We make a copy here because callbacks could, in theory, modify
-        # the list of callbacks.
-        with _module_lock:
-            callbacks_copy = list(callbacks.values())
-
-        for cb in callbacks_copy:
-            cb(s, n)
+        # NOTE: Since _module_rlock is reentrant, this is safe. It will not
+        # deadlock if a callback invokes write() again.
+        with _module_rlock:
+            if _is_writing:
+                return n
+
+            _is_writing = True
+            try:
+                for cb in callbacks.values():
+                    cb(s, n)
+
+            except BaseException as e:
+                # Clear all callbacks on any exception to avoid infinite loops:
+                #
+                # * If we re-raise, an exception handler is likely to print
+                #   the exception to the console and trigger callbacks again
+                # * If we log, we can't guarantee that this doesn't print
+                #   to console.
+                #
+                # This is especially important for KeyboardInterrupt.
+                _stderr_callbacks.clear()
+                _stdout_callbacks.clear()
+
+                if isinstance(e, Exception):
+                    # We suppress Exceptions so that bugs in W&B code don't
+                    # cause the user's print() statements to raise errors.
+                    _logger.exception("Error in console callback, clearing all!")
+                else:
+                    # Re-raise errors like KeyboardInterrupt.
+                    raise
+
+            finally:
+                _is_writing = False
 
         return n
 

wandb/sdk/lib/printer.py

@@ -349,13 +349,13 @@ class _PrinterTerm(Printer):
         text = text or " " * 79
         wandb.termlog(text)
 
-    @override
     @property
+    @override
     def supports_html(self) -> bool:
         return False
 
-    @override
     @property
+    @override
     def supports_unicode(self) -> bool:
         return wandb.util.is_unicode_safe(sys.stderr)
 
@@ -464,11 +464,10 @@ class _PrinterJupyter(Printer):
 
         if handle:
             yield _DynamicJupyterText(handle)
+            handle.update(self._ipython_display.HTML(""))
         else:
             yield None
 
-        handle.update(self._ipython_display.HTML(""))
-
     @override
     def display(
         self,
@@ -483,13 +482,13 @@ class _PrinterJupyter(Printer):
         text = "<br>".join(text.splitlines())
         self._ipython_display.display(self._ipython_display.HTML(text))
 
-    @override
     @property
+    @override
     def supports_html(self) -> bool:
         return True
 
-    @override
     @property
+    @override
     def supports_unicode(self) -> bool:
         return True
 
@@ -540,7 +539,7 @@ class _PrinterJupyter(Printer):
         self._progress.update(percent_done, text)
 
     @override
-    def progress_close(self, _: str | None = None) -> None:
+    def progress_close(self, text: str | None = None) -> None:
         if self._progress:
             self._progress.close()
 

wandb/sdk/lib/progress.py

@@ -7,7 +7,6 @@ import contextlib
 import time
 from typing import Iterable, Iterator, NoReturn
 
-from wandb import env
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.sdk.interface import interface
 from wandb.sdk.lib import asyncio_compat
@@ -107,8 +106,7 @@ class ProgressPrinter:
         progress_text_area: p.DynamicText | None,
         default_text: str,
     ) -> None:
-        # Not implemented by the legacy service.
-        self._show_operation_stats = not env.is_require_legacy_service()
+        self._show_operation_stats = True
         self._printer = printer
         self._progress_text_area = progress_text_area
         self._default_text = default_text

wandb/sdk/lib/service/ipc_support.py

@@ -0,0 +1,13 @@
+"""Constants determining what IPC methods are supported."""
+
+import socket
+
+SUPPORTS_UNIX = hasattr(socket, "AF_UNIX")
+"""Whether Unix sockets are supported.
+
+AF_UNIX is not supported on Windows:
+https://github.com/python/cpython/issues/77589
+
+Windows has supported Unix sockets since ~2017, but support in Python is
+missing as of 2025.
+"""