wandb 0.20.1__py3-none-win32.whl → 0.20.2rc20250616__py3-none-win32.whl

wandb/sdk/internal/system/system_monitor.py

@@ -1,224 +0,0 @@
-import datetime
-import logging
-import queue
-import threading
-from collections import defaultdict, deque
-from typing import TYPE_CHECKING, Deque, Dict, List, Optional, Tuple
-
-from .assets.asset_registry import asset_registry
-from .assets.interfaces import Asset, Interface
-from .assets.open_metrics import OpenMetrics
-from .system_info import SystemInfo
-
-if TYPE_CHECKING:
-    from wandb.proto.wandb_telemetry_pb2 import TelemetryRecord
-    from wandb.sdk.interface.interface import FilesDict
-    from wandb.sdk.internal.settings_static import SettingsStatic
-
-
-logger = logging.getLogger(__name__)
-
-
-class AssetInterface:
-    def __init__(self) -> None:
-        self.metrics_queue: queue.Queue[dict] = queue.Queue()
-        self.telemetry_queue: queue.Queue[TelemetryRecord] = queue.Queue()
-
-    def publish_stats(self, stats: dict) -> None:
-        self.metrics_queue.put(stats)
-
-    def _publish_telemetry(self, telemetry: "TelemetryRecord") -> None:
-        self.telemetry_queue.put(telemetry)
-
-    def publish_files(self, files_dict: "FilesDict") -> None:
-        pass
-
-
-class SystemMonitor:
-    # SystemMonitor is responsible for managing system metrics data.
-
-    # if joining assets, wait for publishing_interval times this many seconds
-    PUBLISHING_INTERVAL_DELAY_FACTOR = 2
-
-    def __init__(
-        self,
-        settings: "SettingsStatic",
-        interface: "Interface",
-    ) -> None:
-        self._shutdown_event: threading.Event = threading.Event()
-        self._process: Optional[threading.Thread] = None
-
-        self.settings = settings
-
-        # settings._stats_join_assets controls whether we should join stats from different assets
-        # before publishing them to the backend. If set to False, we will publish stats from each
-        # asset separately, using the backend interface. If set to True, we will aggregate stats from
-        # all assets before publishing them, using an internal queue interface, and then publish
-        # them using the interface to the backend.
-        # This is done to improve compatibility with older versions of the backend as it used to
-        # collect the names of the metrics to be displayed in the UI from the first stats message.
-
-        # compute the global publishing interval if _stats_join_assets is requested
-        sampling_interval: float = float(
-            max(0.1, self.settings.x_stats_sampling_interval)
-        )  # seconds
-        samples_to_aggregate: int = 1
-        self.publishing_interval: float = sampling_interval * samples_to_aggregate
-        self.join_assets: bool = False
-
-        self.backend_interface = interface
-        self.asset_interface: Optional[AssetInterface] = (
-            AssetInterface() if self.join_assets else None
-        )
-
-        # hardware assets
-        self.assets: List[Asset] = self._get_assets()
-
-        # OpenMetrics/Prometheus-compatible endpoints
-        self.assets.extend(self._get_open_metrics_assets())
-
-        # static system info, both hardware and software
-        self.system_info: SystemInfo = SystemInfo(
-            settings=self.settings, interface=interface
-        )
-
-        self.buffer: Dict[str, Deque[Tuple[float, float]]] = defaultdict(
-            lambda: deque([], maxlen=self.settings.x_stats_buffer_size)
-        )
-
-    def _get_assets(self) -> List["Asset"]:
-        return [
-            asset_class(
-                interface=self.asset_interface or self.backend_interface,
-                settings=self.settings,
-                shutdown_event=self._shutdown_event,
-            )
-            for asset_class in asset_registry
-        ]
-
-    def _get_open_metrics_assets(self) -> List["Asset"]:
-        open_metrics_endpoints = self.settings.x_stats_open_metrics_endpoints
-        if not open_metrics_endpoints:
-            return []
-
-        assets: List[Asset] = []
-        for name, endpoint in open_metrics_endpoints.items():
-            if not OpenMetrics.is_available(url=endpoint):
-                continue
-            logger.debug(f"Monitoring OpenMetrics endpoint: {endpoint}")
-            open_metrics = OpenMetrics(
-                interface=self.asset_interface or self.backend_interface,
-                settings=self.settings,
-                shutdown_event=self._shutdown_event,
-                name=name,
-                url=endpoint,
-            )
-            assets.append(open_metrics)  # type: ignore
-
-        return assets
-
-    def aggregate_and_publish_asset_metrics(self) -> None:
-        if self.asset_interface is None:
-            return None
-        # only extract as many items as are available in the queue at the moment
-        size = self.asset_interface.metrics_queue.qsize()
-
-        aggregated_metrics = {}
-        for _ in range(size):
-            item = self.asset_interface.metrics_queue.get()
-            aggregated_metrics.update(item)
-
-        if aggregated_metrics:
-            # update buffer:
-            # todo: get it from publish_stats instead?
-            #  either is not too accurate, just use wandb-core!
-            t = datetime.datetime.now().timestamp()
-            for k, v in aggregated_metrics.items():
-                self.buffer[k].append((t, v))
-            # publish aggregated metrics
-            self.backend_interface.publish_stats(aggregated_metrics)
-
-    def publish_telemetry(self) -> None:
-        if self.asset_interface is None:
-            return None
-        # get everything from the self.asset_interface.telemetry_queue,
-        # merge into a single dictionary and publish on the backend_interface
-        while not self.asset_interface.telemetry_queue.empty():
-            telemetry_record = self.asset_interface.telemetry_queue.get()
-            self.backend_interface._publish_telemetry(telemetry_record)
-
-    def _start(self) -> None:
-        logger.info("Starting system asset monitoring threads")
-        for asset in self.assets:
-            asset.start()
-
-        # compatibility mode: join stats from different assets before publishing
-        if not (self.join_assets and self.asset_interface is not None):
-            return None
-
-        # give the assets a chance to accumulate and publish their first stats
-        # this will provide a constant offset for the following accumulation events below
-        self._shutdown_event.wait(
-            self.publishing_interval * self.PUBLISHING_INTERVAL_DELAY_FACTOR
-        )
-
-        logger.debug("Starting system metrics aggregation loop")
-
-        while not self._shutdown_event.is_set():
-            self.publish_telemetry()
-            self.aggregate_and_publish_asset_metrics()
-            self._shutdown_event.wait(self.publishing_interval)
-
-        logger.debug("Finished system metrics aggregation loop")
-
-        # try to publish the last batch of metrics + telemetry
-        try:
-            logger.debug("Publishing last batch of metrics")
-            # publish telemetry
-            self.publish_telemetry()
-            self.aggregate_and_publish_asset_metrics()
-        except Exception:
-            logger.exception("Error publishing last batch of metrics.")
-
-    def start(self) -> None:
-        self._shutdown_event.clear()
-        if self._process is not None:
-            return None
-        logger.info("Starting system monitor")
-        self._process = threading.Thread(
-            target=self._start, daemon=True, name="SystemMonitor"
-        )
-        self._process.start()
-
-    def finish(self) -> None:
-        if self._process is None:
-            return None
-        logger.info("Stopping system monitor")
-        self._shutdown_event.set()
-        for asset in self.assets:
-            asset.finish()
-        try:
-            self._process.join()
-        except Exception:
-            logger.exception("Error joining system monitor process.")
-        self._process = None
-
-    def probe(self, publish: bool = True) -> dict:
-        logger.info("Collecting system info")
-        # collect static info about the hardware from registered assets
-        hardware_info: dict = {
-            k: v for d in [asset.probe() for asset in self.assets] for k, v in d.items()
-        }
-        # collect static info about the software environment
-        software_info: dict = self.system_info.probe()
-        # merge the two dictionaries
-        system_info = {**software_info, **hardware_info}
-        logger.debug(system_info)
-        logger.info("Finished collecting system info")
-
-        if publish:
-            logger.info("Publishing system info")
-            self.system_info.publish(system_info)
-            logger.info("Finished publishing system info")
-
-        return system_info

