wandb 0.20.1rc20250604__py3-none-any.whl → 0.21.0__py3-none-any.whl

wandb/sdk/lib/{service_connection.py → service/service_connection.py}

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import atexit
-import os
 from typing import Callable
 
 from wandb.proto import wandb_internal_pb2 as pb
@@ -11,77 +10,42 @@ from wandb.sdk import wandb_settings
 from wandb.sdk.interface.interface import InterfaceBase
 from wandb.sdk.interface.interface_sock import InterfaceSock
 from wandb.sdk.interface.router_sock import MessageSockRouter
-from wandb.sdk.lib import service_token
 from wandb.sdk.lib.exit_hooks import ExitHooks
 from wandb.sdk.lib.sock_client import SockClient, SockClientClosedError
 from wandb.sdk.mailbox import HandleAbandonedError, Mailbox, MailboxClosedError
-from wandb.sdk.service import service
 
-
-class WandbServiceConnectionError(Exception):
-    """Raised on failure to connect to the service process."""
+from . import service_process, service_token
 
 
 class WandbAttachFailedError(Exception):
-    """Raised if attaching to a run fails."""
+    """Failed to attach to a run."""
 
 
 def connect_to_service(
     settings: wandb_settings.Settings,
 ) -> ServiceConnection:
-    """Connects to the service process, starting one up if necessary."""
-    conn = _try_connect_to_existing_service()
-    if conn:
-        return conn
-
-    return _start_and_connect_service(settings)
-
-
-def _try_connect_to_existing_service() -> ServiceConnection | None:
-    """Attempts to connect to an existing service process."""
-    token = service_token.get_service_token()
-    if not token:
-        return None
-
-    # Only localhost sockets are supported below.
-    assert token.host == "localhost"
-    client = SockClient()
+    """Connect to the service process, starting one up if necessary."""
+    token = service_token.from_env()
 
-    try:
-        # TODO: This may block indefinitely if the service is unhealthy.
-        client.connect(token.port)
-
-    except Exception as e:
-        raise WandbServiceConnectionError(
-            "Failed to connect to internal service."
-        ) from e
-
-    return ServiceConnection(client=client, proc=None)
+    if token:
+        return ServiceConnection(client=token.connect(), proc=None)
+    else:
+        return _start_and_connect_service(settings)
 
 
 def _start_and_connect_service(
     settings: wandb_settings.Settings,
 ) -> ServiceConnection:
-    """Starts a service process and returns a connection to it.
+    """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._Service(settings)
-    proc.start()
-
-    port = proc.sock_port
-    assert port
-    client = SockClient()
-    client.connect(port)
-
-    service_token.set_service_token(
-        parent_pid=os.getpid(),
-        transport="tcp",
-        host="localhost",
-        port=port,
-    )
+    proc = service_process.start(settings)
+
+    client = proc.token.connect()
+    proc.token.save_to_env()
 
     hooks = ExitHooks()
     hooks.hook()
@@ -106,7 +70,7 @@ class ServiceConnection:
     def __init__(
         self,
         client: SockClient,
-        proc: service._Service | None,
+        proc: service_process.ServiceProcess | None,
         cleanup: Callable[[], None] | None = None,
     ):
         """Returns a new ServiceConnection.
@@ -132,7 +96,7 @@ class ServiceConnection:
         return InterfaceSock(self._client, self._mailbox, stream_id=stream_id)
 
     def send_record(self, record: pb.Record) -> None:
-        """Sends data to the service."""
+        """Send data to the service."""
         self._client.send_record_publish(record)
 
     def inform_init(
@@ -140,14 +104,14 @@ class ServiceConnection:
         settings: wandb_settings_pb2.Settings,
         run_id: str,
     ) -> None:
-        """Sends an init request to the service."""
+        """Send an init request to the service."""
         request = spb.ServerInformInitRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
         self._client.send_server_request(spb.ServerRequest(inform_init=request))
 
     def inform_finish(self, run_id: str) -> None:
-        """Sends an finish request to the service."""
+        """Send an finish request to the service."""
         request = spb.ServerInformFinishRequest()
         request._info.stream_id = run_id
         self._client.send_server_request(spb.ServerRequest(inform_finish=request))
@@ -156,7 +120,7 @@ class ServiceConnection:
         self,
         attach_id: str,
     ) -> wandb_settings_pb2.Settings:
-        """Sends an attach request to the service.
+        """Send an attach request to the service.
 
         Raises a WandbAttachFailedError if attaching is not possible.
         """
@@ -188,7 +152,7 @@ class ServiceConnection:
         settings: wandb_settings_pb2.Settings,
         run_id: str,
     ) -> None:
-        """Sends a start request to the service."""
+        """Send a start request to the service."""
         request = spb.ServerInformStartRequest()
         request.settings.CopyFrom(settings)
         request._info.stream_id = run_id
@@ -221,7 +185,7 @@ class ServiceConnection:
             return None
 
         # Clear the service token to prevent new connections to the process.
