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

wandb/sdk/service/server.py

@@ -1,107 +0,0 @@
-"""wandb server.
-
-Start up socket transport servers.
-"""
-
-import logging
-import os
-import sys
-from typing import Optional
-
-import wandb
-
-from . import _startup_debug, port_file
-from .server_sock import SocketServer
-from .streams import StreamMux
-
-
-class WandbServer:
-    _pid: Optional[int]
-    _sock_port: Optional[int]
-    _debug: bool
-    _sock_server: Optional[SocketServer]
-    _startup_debug_enabled: bool
-
-    def __init__(
-        self,
-        sock_port: Optional[int] = None,
-        port_fname: Optional[str] = None,
-        address: Optional[str] = None,
-        pid: Optional[int] = None,
-        debug: bool = True,
-    ) -> None:
-        self._sock_port = sock_port
-        self._port_fname = port_fname
-        self._address = address
-        self._pid = pid
-        self._debug = debug
-        self._sock_server = None
-        self._startup_debug_enabled = _startup_debug.is_enabled()
-
-        if debug:
-            logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
-
-    def _inform_used_ports(self, sock_port: Optional[int]) -> None:
-        if not self._port_fname:
-            return
-        pf = port_file.PortFile(sock_port=sock_port)
-        pf.write(self._port_fname)
-
-    def _start_sock(self, mux: StreamMux) -> int:
-        address: str = self._address or "127.0.0.1"
-        port: int = self._sock_port or 0
-        self._sock_server = SocketServer(mux=mux, address=address, port=port)
-        try:
-            self._sock_server.start()
-            port = self._sock_server.port
-            if self._pid:
-                mux.set_pid(self._pid)
-        except KeyboardInterrupt:
-            mux.cleanup()
-            raise
-        except Exception:
-            mux.cleanup()
-            raise
-        return port
-
-    def _stop_servers(self) -> None:
-        if self._sock_server:
-            self._sock_server.stop()
-
-    def _startup_debug_print(self, message: str) -> None:
-        if not self._startup_debug_enabled:
-            return
-        _startup_debug.print_message(message)
-
-    def _setup_proctitle(self, sock_port: Optional[int]) -> None:
-        # TODO: the internal_process should have a better way to have access to
-        # settings.
-        disable_setproctitle = os.environ.get("WANDB_X_DISABLE_SETPROCTITLE")
-        if disable_setproctitle:
-            return
-
-        setproctitle = wandb.util.get_optional_module("setproctitle")
-        if setproctitle:
-            service_ver = 2
-            pid = str(self._pid or 0)
-            transport = "s" if sock_port else "g"
-            port = sock_port or 0
-            # this format is similar to the service token, but it's purely informative now
-            # (consider unifying this in the future)
-            service_id = f"{service_ver}-{pid}-{transport}-{port}"
-            proc_title = f"wandb-service({service_id})"
-            self._startup_debug_print("before_setproctitle")
-            setproctitle.setproctitle(proc_title)
-            self._startup_debug_print("after_setproctitle")
-
-    def serve(self) -> None:
-        mux = StreamMux()
-        self._startup_debug_print("before_network")
-        sock_port = self._start_sock(mux=mux)
-        self._startup_debug_print("after_network")
-        self._inform_used_ports(sock_port=sock_port)
-        self._startup_debug_print("after_inform")
-        self._setup_proctitle(sock_port=sock_port)
-        self._startup_debug_print("before_loop")
-        mux.loop()
-        self._stop_servers()

wandb/sdk/service/server_sock.py

