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

wandb/sdk/internal/flow_control.py

@@ -1,263 +0,0 @@
-"""Flow Control.
-
-States:
-    FORWARDING
-    PAUSING
-
-New messages:
-    pb.SenderMarkRequest    writer -> sender (empty message)
-    pb.StatusReportRequest  sender -> writer (reports current sender progress)
-    pb.SenderReadRequest    writer -> sender (requests read of transaction log)
-
-Thresholds:
-    Threshold_High_MaxOutstandingData      - When above this, stop sending requests to sender
-    Threshold_Mid_StartSendingReadRequests - When below this, start sending read requests
-    Threshold_Low_RestartSendingData       - When below this, start sending normal records
-
-State machine:
-    FORWARDING
-      -> PAUSED if should_pause
-         There is too much work outstanding to the sender thread, after the current request
-         lets stop sending data.
-    PAUSING
-      -> FORWARDING if should_unpause
-      -> PAUSING if should_recover
-      -> PAUSING if should_quiesce
-
-"""
-
-import logging
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Optional
-
-from wandb.proto import wandb_internal_pb2 as pb
-from wandb.sdk.lib import fsm
-
-from .settings_static import SettingsStatic
-
-if TYPE_CHECKING:
-    from wandb.proto.wandb_internal_pb2 import Record
-
-logger = logging.getLogger(__name__)
-
-# By default we will allow 400 MiB of requests in the sender queue
-# before falling back to the transaction log.
-DEFAULT_THRESHOLD = 128 * 1024 * 1024  # 128 MiB
-
-
-def _get_request_type(record: "Record") -> Optional[str]:
-    record_type = record.WhichOneof("record_type")
-    if record_type != "request":
-        return None
-    request_type = record.request.WhichOneof("request_type")
-    return request_type
-
-
-def _is_control_record(record: "Record") -> bool:
-    return record.control.flow_control
-
-
-def _is_local_non_control_record(record: "Record") -> bool:
-    return record.control.local and not record.control.flow_control
-
-
-@dataclass
-class StateContext:
-    last_forwarded_offset: int = 0
-    last_sent_offset: int = 0
-    last_written_offset: int = 0
-
-
-class FlowControl:
-    _fsm: fsm.FsmWithContext["Record", StateContext]
-
-    def __init__(
-        self,
-        settings: SettingsStatic,
-        forward_record: Callable[["Record"], None],
-        write_record: Callable[["Record"], int],
-        pause_marker: Callable[[], None],
-        recover_records: Callable[[int, int], None],
-        _threshold_bytes_high: int = 0,
-        _threshold_bytes_mid: int = 0,
-        _threshold_bytes_low: int = 0,
-    ) -> None:
-        # thresholds to define when to PAUSE, RESTART, FORWARDING
-        if (
-            _threshold_bytes_high == 0
-            or _threshold_bytes_mid == 0
-            or _threshold_bytes_low == 0
-        ):
-            threshold = settings.x_network_buffer or DEFAULT_THRESHOLD
-            _threshold_bytes_high = threshold
-            _threshold_bytes_mid = threshold // 2
-            _threshold_bytes_low = threshold // 4
-        assert _threshold_bytes_high > _threshold_bytes_mid > _threshold_bytes_low
-
-        # FSM definition
-        state_forwarding = StateForwarding(
-            forward_record=forward_record,
-            pause_marker=pause_marker,
-            threshold_pause=_threshold_bytes_high,
-        )
-        state_pausing = StatePausing(
-            forward_record=forward_record,
-            recover_records=recover_records,
-            threshold_recover=_threshold_bytes_mid,
-            threshold_forward=_threshold_bytes_low,
-        )
-        self._fsm = fsm.FsmWithContext(
-            states=[state_forwarding, state_pausing],
-            table={
-                StateForwarding: [
-                    fsm.FsmEntry(
-                        state_forwarding._should_pause,
-                        StatePausing,
-                        state_forwarding._pause,
-                    ),
-                ],
-                StatePausing: [
-                    fsm.FsmEntry(
-                        state_pausing._should_unpause,
-                        StateForwarding,
-                        state_pausing._unpause,
-                    ),
-                    fsm.FsmEntry(
-                        state_pausing._should_recover,
-                        StatePausing,
-                        state_pausing._recover,
-                    ),
-                    fsm.FsmEntry(
-                        state_pausing._should_quiesce,
-                        StatePausing,
-                        state_pausing._quiesce,
-                    ),
-                ],
-            },
-        )
-
-    def flush(self) -> None:
-        # TODO(mempressure): what do we do here, how do we make sure we dont have work in pause state
-        pass
-
-    def flow(self, record: "Record") -> None:
-        self._fsm.input(record)
-
-
-class StateShared:
-    _context: StateContext
-
-    def __init__(self) -> None:
-        self._context = StateContext()
-
-    def _update_written_offset(self, record: "Record") -> None:
-        end_offset = record.control.end_offset
-        if end_offset:
-            self._context.last_written_offset = end_offset
-
-    def _update_forwarded_offset(self) -> None:
-        self._context.last_forwarded_offset = self._context.last_written_offset
-
-    def _process(self, record: "Record") -> None:
-        request_type = _get_request_type(record)
-        if not request_type:
-            return
-        process_str = f"_process_{request_type}"
-        process_handler: Optional[Callable[[pb.Record], None]] = getattr(
-            self, process_str, None
-        )
-        if not process_handler:
-            return
-        process_handler(record)
-
-    def _process_status_report(self, record: "Record") -> None:
-        sent_offset = record.request.status_report.sent_offset
-        self._context.last_sent_offset = sent_offset
-
-    def on_exit(self, record: "Record") -> StateContext:
-        return self._context
-
-    def on_enter(self, record: "Record", context: StateContext) -> None:
-        self._context = context
-
-    @property
-    def _behind_bytes(self) -> int:
-        return self._context.last_forwarded_offset - self._context.last_sent_offset
-
-
-class StateForwarding(StateShared):
-    _forward_record: Callable[["Record"], None]
-    _pause_marker: Callable[[], None]
-    _threshold_pause: int
-
-    def __init__(
-        self,
-        forward_record: Callable[["Record"], None],
-        pause_marker: Callable[[], None],
-        threshold_pause: int,
-    ) -> None:
-        super().__init__()
-        self._forward_record = forward_record
-        self._pause_marker = pause_marker
-        self._threshold_pause = threshold_pause
-
-    def _should_pause(self, record: "Record") -> bool:
-        return self._behind_bytes >= self._threshold_pause
-
-    def _pause(self, record: "Record") -> None:
-        self._pause_marker()
-
-    def on_check(self, record: "Record") -> None:
-        self._update_written_offset(record)
-        self._process(record)
-        if not _is_control_record(record):
-            self._forward_record(record)
-        self._update_forwarded_offset()
-
-
-class StatePausing(StateShared):
-    _forward_record: Callable[["Record"], None]
-    _recover_records: Callable[[int, int], None]
-    _threshold_recover: int
-    _threshold_forward: int
-
-    def __init__(
-        self,
-        forward_record: Callable[["Record"], None],
-        recover_records: Callable[[int, int], None],
-        threshold_recover: int,
-        threshold_forward: int,
-    ) -> None:
-        super().__init__()
-        self._forward_record = forward_record
-        self._recover_records = recover_records
-        self._threshold_recover = threshold_recover
-        self._threshold_forward = threshold_forward
-
-    def _should_unpause(self, record: "Record") -> bool:
-        return self._behind_bytes < self._threshold_forward
-
-    def _unpause(self, record: "Record") -> None:
-        self._quiesce(record)
-
-    def _should_recover(self, record: "Record") -> bool:
-        return self._behind_bytes < self._threshold_recover
-
-    def _recover(self, record: "Record") -> None:
-        self._quiesce(record)
-
-    def _should_quiesce(self, record: "Record") -> bool:
-        return _is_local_non_control_record(record)
-
-    def _quiesce(self, record: "Record") -> None:
-        start = self._context.last_forwarded_offset
-        end = self._context.last_written_offset
-        if start != end:
-            self._recover_records(start, end)
-        if _is_local_non_control_record(record):
-            self._forward_record(record)
-        self._update_forwarded_offset()
-
-    def on_check(self, record: "Record") -> None:
-        self._update_written_offset(record)
-        self._process(record)

