wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/lib/service/service_connection.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import atexit
+from typing import Callable
+
+from wandb.proto import wandb_server_pb2 as spb
+from wandb.proto import wandb_settings_pb2
+from wandb.sdk import wandb_settings
+from wandb.sdk.interface.interface import InterfaceBase
+from wandb.sdk.interface.interface_sock import InterfaceSock
+from wandb.sdk.lib import asyncio_manager
+from wandb.sdk.lib.exit_hooks import ExitHooks
+from wandb.sdk.lib.service.service_client import ServiceClient
+from wandb.sdk.mailbox import HandleAbandonedError, MailboxClosedError
+
+from . import service_process, service_token
+
+
+class WandbAttachFailedError(Exception):
+    """Failed to attach to a run."""
+
+
+def connect_to_service(
+    asyncer: asyncio_manager.AsyncioManager,
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Connect to the service process, starting one up if necessary."""
+    token = service_token.from_env()
+
+    if token:
+        return ServiceConnection(
+            client=token.connect(asyncer=asyncer),
+            proc=None,
+        )
+    else:
+        return _start_and_connect_service(asyncer, settings)
+
+
+def _start_and_connect_service(
+    asyncer: asyncio_manager.AsyncioManager,
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Start a service process and returns a connection to it.
+
+    An atexit hook is registered to tear down the service process and wait for
+    it to complete. The hook does not run in processes started using the
+    multiprocessing module.
+    """
+    proc = service_process.start(settings)
+
+    client = proc.token.connect(asyncer=asyncer)
+    proc.token.save_to_env()
+
+    hooks = ExitHooks()
+    hooks.hook()
+
+    def teardown_atexit():
+        conn.teardown(hooks.exit_code)
+
+    conn = ServiceConnection(
+        client=client,
+        proc=proc,
+        cleanup=lambda: atexit.unregister(teardown_atexit),
+    )
+
+    atexit.register(teardown_atexit)
+
+    return conn
+
+
+class ServiceConnection:
+    """A connection to the W&B internal service process."""
+
+    def __init__(
+        self,
+        client: ServiceClient,
+        proc: service_process.ServiceProcess | None,
+        cleanup: Callable[[], None] | None = None,
+    ):
+        """Returns a new ServiceConnection.
+
+        Args:
+            mailbox: The mailbox to use for all communication over the socket.
+            router: A handle to the thread that reads from the socket and
+                updates the mailbox.
+            client: A socket that's connected to the service.
+            proc: The service process if we own it, or None otherwise.
+            cleanup: A callback to run on teardown before doing anything.
+        """
+        self._client = client
+        self._proc = proc
+        self._torn_down = False
+        self._cleanup = cleanup
+
+    def make_interface(self, stream_id: str) -> InterfaceBase:
+        """Returns an interface for communicating with the service."""
+        return InterfaceSock(self._client, stream_id=stream_id)
+
+    def inform_init(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Send an init request to the service."""
+        request = spb.ServerInformInitRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.publish(spb.ServerRequest(inform_init=request))
+
+    def inform_finish(self, run_id: str) -> None:
+        """Send an finish request to the service."""
+        request = spb.ServerInformFinishRequest()
+        request._info.stream_id = run_id
+        self._client.publish(spb.ServerRequest(inform_finish=request))
+
+    def inform_attach(
+        self,
+        attach_id: str,
+    ) -> wandb_settings_pb2.Settings:
+        """Send an attach request to the service.
+
+        Raises a WandbAttachFailedError if attaching is not possible.
+        """
+        request = spb.ServerRequest()
+        request.inform_attach._info.stream_id = attach_id
+
+        try:
+            handle = self._client.deliver(request)
+            response = handle.wait_or(timeout=10)
+
+        except (MailboxClosedError, HandleAbandonedError):
+            raise WandbAttachFailedError(
+                "Failed to attach: the service process is not running.",
+            ) from None
+
+        except TimeoutError:
+            raise WandbAttachFailedError(
+                "Failed to attach because the run does not belong to"
+                " the current service process, or because the service"
+                " process is busy (unlikely)."
+            ) from None
+
+        else:
+            return response.inform_attach_response.settings
+
+    def inform_start(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Send a start request to the service."""
+        request = spb.ServerInformStartRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.publish(spb.ServerRequest(inform_start=request))
+
+    def teardown(self, exit_code: int) -> int | None:
+        """Close the connection.
+
+        Stop reading responses on the connection, and if this connection owns
+        the service process, send a teardown message and wait for it to shut
+        down.
+
+        This may only be called once.
+
+        Returns:
+            The exit code of the service process, or None if the process was
+            not owned by this connection.
+        """
+        if self._torn_down:
+            raise AssertionError("Already torn down.")
+        self._torn_down = True
+
+        if self._cleanup:
+            self._cleanup()
+
+        if not self._proc:
+            return None
+
+        # Clear the service token to prevent new connections to the process.
+        service_token.clear_service_in_env()
+
+        self._client.publish(
+            spb.ServerRequest(
+                inform_teardown=spb.ServerInformTeardownRequest(
+                    exit_code=exit_code,
+                )
+            ),
+        )
+        self._client.close()
+
+        return self._proc.join()

wandb/sdk/lib/service/service_port_file.py

@@ -0,0 +1,105 @@
+"""Module for figuring out how to connect to the service process."""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import re
+import subprocess
+import time
+
+import wandb
+
+from . import ipc_support, service_token
+
+# Time functions are monkeypatched in unit tests.
+_MONOTONIC = time.monotonic
+_SLEEP = time.sleep
+
+
+class ServicePollForTokenError(wandb.Error):
+    """Failed to discover how to connect to the service."""
+
+
+def poll_for_token(
+    file: pathlib.Path,
+    proc: subprocess.Popen,
+    *,
+    timeout: float,
+) -> service_token.ServiceToken:
+    """Poll the 'port' file to discover how to connect to the service.
+
+    Args:
+        file: The file path that should eventually contain this information.
+        proc: The process that's supposed to generate the file.
+            If the process dies, this raises an error.
+        timeout: A timeout in seconds after which to raise an error.
+
+    Returns:
+        A token specifying how to connect to the service process.
+
+    Raises:
+        ServicePollForTokenError: if the service process dies, a timeout
+            occurs, or there's an issue reading the port file.
+    """
+    end_time = _MONOTONIC() + timeout
+
+    while _MONOTONIC() < end_time:
+        if (code := proc.poll()) is not None:
+            raise ServicePollForTokenError(
+                f"wandb-core exited with code {code}",
+                context={
+                    "command": proc.args,
+                    "proc_out": proc.stdout.read() if proc.stdout else "",
+                    "proc_err": proc.stderr.read() if proc.stderr else "",
+                },
+            )
+
+        if token := _poll_once(file):
+            return token
+
+        _SLEEP(max(0, min(0.2, end_time - _MONOTONIC())))
+
+    raise ServicePollForTokenError(
+        f"Failed to read port info after {timeout} seconds.",
+    )
+
+
+_UNIX_NAME_RE = re.compile(r"unix=(.+)")
+_TCP_PORT_RE = re.compile(r"sock=(\d+)")
+
+
+def _poll_once(file: pathlib.Path) -> service_token.ServiceToken | None:
+    """Try to read the port file.
+
+    Returns:
+        A connection token on success. Otherwise, returns None.
+
+    Raises:
+        ServicePollForTokenError: if the file contains no known
+            connection method.
+    """
+    try:
+        text = file.read_text()
+    except OSError:
+        return None
+
+    lines = text.splitlines()
+    if lines[-1] != "EOF":
+        return None
+
+    for line in lines:
+        if ipc_support.SUPPORTS_UNIX and (match := _UNIX_NAME_RE.fullmatch(line)):
+            return service_token.UnixServiceToken(
+                parent_pid=os.getpid(),
+                path=match.group(1),
+            )
+        elif match := _TCP_PORT_RE.fullmatch(line):
+            return service_token.TCPServiceToken(
+                parent_pid=os.getpid(),
+                port=int(match.group(1)),
+            )
+
+    raise ServicePollForTokenError(
+        f"No known connection method in {file}:\n{text}",
+    )

wandb/sdk/lib/service/service_process.py

@@ -0,0 +1,111 @@
+"""Module for starting up the service process (wandb-core)."""
+
+from __future__ import annotations
+
+import os
+import pathlib
+import platform
+import subprocess
+import tempfile
+from typing import TYPE_CHECKING
+
+from wandb import _sentry
+from wandb.env import core_debug, dcgm_profiling_enabled, error_reporting_enabled
+from wandb.errors import WandbCoreNotAvailableError
+from wandb.sdk.lib.service import ipc_support
+from wandb.util import get_core_path
+
+from . import service_port_file, service_token
+
+if TYPE_CHECKING:
+    from wandb.sdk.wandb_settings import Settings
+
+
+def start(settings: Settings) -> ServiceProcess:
+    """Start the internal service process.
+
+    Returns:
+        A handle to the process.
+    """
+    _sentry.configure_scope(tags=dict(settings), process_context="service")
+
+    try:
+        return _launch_server(settings)
+    except Exception as e:
+        _sentry.reraise(e)
+
+
+class ServiceProcess:
+    """A handle to a process running the internal service."""
+
+    def __init__(
+        self,
+        *,
+        connection_token: service_token.ServiceToken,
+        process: subprocess.Popen,
+    ) -> None:
+        self._token = connection_token
+        self._process = process
+
+    @property
+    def token(self) -> service_token.ServiceToken:
+        """A token for connecting to the process."""
+        return self._token
+
+    def join(self) -> int:
+        """Wait for the process to end and return its exit code."""
+        return self._process.wait()
+
+
+def _launch_server(settings: Settings) -> ServiceProcess:
+    """Launch server and set ports."""
+    if platform.system() == "Windows":
+        creationflags: int = subprocess.CREATE_NEW_PROCESS_GROUP  # type: ignore[attr-defined]
+        start_new_session = False
+    else:
+        creationflags = 0
+        start_new_session = True
+
+    pid = str(os.getpid())
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        port_file = pathlib.Path(tmpdir, f"port-{pid}.txt")
+        service_args: list[str] = []
+
+        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")
+
+        service_args.extend(["--port-filename", str(port_file)])
+        service_args.extend(["--pid", pid])
+
+        if not ipc_support.SUPPORTS_UNIX:
+            service_args.append("--listen-on-localhost")
+
+        proc = subprocess.Popen(
+            service_args,
+            env=os.environ,
+            close_fds=True,
+            creationflags=creationflags,
+            start_new_session=start_new_session,
+        )
+
+        token = service_port_file.poll_for_token(
+            port_file,
+            proc,
+            timeout=settings.x_service_wait,
+        )
+
+        return ServiceProcess(connection_token=token, process=proc)

wandb/sdk/lib/service/service_token.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import abc
+import asyncio
+import os
+import re
+
+from typing_extensions import final, override
+
+from wandb import env
+from wandb.sdk.lib import asyncio_manager
+from wandb.sdk.lib.service import ipc_support
+
+from .service_client import ServiceClient
+
+_CURRENT_VERSION = "3"
+
+# Token formats:
+_UNIX_TOKEN_RE = re.compile(rf"{_CURRENT_VERSION}-(\d+)-unix-(.+)")
+_TCP_TOKEN_RE = re.compile(rf"{_CURRENT_VERSION}-(\d+)-tcp-localhost-(\d+)")
+
+
+class WandbServiceConnectionError(Exception):
+    """Failed to connect to the service process."""
+
+
+def clear_service_in_env() -> None:
+    """Clear the environment variable that stores the service token."""
+    os.environ.pop(env.SERVICE, None)
+
+
+def from_env() -> ServiceToken | None:
+    """Read 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
+
+    if unix_token := UnixServiceToken.from_env_string(token):
+        return unix_token
+    if tcp_token := TCPServiceToken.from_env_string(token):
+        return tcp_token
+
+    raise ValueError(f"Failed to parse {env.SERVICE}={token!r}")
+
+
+class ServiceToken(abc.ABC):
+    """A way of connecting to a running service process."""
+
+    @abc.abstractmethod
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
+        """Connect to the service process.
+
+        Args:
+            asyncer: A started AsyncioManager for asyncio operations.
+
+        Returns:
+            A socket object for communicating with the service.
+
+        Raises:
+            WandbServiceConnectionError: on failure to connect.
+        """
+
+    def save_to_env(self) -> None:
+        """Save the token in this process's environment variables."""
+        os.environ[env.SERVICE] = self._as_env_string()
+
+    @abc.abstractmethod
+    def _as_env_string(self) -> str:
+        """Returns a string representation of this token."""
+
+
+@final
+class UnixServiceToken(ServiceToken):
+    """Connects to the service using a Unix domain socket."""
+
+    def __init__(self, *, parent_pid: int, path: str) -> None:
+        self._parent_pid = parent_pid
+        self._path = path
+
+    @override
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
+        if not ipc_support.SUPPORTS_UNIX:
+            raise WandbServiceConnectionError("AF_UNIX socket not supported")
+
+        try:
+            # TODO: This may block indefinitely if the service is unhealthy.
+            reader, writer = asyncer.run(
+                lambda: asyncio.open_unix_connection(self._path),
+            )
+        except Exception as e:
+            raise WandbServiceConnectionError(
+                f"Failed to connect to service on socket {self._path}",
+            ) from e
+
+        return ServiceClient(asyncer, reader, writer)
+
+    @override
+    def _as_env_string(self):
+        return "-".join(
+            (
+                _CURRENT_VERSION,
+                str(self._parent_pid),
+                "unix",
+                str(self._path),
+            )
+        )
+
+    @staticmethod
+    def from_env_string(token: str) -> UnixServiceToken | None:
+        """Returns a Unix service token parsed from the env var."""
+        match = _UNIX_TOKEN_RE.fullmatch(token)
+        if not match:
+            return None
+
+        parent_pid, path = match.groups()
+        return UnixServiceToken(parent_pid=int(parent_pid), path=path)
+
+
+@final
+class TCPServiceToken(ServiceToken):
+    """Connects to the service using TCP over a localhost socket."""
+
+    def __init__(self, *, parent_pid: int, port: int) -> None:
+        self._parent_pid = parent_pid
+        self._port = port
+
+    @override
+    def connect(
+        self,
+        *,
+        asyncer: asyncio_manager.AsyncioManager,
+    ) -> ServiceClient:
+        try:
+            # TODO: This may block indefinitely if the service is unhealthy.
+            reader, writer = asyncer.run(
+                lambda: asyncio.open_connection("localhost", self._port),
+            )
+        except Exception as e:
+            raise WandbServiceConnectionError(
+                f"Failed to connect to service on port {self._port}",
+            ) from e
+
+        return ServiceClient(asyncer, reader, writer)
+
+    @override
+    def _as_env_string(self):
+        return "-".join(
+            (
+                _CURRENT_VERSION,
+                str(self._parent_pid),
+                "tcp",
+                "localhost",
+                str(self._port),
+            )
+        )
+
+    @staticmethod
+    def from_env_string(token: str) -> TCPServiceToken | None:
+        """Returns a TCP service token parsed from the env var."""
+        match = _TCP_TOKEN_RE.fullmatch(token)
+        if not match:
+            return None
+
+        parent_pid, port = match.groups()
+        return TCPServiceToken(parent_pid=int(parent_pid), port=int(port))

wandb/sdk/lib/sparkline.py

@@ -0,0 +1,44 @@
+#
+# From pysparklines (BSD License): https://pypi.python.org/pypi/pysparklines
+
+import math
+from typing import List, Union
+
+spark_chars = "▁▂▃▄▅▆▇█"
+
+
+# math.isfinite doesn't exist in python2, so provider our own
+def isfinite(f):
+    return not (math.isinf(f) or math.isnan(f))
+
+
+def sparkify(series: List[Union[float, int]]) -> str:
+    """Convert <series> to a sparkline string.
+
+    Example:
+    >>> sparkify([0.5, 1.2, 3.5, 7.3, 8.0, 12.5, 13.2, 15.0, 14.2, 11.8, 6.1, 1.9])
+    u'▁▁▂▄▅▇▇██▆▄▂'
+
+    >>> sparkify([1, 1, -2, 3, -5, 8, -13])
+    u'▆▆▅▆▄█▁'
+
+    Raises ValueError if input data cannot be converted to float.
+    Raises TypeError if series is not an iterable.
+    """
+    series = [float(i) for i in series]
+    finite_series = [x for x in series if isfinite(x)]
+    if not finite_series:
+        return ""
+    minimum = min(finite_series)
+    maximum = max(finite_series)
+    data_range = maximum - minimum
+    if data_range == 0.0:
+        # Graph a baseline if every input value is equal.
+        return "".join([spark_chars[0] if isfinite(x) else " " for x in series])
+    coefficient = (len(spark_chars) - 1.0) / data_range
+    return "".join(
+        [
+            spark_chars[int(round((x - minimum) * coefficient))] if isfinite(x) else " "
+            for x in series
+        ]
+    )