@@ -1,286 +0,0 @@
-import queue
-import socket
-import threading
-import time
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
-
-import wandb
-from wandb.proto import wandb_internal_pb2 as pb
-from wandb.proto import wandb_server_pb2 as spb
-from wandb.sdk.internal.settings_static import SettingsStatic
-
-from ..lib.sock_client import SockClient, SockClientClosedError
-from .streams import StreamMux
-
-if TYPE_CHECKING:
-    from threading import Event
-
-    from ..interface.interface_relay import InterfaceRelay
-
-
-class ClientDict:
-    _client_dict: Dict[str, SockClient]
-    _lock: threading.Lock
-
-    def __init__(self) -> None:
-        self._client_dict = {}
-        self._lock = threading.Lock()
-
-    def get_client(self, client_id: str) -> Optional[SockClient]:
-        with self._lock:
-            client = self._client_dict.get(client_id)
-        return client
-
-    def add_client(self, client: SockClient) -> None:
-        with self._lock:
-            self._client_dict[client._sockid] = client
-
-    def del_client(self, client: SockClient) -> None:
-        with self._lock:
-            del self._client_dict[client._sockid]
-
-
-class SockServerInterfaceReaderThread(threading.Thread):
-    _socket_client: SockClient
-    _stopped: "Event"
-
-    def __init__(
-        self,
-        clients: ClientDict,
-        iface: "InterfaceRelay",
-        stopped: "Event",
-    ) -> None:
-        self._iface = iface
-        self._clients = clients
-        threading.Thread.__init__(self)
-        self.name = "SockSrvIntRdThr"
-        self._stopped = stopped
-
-    def run(self) -> None:
-        while not self._stopped.is_set():
-            try:
-                result = self._iface.relay_q.get(timeout=1)
-            except queue.Empty:
-                continue
-            except OSError:
-                # handle is closed
-                break
-            except ValueError:
-                # queue is closed
-                break
-            sockid = result.control.relay_id
-            assert sockid
-            sock_client = self._clients.get_client(sockid)
-            assert sock_client
-            sresp = spb.ServerResponse()
-            sresp.request_id = result.control.mailbox_slot
-            sresp.result_communicate.CopyFrom(result)
-            sock_client.send_server_response(sresp)
-
-
-class SockServerReadThread(threading.Thread):
-    _sock_client: SockClient
-    _mux: StreamMux
-    _stopped: "Event"
-    _clients: ClientDict
-
-    def __init__(
-        self, conn: socket.socket, mux: StreamMux, clients: ClientDict
-    ) -> None:
-        self._mux = mux
-        threading.Thread.__init__(self)
-        self.name = "SockSrvRdThr"
-        sock_client = SockClient()
-        sock_client.set_socket(conn)
-        self._sock_client = sock_client
-        self._stopped = mux._get_stopped_event()
-        self._clients = clients
-
-    def run(self) -> None:
-        while not self._stopped.is_set():
-            try:
-                sreq = self._sock_client.read_server_request()
-            except SockClientClosedError:
-                # socket has been closed
-                # TODO: shut down other threads serving this socket?
-                break
-            assert sreq, "read_server_request should never timeout"
-            sreq_type = sreq.WhichOneof("server_request_type")
-            shandler_str = "server_" + sreq_type  # type: ignore
-            shandler: Callable[[spb.ServerRequest], None] = getattr(  # type: ignore
-                self, shandler_str, None
-            )
-            assert shandler, f"unknown handle: {shandler_str}"  # type: ignore
-            shandler(sreq)
-
-    def stop(self) -> None:
-        try:
-            # See shutdown notes in class SocketServer for a discussion about this mechanism
-            self._sock_client.shutdown(socket.SHUT_RDWR)
-        except OSError:
-            pass
-        self._sock_client.close()
-
-    def server_inform_init(self, sreq: "spb.ServerRequest") -> None:
-        request = sreq.inform_init
-        stream_id = request._info.stream_id
-        settings = SettingsStatic(request.settings)
-        self._mux.add_stream(stream_id, settings=settings)
-
-        iface = self._mux.get_stream(stream_id).interface
-        self._clients.add_client(self._sock_client)
-        iface_reader_thread = SockServerInterfaceReaderThread(
-            clients=self._clients,
-            iface=iface,
-            stopped=self._stopped,
-        )
-        iface_reader_thread.start()
-
-    def server_inform_start(self, sreq: "spb.ServerRequest") -> None:
-        request = sreq.inform_start
-        stream_id = request._info.stream_id
-        settings = SettingsStatic(request.settings)
-        self._mux.update_stream(stream_id, settings=settings)
-        self._mux.start_stream(stream_id)
-
-    def server_inform_attach(self, sreq: "spb.ServerRequest") -> None:
-        request = sreq.inform_attach
-        stream_id = request._info.stream_id
-
-        self._clients.add_client(self._sock_client)
-        inform_attach_response = spb.ServerInformAttachResponse()
-        inform_attach_response.settings.CopyFrom(
-            self._mux._streams[stream_id]._settings._proto,
-        )
-        response = spb.ServerResponse(
-            request_id=sreq.request_id,
-            inform_attach_response=inform_attach_response,
-        )
-        self._sock_client.send_server_response(response)
-
-    def server_record_communicate(self, sreq: "spb.ServerRequest") -> None:
-        self._put_record(sreq.record_communicate)
-
-    def server_record_publish(self, sreq: "spb.ServerRequest") -> None:
-        self._put_record(sreq.record_publish)
-
-    def _put_record(self, record: "pb.Record") -> None:
-        # encode relay information so the right socket picks up the data
-        record.control.relay_id = self._sock_client._sockid
-        stream_id = record._info.stream_id
-
-        try:
-            iface = self._mux.get_stream(stream_id).interface
-
-        except KeyError:
-            # We should log the error but cannot because it may print to console
-            # due to how logging is set up. This error usually happens if
-            # a record is sent when no run is active, but during this time the
-            # logger prints to the console.
-            pass
-
-        else:
-            assert iface.record_q
-            iface.record_q.put(record)
-
-    def server_inform_finish(self, sreq: "spb.ServerRequest") -> None:
-        request = sreq.inform_finish
-        stream_id = request._info.stream_id
-        self._mux.drop_stream(stream_id)
-
-    def server_inform_teardown(self, sreq: "spb.ServerRequest") -> None:
-        request = sreq.inform_teardown
-        exit_code = request.exit_code
-        self._mux.teardown(exit_code)
-
-
-class SockAcceptThread(threading.Thread):
-    _sock: socket.socket
-    _mux: StreamMux
-    _stopped: "Event"
-    _clients: ClientDict
-
-    def __init__(self, sock: socket.socket, mux: StreamMux) -> None:
-        self._sock = sock
-        self._mux = mux
-        self._stopped = mux._get_stopped_event()
-        threading.Thread.__init__(self)
-        self.name = "SockAcceptThr"
-        self._clients = ClientDict()
-
-    def run(self) -> None:
-        read_threads = []
-
-        while not self._stopped.is_set():
-            try:
-                conn, addr = self._sock.accept()
-            except ConnectionAbortedError:
-                break
-            except OSError:
-                # on shutdown
-                break
-            sr = SockServerReadThread(conn=conn, mux=self._mux, clients=self._clients)
-            sr.start()
-            read_threads.append(sr)
-
-        for rt in read_threads:
-            rt.stop()
-
-
-class DebugThread(threading.Thread):
-    def __init__(self, mux: "StreamMux") -> None:
-        threading.Thread.__init__(self)
-        self.daemon = True
-        self.name = "DebugThr"
-
-    def run(self) -> None:
-        while True:
-            time.sleep(30)
-            for thread in threading.enumerate():
-                wandb.termwarn(f"DEBUG: {thread.name}")
-
-
-class SocketServer:
-    _mux: StreamMux
-    _address: str
-    _port: int
-    _sock: socket.socket
-
-    def __init__(self, mux: Any, address: str, port: int) -> None:
-        self._mux = mux
-        self._address = address
-        self._port = port
-        # This is the server socket that we accept new connections from
-        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-
-    def _bind(self) -> None:
-        self._sock.bind((self._address, self._port))
-        self._port = self._sock.getsockname()[1]
-
-    @property
-    def port(self) -> int:
-        return self._port
-
-    def start(self) -> None:
-        self._bind()
-        self._sock.listen(5)
-        self._thread = SockAcceptThread(sock=self._sock, mux=self._mux)
-        self._thread.start()
-        # Note: Uncomment to figure out what thread is not exiting properly
-        # self._dbg_thread = DebugThread(mux=self._mux)
-        # self._dbg_thread.start()
-
-    def stop(self) -> None:
-        if self._sock:
-            # we need to stop the SockAcceptThread
-            try:
-                # TODO(jhr): consider a more graceful shutdown in the future
-                # socket.shutdown() is a more heavy handed approach to interrupting socket.accept()
-                # in the future we might want to consider a more graceful shutdown which would involve setting
-                # a threading Event and then initiating one last connection just to close down the thread
-                # The advantage of the heavy handed approach is that it does not depend on the threads functioning
-                # properly, that is, if something has gone wrong, we probably want to use this hammer to shut things down
-                self._sock.shutdown(socket.SHUT_RDWR)
-            except OSError:
-                pass
-            self._sock.close()