wandb/sdk/internal/internal.py

@@ -1,401 +0,0 @@
-#
-"""Internal process.
-
-This module implements the entrypoint for the internal process. The internal process
-is responsible for handling "record" requests, and responding with "results". Data is
-passed to the process over multiprocessing queues.
-
-Threads:
-    HandlerThread -- read from record queue and call handlers
-    SenderThread -- send to network
-    WriterThread -- write to disk
-
-"""
-
-import atexit
-import logging
-import os
-import queue
-import sys
-import threading
-import time
-import traceback
-from datetime import datetime
-from typing import TYPE_CHECKING, Any, List, Optional
-
-import psutil
-
-import wandb
-
-from ..interface.interface_queue import InterfaceQueue
-from . import context, handler, internal_util, sender, writer
-
-if TYPE_CHECKING:
-    from queue import Queue
-    from threading import Event
-
-    from wandb.proto.wandb_internal_pb2 import Record, Result
-
-    from .internal_util import RecordLoopThread
-    from .settings_static import SettingsStatic
-
-
-logger = logging.getLogger(__name__)
-
-
-def wandb_internal(
-    settings: "SettingsStatic",
-    record_q: "Queue[Record]",
-    result_q: "Queue[Result]",
-    port: Optional[int] = None,
-    user_pid: Optional[int] = None,
-) -> None:
-    """Internal process function entrypoint.
-
-    Read from record queue and dispatch work to various threads.
-
-    Args:
-        settings: settings object
-        record_q: records to be handled
-        result_q: for sending results back
-
-    """
-    # mark this process as internal
-    wandb._set_internal_process()  # type: ignore
-    started = time.time()
-
-    # any sentry events in the internal process will be tagged as such
-    wandb._sentry.configure_scope(process_context="internal", tags=dict(settings))
-
-    # register the exit handler only when wandb_internal is called, not on import
-    @atexit.register
-    def handle_exit(*args: "Any") -> None:
-        logger.info("Internal process exited")
-
-    # Let's make sure we don't modify settings so use a static object
-    _settings = settings
-    if _settings.log_internal:
-        configure_logging(_settings.log_internal, _settings.x_log_level)
-
-    user_pid = user_pid or os.getppid()
-    pid = os.getpid()
-
-    logger.info(
-        "W&B internal server running at pid: %s, started at: %s",
-        pid,
-        datetime.fromtimestamp(started),
-    )
-
-    publish_interface = InterfaceQueue(record_q=record_q)
-
-    stopped = threading.Event()
-    threads: List[RecordLoopThread] = []
-
-    context_keeper = context.ContextKeeper()
-
-    send_record_q: Queue[Record] = queue.Queue()
-
-    write_record_q: Queue[Record] = queue.Queue()
-
-    record_sender_thread = SenderThread(
-        settings=_settings,
-        record_q=send_record_q,
-        result_q=result_q,
-        stopped=stopped,
-        interface=publish_interface,
-        debounce_interval_ms=5000,
-        context_keeper=context_keeper,
-    )
-    threads.append(record_sender_thread)
-
-    record_writer_thread = WriterThread(
-        settings=_settings,
-        record_q=write_record_q,
-        result_q=result_q,
-        stopped=stopped,
-        interface=publish_interface,
-        sender_q=send_record_q,
-        context_keeper=context_keeper,
-    )
-    threads.append(record_writer_thread)
-
-    record_handler_thread = HandlerThread(
-        settings=_settings,
-        record_q=record_q,
-        result_q=result_q,
-        stopped=stopped,
-        writer_q=write_record_q,
-        interface=publish_interface,
-        context_keeper=context_keeper,
-    )
-    threads.append(record_handler_thread)
-
-    process_check = ProcessCheck(settings=_settings, user_pid=user_pid)
-
-    for thread in threads:
-        thread.start()
-
-    interrupt_count = 0
-    while not stopped.is_set():
-        try:
-            # wait for stop event
-            while not stopped.is_set():
-                time.sleep(1)
-                if process_check.is_dead():
-                    logger.error("Internal process shutdown.")
-                    stopped.set()
-        except KeyboardInterrupt:
-            interrupt_count += 1
-            logger.warning(f"Internal process interrupt: {interrupt_count}")
-        finally:
-            if interrupt_count >= 2:
-                logger.error("Internal process interrupted.")
-                stopped.set()
-
-    for thread in threads:
-        thread.join()
-
-    def close_internal_log() -> None:
-        root = logging.getLogger("wandb")
-        for _handler in root.handlers[:]:
-            _handler.close()
-            root.removeHandler(_handler)
-
-    for thread in threads:
-        exc_info = thread.get_exception()
-        if exc_info:
-            logger.error(f"Thread {thread.name}:", exc_info=exc_info)
-            print(f"Thread {thread.name}:", file=sys.stderr)  # noqa: T201
-            traceback.print_exception(*exc_info)
-            wandb._sentry.exception(exc_info)
-            wandb.termerror("Internal wandb error: file data was not synced")
-            # TODO: We can make this more graceful by returning an error to streams.py
-            # and potentially just fail the one stream.
-            os._exit(-1)
-
-    close_internal_log()
-
-
-def configure_logging(
-    log_fname: str, log_level: int, run_id: Optional[str] = None
-) -> None:
-    # TODO: we may want make prints and stdout make it into the logs
-    # sys.stdout = open(settings.log_internal, "a")
-    # sys.stderr = open(settings.log_internal, "a")
-    log_handler = logging.FileHandler(log_fname)
-    log_handler.setLevel(log_level)
-
-    class WBFilter(logging.Filter):
-        def filter(self, record: "Any") -> bool:
-            record.run_id = run_id
-            return True
-
-    if run_id:
-        formatter = logging.Formatter(
-            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d "
-            "[%(run_id)s:%(filename)s:%(funcName)s():%(lineno)s] %(message)s"
-        )
-    else:
-        formatter = logging.Formatter(
-            "%(asctime)s %(levelname)-7s %(threadName)-10s:%(process)d "
-            "[%(filename)s:%(funcName)s():%(lineno)s] %(message)s"
-        )
-
-    log_handler.setFormatter(formatter)
-    if run_id:
-        log_handler.addFilter(WBFilter())
-    # If this is called without "wandb", backend logs from this module
-    # are not streamed to `debug-internal.log` when we spawn with fork
-    # TODO: (cvp) we should really take another pass at logging in general
-    root = logging.getLogger("wandb")
-    root.propagate = False
-    root.setLevel(logging.DEBUG)
-    root.addHandler(log_handler)
-
-
-class HandlerThread(internal_util.RecordLoopThread):
-    """Read records from queue and dispatch to handler routines."""
-
-    _record_q: "Queue[Record]"
-    _result_q: "Queue[Result]"
-    _stopped: "Event"
-    _context_keeper: context.ContextKeeper
-
-    def __init__(
-        self,
-        settings: "SettingsStatic",
-        record_q: "Queue[Record]",
-        result_q: "Queue[Result]",
-        stopped: "Event",
-        writer_q: "Queue[Record]",
-        interface: "InterfaceQueue",
-        context_keeper: context.ContextKeeper,
-        debounce_interval_ms: "float" = 1000,
-    ) -> None:
-        super().__init__(
-            input_record_q=record_q,
-            result_q=result_q,
-            stopped=stopped,
-            debounce_interval_ms=debounce_interval_ms,
-        )
-        self.name = "HandlerThread"
-        self._settings = settings
-        self._record_q = record_q
-        self._result_q = result_q
-        self._stopped = stopped
-        self._writer_q = writer_q
-        self._interface = interface
-        self._context_keeper = context_keeper
-
-    def _setup(self) -> None:
-        self._hm = handler.HandleManager(
-            settings=self._settings,
-            record_q=self._record_q,
-            result_q=self._result_q,
-            stopped=self._stopped,
-            writer_q=self._writer_q,
-            interface=self._interface,
-            context_keeper=self._context_keeper,
-        )
-
-    def _process(self, record: "Record") -> None:
-        self._hm.handle(record)
-
-    def _finish(self) -> None:
-        self._hm.finish()
-
-    def _debounce(self) -> None:
-        self._hm.debounce()
-
-
-class SenderThread(internal_util.RecordLoopThread):
-    """Read records from queue and dispatch to sender routines."""
-
-    _record_q: "Queue[Record]"
-    _result_q: "Queue[Result]"
-    _context_keeper: context.ContextKeeper
-
-    def __init__(
-        self,
-        settings: "SettingsStatic",
-        record_q: "Queue[Record]",
-        result_q: "Queue[Result]",
-        stopped: "Event",
-        interface: "InterfaceQueue",
-        context_keeper: context.ContextKeeper,
-        debounce_interval_ms: "float" = 5000,
-    ) -> None:
-        super().__init__(
-            input_record_q=record_q,
-            result_q=result_q,
-            stopped=stopped,
-            debounce_interval_ms=debounce_interval_ms,
-        )
-        self.name = "SenderThread"
-        self._settings = settings
-        self._record_q = record_q
-        self._result_q = result_q
-        self._interface = interface
-        self._context_keeper = context_keeper
-
-    def _setup(self) -> None:
-        self._sm = sender.SendManager(
-            settings=self._settings,
-            record_q=self._record_q,
-            result_q=self._result_q,
-            interface=self._interface,
-            context_keeper=self._context_keeper,
-        )
-
-    def _process(self, record: "Record") -> None:
-        self._sm.send(record)
-
-    def _finish(self) -> None:
-        self._sm.finish()
-
-    def _debounce(self) -> None:
-        self._sm.debounce()
-
-
-class WriterThread(internal_util.RecordLoopThread):
-    """Read records from queue and dispatch to writer routines."""
-
-    _record_q: "Queue[Record]"
-    _result_q: "Queue[Result]"
-    _context_keeper: context.ContextKeeper
-
-    def __init__(
-        self,
-        settings: "SettingsStatic",
-        record_q: "Queue[Record]",
-        result_q: "Queue[Result]",
-        stopped: "Event",
-        interface: "InterfaceQueue",
-        sender_q: "Queue[Record]",
-        context_keeper: context.ContextKeeper,
-        debounce_interval_ms: "float" = 1000,
-    ) -> None:
-        super().__init__(
-            input_record_q=record_q,
-            result_q=result_q,
-            stopped=stopped,
-            debounce_interval_ms=debounce_interval_ms,
-        )
-        self.name = "WriterThread"
-        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
-
-    def _setup(self) -> None:
-        self._wm = writer.WriteManager(
-            settings=self._settings,
-            record_q=self._record_q,
-            result_q=self._result_q,
-            sender_q=self._sender_q,
-            interface=self._interface,
-            context_keeper=self._context_keeper,
-        )
-
-    def _process(self, record: "Record") -> None:
-        self._wm.write(record)
-
-    def _finish(self) -> None:
-        self._wm.finish()
-
-    def _debounce(self) -> None:
-        self._wm.debounce()
-
-
-class ProcessCheck:
-    """Class to help watch a process id to detect when it is dead."""
-
-    check_process_last: Optional[float]
-
-    def __init__(self, settings: "SettingsStatic", user_pid: Optional[int]) -> None:
-        self.settings = settings
-        self.pid = user_pid
-        self.check_process_last = None
-        self.check_process_interval = settings.x_internal_check_process
-
-    def is_dead(self) -> bool:
-        if not self.check_process_interval or not self.pid:
-            return False
-        time_now = time.time()
-        if (
-            self.check_process_last
-            and time_now < self.check_process_last + self.check_process_interval
-        ):
-            return False
-        self.check_process_last = time_now
-
-        # TODO(jhr): check for os.getppid on unix being 1?
-        exists = psutil.pid_exists(self.pid)
-        if not exists:
-            logger.warning(
-                f"Internal process exiting, parent pid {self.pid} disappeared"
-            )
-            return True
-        return False