-        service_token.clear_service_token()
+        service_token.clear_service_in_env()
 
         self._client.send_server_request(
             spb.ServerRequest(

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,164 @@
+from __future__ import annotations
+
+import abc
+import os
+import re
+import socket
+
+from typing_extensions import final, override
+
+from wandb import env
+from wandb.sdk.lib.service import ipc_support
+from wandb.sdk.lib.sock_client import SockClient
+
+_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) -> SockClient:
+        """Connect to the service process.
+
+        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) -> SockClient:
+        if not ipc_support.SUPPORTS_UNIX:
+            raise WandbServiceConnectionError("AF_UNIX socket not supported")
+
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+
+        try:
+            # TODO: This may block indefinitely if the service is unhealthy.
+            sock.connect(self._path)
+        except Exception as e:
+            raise WandbServiceConnectionError(
+                f"Failed to connect to service on socket {self._path}",
+            ) from e
+
+        return SockClient(sock)
+
+    @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) -> SockClient:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+        try:
+            # TODO: This may block indefinitely if the service is unhealthy.
+            sock.connect(("localhost", self._port))
+        except Exception as e:
+            raise WandbServiceConnectionError(
+                f"Failed to connect to service on port {self._port}",
+            ) from e
+
+        return SockClient(sock)
+
+    @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/sock_client.py

@@ -79,17 +79,17 @@ class SockBuffer:
 
 
 class SockClient:
-    _sock: socket.socket
-    _sockid: str
-    _retry_delay: float
-    _lock: "threading.Lock"
-    _bufsize: int
-    _buffer: SockBuffer
-
     # current header is magic byte "W" followed by 4 byte length of the message
     HEADLEN = 1 + 4
 
-    def __init__(self) -> None:
+    def __init__(self, sock: socket.socket) -> None:
+        """Create a SockClient.
+
+        Args:
+            sock: A connected socket.
+        """
+        self._sock = sock
+
         # TODO: use safe uuid's (python3.7+) or emulate this
         self._sockid = uuid.uuid4().hex
         self._retry_delay = 0.1
@@ -97,10 +97,6 @@ class SockClient:
         self._bufsize = 4096
         self._buffer = SockBuffer()
 
-    def connect(self, port: int) -> None:
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.connect(("localhost", port))
-        self._sock = s
         self._detect_bufsize()
 
     def _detect_bufsize(self) -> None:

wandb/sdk/wandb_init.py

@@ -740,7 +740,6 @@ class _WandbInit:
         drun._config.update(config.sweep_no_artifacts)
         drun._config.update(config.base_no_artifacts)
         drun.summary = SummaryDisabled()  # type: ignore
-        drun._Run__metadata = wandb.sdk.wandb_metadata.Metadata()
 
         # methods
         drun.log = lambda data, *_, **__: drun.summary.update(data)  # type: ignore[method-assign]
@@ -925,8 +924,6 @@ class _WandbInit:
                 tel.feature.flow_control_disabled = True
             if settings.x_flow_control_custom:
                 tel.feature.flow_control_custom = True
-            if not settings.x_require_legacy_service:
-                tel.feature.core = True
             if settings._shared:
                 wandb.termwarn(
                     "The `shared` mode feature is experimental and may change. "
@@ -1367,7 +1364,7 @@ def init(  # noqa: C901
             the UI.
             If resuming a run, the tags provided here will replace any existing
             tags. To add tags to a resumed run without overwriting the current
-            tags, use `run.tags += ["new_tag"]` after calling `run = wandb.init()`.
+            tags, use `run.tags += ("new_tag",)` after calling `run = wandb.init()`.
         config: Sets `wandb.config`, a dictionary-like object for storing input
             parameters to your run, such as model hyperparameters or data
             preprocessing settings.
@@ -1621,4 +1618,3 @@ def init(  # noqa: C901
         # Need to build delay into this sentry capture because our exit hooks
         # mess with sentry's ability to send out errors before the program ends.
         wandb._sentry.reraise(e)
-        raise AssertionError()  # should never get here

wandb/sdk/wandb_require.py

@@ -11,13 +11,10 @@ Example:
 
 from __future__ import annotations
 
-import os
 from typing import Iterable
 
 import wandb
-from wandb.env import _REQUIRE_LEGACY_SERVICE
 from wandb.errors import UnsupportedError
-from wandb.sdk import wandb_run
 
 
 class _Requires:
@@ -33,22 +30,21 @@ class _Requires:
     def require_require(self) -> None:
         pass
 
-    def _require_service(self) -> None:
-        wandb.teardown = wandb._teardown  # type: ignore
-        wandb.attach = wandb._attach  # type: ignore
-        wandb_run.Run.detach = wandb_run.Run._detach  # type: ignore
-
     def require_service(self) -> None:
-        self._require_service()
+        # Legacy no-op kept solely for backward compatibility:
+        # some integrations (e.g. PyTorch Lightning) still call
+        # `wandb.require('service')`, which routes here.
+        wandb.termwarn(
+            "`wandb.require('service')` is a no-op as it is now the default behavior."
+        )
 
     def require_core(self) -> None:
+        # Legacy no-op kept solely for backward compatibility:
+        # many public codebases still call `wandb.require('core')`.
         wandb.termwarn(
-            "`wandb.require('core')` is redundant as it is now the default behavior."
+            "`wandb.require('core')` is a no-op as it is now the default behavior."
         )
 
-    def require_legacy_service(self) -> None:
-        os.environ[_REQUIRE_LEGACY_SERVICE] = "true"
-
     def apply(self) -> None:
         """Call require_* method for supported features."""
         last_message: str = ""
@@ -64,7 +60,6 @@ class _Requires:
             func()
 
         if last_message:
-            wandb.termwarn("Supported requirements are: `legacy-service`, `service`.")
             raise UnsupportedError(last_message)
 
 
@@ -91,10 +86,3 @@ def require(
 
     f = _Requires(features=features)
     f.apply()
-
-
-def _import_module_hook() -> None:
-    """On wandb import, setup anything needed based on parent process require calls."""
-    # TODO: optimize by caching which pids this has been done for or use real import hooks
-    # TODO: make this more generic, but for now this works
-    require("service")