wandb/sdk/internal/writer.py

@@ -1,204 +0,0 @@
-"""Writer thread."""
-
-import logging
-from typing import TYPE_CHECKING, Callable, Optional
-
-from wandb.proto import wandb_internal_pb2 as pb
-from wandb.proto import wandb_telemetry_pb2 as tpb
-
-from ..interface.interface_queue import InterfaceQueue
-from ..lib import proto_util, telemetry
-from . import context, datastore, flow_control
-from .settings_static import SettingsStatic
-
-if TYPE_CHECKING:
-    from queue import Queue
-
-
-logger = logging.getLogger(__name__)
-
-
-class WriteManager:
-    _settings: SettingsStatic
-    _record_q: "Queue[pb.Record]"
-    _result_q: "Queue[pb.Result]"
-    _sender_q: "Queue[pb.Record]"
-    _interface: InterfaceQueue
-    _context_keeper: context.ContextKeeper
-
-    _ds: Optional[datastore.DataStore]
-    _flow_control: Optional[flow_control.FlowControl]
-    _status_report: Optional["pb.StatusReportRequest"]
-    _record_num: int
-    _telemetry_obj: tpb.TelemetryRecord
-    _telemetry_overflow: bool
-    _use_flow_control: bool
-
-    # TODO(cancel_paused): implement me
-    # _sender_cancel_set: Set[str]
-
-    def __init__(
-        self,
-        settings: SettingsStatic,
-        record_q: "Queue[pb.Record]",
-        result_q: "Queue[pb.Result]",
-        sender_q: "Queue[pb.Record]",
-        interface: InterfaceQueue,
-        context_keeper: context.ContextKeeper,
-    ):
-        self._settings = settings
-        self._record_q = record_q
-        self._result_q = result_q
-        self._sender_q = sender_q
-        self._interface = interface
-        self._context_keeper = context_keeper
-
-        # TODO(cancel_paused): implement me
-        # self._sender_cancel_set = set()
-
-        self._ds = None
-        self._flow_control = None
-        self._status_report = None
-        self._record_num = 0
-        self._telemetry_obj = tpb.TelemetryRecord()
-        self._telemetry_overflow = False
-        self._use_flow_control = not (
-            self._settings.x_flow_control_disabled or self._settings._offline
-        )
-
-    def open(self) -> None:
-        self._ds = datastore.DataStore()
-        self._ds.open_for_write(self._settings.sync_file)
-        self._flow_control = flow_control.FlowControl(
-            settings=self._settings,
-            write_record=self._write_record,
-            forward_record=self._forward_record,
-            pause_marker=self._pause_marker,
-            recover_records=self._recover_records,
-        )
-
-    def _forward_record(self, record: "pb.Record") -> None:
-        self._context_keeper.add_from_record(record)
-        self._sender_q.put(record)
-
-    def _send_mark(self) -> None:
-        sender_mark = pb.SenderMarkRequest()
-        record = self._interface._make_request(sender_mark=sender_mark)
-        self._forward_record(record)
-
-    def _maybe_send_telemetry(self) -> None:
-        if self._telemetry_overflow:
-            return
-        self._telemetry_overflow = True
-        with telemetry.context(obj=self._telemetry_obj) as tel:
-            tel.feature.flow_control_overflow = True
-        telemetry_record = pb.TelemetryRecordRequest(telemetry=self._telemetry_obj)
-        record = self._interface._make_request(telemetry_record=telemetry_record)
-        self._forward_record(record)
-
-    def _pause_marker(self) -> None:
-        self._maybe_send_telemetry()
-        self._send_mark()
-
-    def _write_record(self, record: "pb.Record") -> int:
-        assert self._ds
-
-        self._record_num += 1
-        proto_util._assign_record_num(record, self._record_num)
-        ret = self._ds.write(record)
-        assert ret is not None
-
-        _start_offset, end_offset, _flush_offset = ret
-        proto_util._assign_end_offset(record, end_offset)
-        return end_offset
-
-    def _ensure_flushed(self, offset: int) -> None:
-        if self._ds:
-            self._ds.ensure_flushed(offset)
-
-    def _recover_records(self, start: int, end: int) -> None:
-        sender_read = pb.SenderReadRequest(start_offset=start, final_offset=end)
-        # TODO(cancel_paused): implement me
-        # for cancel_id in self._sender_cancel_set:
-        #     sender_read.cancel_list.append(cancel_id)
-        record = self._interface._make_request(sender_read=sender_read)
-        self._ensure_flushed(end)
-        self._forward_record(record)
-
-    def _write(self, record: "pb.Record") -> None:
-        if not self._ds:
-            self.open()
-        assert self._flow_control
-
-        if not record.control.local:
-            self._write_record(record)
-
-        if self._use_flow_control:
-            self._flow_control.flow(record)
-        elif not self._settings._offline or record.control.always_send:
-            # when flow_control is disabled we pass through all records to
-            # the sender as long as we are online.  The exception is there
-            # are special records that we always pass to the sender
-            # (namely the exit record so we can trigger the defer shutdown
-            # state machine)
-            self._forward_record(record)
-
-    def write(self, record: "pb.Record") -> None:
-        record_type = record.WhichOneof("record_type")
-        assert record_type
-        writer_str = "write_" + record_type
-        write_handler: Callable[[pb.Record], None] = getattr(
-            self, writer_str, self._write
-        )
-        write_handler(record)
-
-    def write_request(self, record: "pb.Record") -> None:
-        request_type = record.request.WhichOneof("request_type")
-        assert request_type
-        write_request_str = "write_request_" + request_type
-        write_request_handler: Optional[Callable[[pb.Record], None]] = getattr(
-            self, write_request_str, None
-        )
-        if write_request_handler:
-            return write_request_handler(record)
-        self._write(record)
-
-    def write_request_run_status(self, record: "pb.Record") -> None:
-        result = proto_util._result_from_record(record)
-        if self._status_report:
-            result.response.run_status_response.sync_time.CopyFrom(
-                self._status_report.sync_time
-            )
-            send_record_num = self._status_report.record_num
-            result.response.run_status_response.sync_items_total = self._record_num
-            result.response.run_status_response.sync_items_pending = (
-                self._record_num - send_record_num
-            )
-        self._respond_result(result)
-
-    def write_request_status_report(self, record: "pb.Record") -> None:
-        self._status_report = record.request.status_report
-        self._write(record)
-
-    def write_request_cancel(self, record: "pb.Record") -> None:
-        cancel_id = record.request.cancel.cancel_slot
-        self._context_keeper.cancel(cancel_id)
-
-        # TODO(cancel_paused): implement me
-        # cancelled = self._context_keeper.cancel(cancel_id)
-        # if not cancelled:
-        #     self._sender_cancel_set.add(cancel_id)
-
-    def _respond_result(self, result: "pb.Result") -> None:
-        self._result_q.put(result)
-
-    def finish(self) -> None:
-        if self._flow_control:
-            self._flow_control.flush()
-        if self._ds:
-            self._ds.close()
-        # TODO(debug_context) see context.py
-        # self._context_keeper._debug_print_orphans(print_to_stdout=self._settings._debug)
-
-    def debounce(self) -> None:
-        pass