wandb/sdk/service/service.py

@@ -1,252 +0,0 @@
-"""Reliably launch and connect to backend server process (wandb service).
-
-Backend server process can be connected to using tcp sockets transport.
-"""
-
-import datetime
-import os
-import pathlib
-import platform
-import shutil
-import subprocess
-import sys
-import tempfile
-import time
-from typing import TYPE_CHECKING, Any, Dict, Optional
-
-from wandb import _sentry
-from wandb.env import (
-    core_debug,
-    dcgm_profiling_enabled,
-    error_reporting_enabled,
-    is_require_legacy_service,
-)
-from wandb.errors import Error, WandbCoreNotAvailableError
-from wandb.errors.term import termlog, termwarn
-from wandb.util import get_core_path, get_module
-
-from . import _startup_debug, port_file
-
-if TYPE_CHECKING:
-    from wandb.sdk.wandb_settings import Settings
-
-
-class ServiceStartProcessError(Error):
-    """Raised when a known error occurs when launching wandb service."""
-
-
-class ServiceStartTimeoutError(Error):
-    """Raised when service start times out."""
-
-
-class ServiceStartPortError(Error):
-    """Raised when service start fails to find a port."""
-
-
-class _Service:
-    _settings: "Settings"
-    _sock_port: Optional[int]
-    _internal_proc: Optional[subprocess.Popen]
-    _startup_debug_enabled: bool
-
-    def __init__(
-        self,
-        settings: "Settings",
-    ) -> None:
-        self._settings = settings
-        self._stub = None
-        self._sock_port = None
-        self._internal_proc = None
-        self._startup_debug_enabled = _startup_debug.is_enabled()
-
-        _sentry.configure_scope(tags=dict(settings), process_context="service")
-
-    def _startup_debug_print(self, message: str) -> None:
-        if not self._startup_debug_enabled:
-            return
-        _startup_debug.print_message(message)
-
-    def _wait_for_ports(
-        self, fname: str, proc: Optional[subprocess.Popen] = None
-    ) -> None:
-        """Wait for the service to write the port file and then read it.
-
-        Args:
-            fname: The path to the port file.
-            proc: The process to wait for.
-
-        Raises:
-            ServiceStartTimeoutError: If the service takes too long to start.
-            ServiceStartPortError: If the service writes an invalid port file or unable to read it.
-            ServiceStartProcessError: If the service process exits unexpectedly.
-
-        """
-        time_max = time.monotonic() + self._settings.x_service_wait
-        while time.monotonic() < time_max:
-            if proc and proc.poll():
-                # process finished
-                # define these variables for sentry context grab:
-                # command = proc.args
-                # sys_executable = sys.executable
-                # which_python = shutil.which("python3")
-                # proc_out = proc.stdout.read()
-                # proc_err = proc.stderr.read()
-                context = dict(
-                    command=proc.args,
-                    sys_executable=sys.executable,
-                    which_python=shutil.which("python3"),
-                    proc_out=proc.stdout.read() if proc.stdout else "",
-                    proc_err=proc.stderr.read() if proc.stderr else "",
-                )
-                raise ServiceStartProcessError(
-                    f"The wandb service process exited with {proc.returncode}. "
-                    "Ensure that `sys.executable` is a valid python interpreter. "
-                    "You can override it with the `_executable` setting "
-                    "or with the `WANDB_X_EXECUTABLE` environment variable."
-                    f"\n{context}",
-                    context=context,
-                )
-            if not os.path.isfile(fname):
-                time.sleep(0.2)
-                continue
-            try:
-                pf = port_file.PortFile()
-                pf.read(fname)
-                if not pf.is_valid:
-                    time.sleep(0.2)
-                    continue
-                self._sock_port = pf.sock_port
-            except Exception as e:
-                # todo: point at the docs. this could be due to a number of reasons,
-                #  for example, being unable to write to the port file etc.
-                raise ServiceStartPortError(
-                    f"Failed to allocate port for wandb service: {e}."
-                )
-            return
-        raise ServiceStartTimeoutError(
-            "Timed out waiting for wandb service to start after "
-            f"{self._settings.x_service_wait} seconds. "
-            "Try increasing the timeout with the `_service_wait` setting."
-        )
-
-    def _launch_server(self) -> None:
-        """Launch server and set ports."""
-        # References for starting processes
-        # - https://github.com/wandb/wandb/blob/archive/old-cli/wandb/__init__.py
-        # - https://stackoverflow.com/questions/1196074/how-to-start-a-background-process-in-python
-        self._startup_debug_print("launch")
-
-        kwargs: Dict[str, Any] = dict(close_fds=True)
-        # flags to handle keyboard interrupt signal that is causing a hang
-        if platform.system() == "Windows":
-            kwargs.update(creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)  # type: ignore [attr-defined]
-        else:
-            kwargs.update(start_new_session=True)
-
-        pid = str(os.getpid())
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            fname = os.path.join(tmpdir, f"port-{pid}.txt")
-
-            executable = self._settings.x_executable
-            exec_cmd_list = [executable, "-m"]
-
-            service_args = []
-
-            if not is_require_legacy_service():
-                try:
-                    core_path = get_core_path()
-                except WandbCoreNotAvailableError as e:
-                    _sentry.reraise(e)
-
-                service_args.extend([core_path])
-
-                if not error_reporting_enabled():
-                    service_args.append("--no-observability")
-
-                if core_debug(default="False"):
-                    service_args.extend(["--log-level", "-4"])
-
-                if dcgm_profiling_enabled():
-                    service_args.append("--enable-dcgm-profiling")
-
-                exec_cmd_list = []
-            else:
-                service_args.extend(["wandb", "service", "--debug"])
-                termwarn(
-                    "Using legacy-service, which is deprecated. If this is"
-                    " unintentional, you can fix it by ensuring you do not call"
-                    " `wandb.require('legacy-service')` and do not set the"
-                    " WANDB_X_REQUIRE_LEGACY_SERVICE environment"
-                    " variable."
-                )
-
-            service_args += [
-                "--port-filename",
-                fname,
-                "--pid",
-                pid,
-            ]
-
-            if os.environ.get("WANDB_SERVICE_PROFILE") == "memray":
-                _ = get_module(
-                    "memray",
-                    required=(
-                        "wandb service memory profiling requires memray, "
-                        "install with `pip install memray`"
-                    ),
-                )
-
-                time_tag = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-                output_file = f"wandb_service.memray.{time_tag}.bin"
-                cli_executable = (
-                    pathlib.Path(__file__).parent.parent.parent.parent
-                    / "tools"
-                    / "cli.py"
-                )
-                exec_cmd_list = [
-                    executable,
-                    "-m",
-                    "memray",
-                    "run",
-                    "-o",
-                    output_file,
-                ]
-                service_args[0] = str(cli_executable)
-                termlog(
-                    f"wandb service memory profiling enabled, output file: {output_file}"
-                )
-                termlog(
-                    f"Convert to flamegraph with: `python -m memray flamegraph {output_file}`"
-                )
-
-            try:
-                internal_proc = subprocess.Popen(
-                    exec_cmd_list + service_args,  # type: ignore[arg-type]
-                    env=os.environ,
-                    **kwargs,
-                )
-            except Exception as e:
-                _sentry.reraise(e)
-
-            self._startup_debug_print("wait_ports")
-            try:
-                self._wait_for_ports(fname, proc=internal_proc)
-            except Exception as e:
-                _sentry.reraise(e)
-            self._startup_debug_print("wait_ports_done")
-            self._internal_proc = internal_proc
-        self._startup_debug_print("launch_done")
-
-    def start(self) -> None:
-        self._launch_server()
-
-    @property
-    def sock_port(self) -> Optional[int]:
-        return self._sock_port
-
-    def join(self) -> int:
-        ret = 0
-        if self._internal_proc:
-            ret = self._internal_proc.wait()
-        return ret