wandb 0.18.0rc1__py3-none-macosx_10_13_x86_64.whl

wandb/sdk/service/server_sock.py

@@ -0,0 +1,276 @@
+import queue
+import socket
+import threading
+import time
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.internal.settings_static import SettingsStatic
+
+from ..lib import tracelog
+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:
+        assert self._iface.relay_q
+        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
+            tracelog.log_message_dequeue(result, self._iface.relay_q)
+            sockid = result.control.relay_id
+            assert sockid
+            sock_client = self._clients.get_client(sockid)
+            assert sock_client
+            sresp = spb.ServerResponse()
+            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(inform_attach_response=inform_attach_response)
+        self._sock_client.send_server_response(response)
+        iface = self._mux.get_stream(stream_id).interface
+
+        assert iface
+
+    def server_record_communicate(self, sreq: "spb.ServerRequest") -> None:
+        record = sreq.record_communicate
+        # 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
+        iface = self._mux.get_stream(stream_id).interface
+        assert iface.record_q
+        iface.record_q.put(record)
+
+    def server_record_publish(self, sreq: "spb.ServerRequest") -> None:
+        record = sreq.record_publish
+        # 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
+        iface = self._mux.get_stream(stream_id).interface
+        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():
+                print(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 doesnt 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

@@ -0,0 +1,271 @@
+"""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, termlog
+from wandb.env import (
+    core_debug,
+    core_error_reporting_enabled,
+    is_require_legacy_service,
+)
+from wandb.errors import Error, WandbCoreNotAvailableError
+from wandb.sdk.lib.wburls import wburls
+from wandb.util import get_core_path, get_module
+
+from . import _startup_debug, port_file
+from .service_base import ServiceInterface
+from .service_sock import ServiceSockInterface
+
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
+
+class ServiceStartProcessError(Error):
+    """Raised when a known error occurs when launching wandb service."""
+
+    pass
+
+
+class ServiceStartTimeoutError(Error):
+    """Raised when service start times out."""
+
+    pass
+
+
+class ServiceStartPortError(Error):
+    """Raised when service start fails to find a port."""
+
+    pass
+
+
+class _Service:
+    _settings: "Settings"
+    _sock_port: Optional[int]
+    _service_interface: ServiceInterface
+    _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")
+
+        # current code only supports socket server implementation, in the
+        # future we might be able to support both
+        self._service_interface = ServiceSockInterface()
+
+    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._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__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._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._executable
+            exec_cmd_list = [executable, "-m"]
+            # Add coverage collection if needed
+            if os.environ.get("YEA_RUN_COVERAGE") and os.environ.get("COVERAGE_RCFILE"):
+                exec_cmd_list += ["coverage", "run", "-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 core_error_reporting_enabled(default="True"):
+                    service_args.append("--no-observability")
+
+                if core_debug(default="False"):
+                    service_args.append("--debug")
+
+                trace_filename = os.environ.get("_WANDB_TRACE")
+                if trace_filename is not None:
+                    service_args.extend(["--trace", trace_filename])
+
+                exec_cmd_list = []
+                termlog(
+                    "Using wandb-core as the SDK backend."
+                    f" Please refer to {wburls.get('wandb_core')} for more information.",
+                    repeat=False,
+                )
+            else:
+                service_args.extend(["wandb", "service", "--debug"])
+
+            service_args += [
+                "--port-filename",
+                fname,
+                "--pid",
+                pid,
+            ]
+            service_args.append("--serve-sock")
+
+            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,
+                    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
+
+    @property
+    def service_interface(self) -> ServiceInterface:
+        return self._service_interface
+
+    def join(self) -> int:
+        ret = 0
+        if self._internal_proc:
+            ret = self._internal_proc.wait()
+        return ret

wandb/sdk/service/service_base.py

@@ -0,0 +1,50 @@
+"""Base service abstract class.
+
+Derived classes for socket service interfaces classes should implement
+abstract methods.
+"""
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+from wandb.proto import wandb_server_pb2 as spb
+
+if TYPE_CHECKING:
+    from wandb.proto import wandb_settings_pb2
+
+
+class ServiceInterface:
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def get_transport(self) -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_inform_init(
+        self, settings: "wandb_settings_pb2.Settings", run_id: str
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_inform_start(
+        self, settings: "wandb_settings_pb2.Settings", run_id: str
+    ) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_inform_attach(self, attach_id: str) -> spb.ServerInformAttachResponse:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_inform_finish(self, run_id: Optional[str] = None) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_inform_teardown(self, exit_code: int) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _svc_connect(self, port: int) -> None:
+        raise NotImplementedError

wandb/sdk/service/service_sock.py

@@ -0,0 +1,70 @@
+"""socket service.
+
+Implement ServiceInterface for socket transport.
+"""
+
+from typing import TYPE_CHECKING, Optional
+
+from wandb.proto import wandb_server_pb2 as spb
+
+from ..lib.sock_client import SockClient
+from .service_base import ServiceInterface
+
+if TYPE_CHECKING:
+    from wandb.proto import wandb_settings_pb2
+
+
+class ServiceSockInterface(ServiceInterface):
+    _sock_client: SockClient
+
+    def __init__(self) -> None:
+        self._sock_client = SockClient()
+
+    def get_transport(self) -> str:
+        return "tcp"
+
+    def _get_sock_client(self) -> SockClient:
+        return self._sock_client
+
+    def _svc_connect(self, port: int) -> None:
+        self._sock_client.connect(port=port)
+
+    def _svc_inform_init(
+        self, settings: "wandb_settings_pb2.Settings", run_id: str
+    ) -> None:
+        inform_init = spb.ServerInformInitRequest()
+        inform_init.settings.CopyFrom(settings)
+        inform_init._info.stream_id = run_id
+        assert self._sock_client
+        self._sock_client.send(inform_init=inform_init)
+
+    def _svc_inform_start(
+        self, settings: "wandb_settings_pb2.Settings", run_id: str
+    ) -> None:
+        inform_start = spb.ServerInformStartRequest()
+        inform_start.settings.CopyFrom(settings)
+        inform_start._info.stream_id = run_id
+        assert self._sock_client
+        self._sock_client.send(inform_start=inform_start)
+
+    def _svc_inform_finish(self, run_id: Optional[str] = None) -> None:
+        assert run_id
+        inform_finish = spb.ServerInformFinishRequest()
+        inform_finish._info.stream_id = run_id
+
+        assert self._sock_client
+        self._sock_client.send(inform_finish=inform_finish)
+
+    def _svc_inform_attach(self, attach_id: str) -> spb.ServerInformAttachResponse:
+        inform_attach = spb.ServerInformAttachRequest()
+        inform_attach._info.stream_id = attach_id
+
+        assert self._sock_client
+        response = self._sock_client.send_and_recv(inform_attach=inform_attach)
+        return response.inform_attach_response
+
+    def _svc_inform_teardown(self, exit_code: int) -> None:
+        inform_teardown = spb.ServerInformTeardownRequest(exit_code=exit_code)
+
+        assert self._sock_client
+        self._sock_client.send(inform_teardown=inform_teardown)