wandb/sdk/lib/service_token.py

@@ -1,93 +0,0 @@
-from __future__ import annotations
-
-import dataclasses
-import os
-
-from wandb import env
-
-_CURRENT_VERSION = "2"
-_SUPPORTED_TRANSPORTS = "tcp"
-
-
-def get_service_token() -> ServiceToken | None:
-    """Reads the token from environment variables.
-
-    Returns:
-        The token if the correct environment variable is set, or None.
-
-    Raises:
-        ValueError: If the environment variable is set but cannot be
-            parsed.
-    """
-    token = os.environ.get(env.SERVICE)
-    if not token:
-        return None
-
-    parts = token.split("-")
-    if len(parts) != 5:
-        raise ValueError(f"Invalid token: {token}")
-
-    version, pid_str, transport, host, port_str = parts
-
-    if version != _CURRENT_VERSION:
-        raise ValueError(
-            f"Expected version {_CURRENT_VERSION}, but got {version} (token={token})"
-        )
-    if transport not in _SUPPORTED_TRANSPORTS:
-        raise ValueError(
-            f"Unsupported transport: {transport} (token={token})",
-        )
-
-    try:
-        return ServiceToken(
-            version=version,
-            pid=int(pid_str),
-            transport=transport,
-            host=host,
-            port=int(port_str),
-        )
-    except ValueError as e:
-        raise ValueError(f"Invalid token: {token}") from e
-
-
-def set_service_token(parent_pid: int, transport: str, host: str, port: int) -> None:
-    """Stores a service token in an environment variable.
-
-    Args:
-        parent_pid: The process ID of the process that started the service.
-        transport: The transport used to communicate with the service.
-        host: The host part of the internet address on which the service
-            is listening (e.g. localhost).
-        port: The port the service is listening on.
-
-    Raises:
-        ValueError: If given an unsupported transport.
-    """
-    if transport not in _SUPPORTED_TRANSPORTS:
-        raise ValueError(f"Unsupported transport: {transport}")
-
-    os.environ[env.SERVICE] = "-".join(
-        (
-            _CURRENT_VERSION,
-            str(parent_pid),
-            transport,
-            host,
-            str(port),
-        )
-    )
-
-
-def clear_service_token() -> None:
-    """Clears the environment variable storing the service token."""
-    os.environ.pop(env.SERVICE, None)
-
-
-@dataclasses.dataclass(frozen=True)
-class ServiceToken:
-    """An identifier for a running service process."""
-
-    version: str
-    pid: int
-    transport: str
-    host: str
-    port: int

wandb/sdk/service/_startup_debug.py

@@ -1,22 +0,0 @@
-"""_startup_debug.
-
-Temporary helper to debug issues with wandb service startup
-"""
-
-import os
-import time
-
-
-def is_enabled() -> bool:
-    # This is very temporary to help diagnose problems seen by some
-    # customers which we are having trouble reproducing. It should be
-    # replaced by something more permanent in the future when we have
-    # proper logging for wandb-service
-    if os.environ.get("_WANDB_STARTUP_DEBUG"):
-        return True
-    return False
-
-
-def print_message(message: str) -> None:
-    time_now = time.time()
-    print("WANDB_STARTUP_DEBUG", time_now, message)  # noqa: T201

wandb/sdk/service/port_file.py

@@ -1,53 +0,0 @@
-"""port_file: write/read file containing port info."""
-
-import os
-import tempfile
-from typing import Optional
-
-
-class PortFile:
-    _sock_port: Optional[int]
-    _valid: bool
-
-    SOCK_TOKEN = "sock="
-    EOF_TOKEN = "EOF"
-
-    def __init__(self, sock_port: Optional[int] = None) -> None:
-        self._sock_port = sock_port
-        self._valid = False
-
-    def write(self, fname: str) -> None:
-        dname, bname = os.path.split(fname)
-        f = tempfile.NamedTemporaryFile(prefix=bname, dir=dname, mode="w", delete=False)
-        try:
-            tmp_filename = f.name
-            with f:
-                data = []
-                if self._sock_port:
-                    data.append(f"{self.SOCK_TOKEN}{self._sock_port}")
-                data.append(self.EOF_TOKEN)
-                port_str = "\n".join(data)
-                written = f.write(port_str)
-                assert written == len(port_str)
-            os.rename(tmp_filename, fname)
-        except Exception:
-            os.unlink(tmp_filename)
-            raise
-
-    def read(self, fname: str) -> None:
-        with open(fname) as f:
-            lines = f.readlines()
-            if lines[-1] != self.EOF_TOKEN:
-                return
-            for ln in lines:
-                if ln.startswith(self.SOCK_TOKEN):
-                    self._sock_port = int(ln[len(self.SOCK_TOKEN) :])
-            self._valid = True
-
-    @property
-    def sock_port(self) -> Optional[int]:
-        return self._sock_port
-
-    @property
-    def is_valid(self) -> bool:
-        return self._valid