wandb 0.18.1__py3-none-win_amd64.whl → 0.18.3__py3-none-win_amd64.whl

wandb/sdk/lib/hashutil.py

@@ -1,19 +1,22 @@
+from __future__ import annotations
+
 import base64
 import hashlib
 import mmap
-import os
 import sys
-from pathlib import Path
-from typing import NewType, Union
+from typing import TYPE_CHECKING, NewType
 
 from wandb.sdk.lib.paths import StrPath
 
+if TYPE_CHECKING:
+    import _hashlib  # type: ignore[import-not-found]
+
 ETag = NewType("ETag", str)
 HexMD5 = NewType("HexMD5", str)
 B64MD5 = NewType("B64MD5", str)
 
 
-def _md5(data: bytes = b"") -> "hashlib._Hash":
+def _md5(data: bytes = b"") -> _hashlib.HASH:
     """Allow FIPS-compliant md5 hash when supported."""
     if sys.version_info >= (3, 9):
         return hashlib.md5(data, usedforsecurity=False)
@@ -25,7 +28,7 @@ def md5_string(string: str) -> B64MD5:
     return _b64_from_hasher(_md5(string.encode("utf-8")))
 
 
-def _b64_from_hasher(hasher: "hashlib._Hash") -> B64MD5:
+def _b64_from_hasher(hasher: _hashlib.HASH) -> B64MD5:
     return B64MD5(base64.b64encode(hasher.digest()).decode("ascii"))
 
 
@@ -33,7 +36,7 @@ def b64_to_hex_id(string: B64MD5) -> HexMD5:
     return HexMD5(base64.standard_b64decode(string).hex())
 
 
-def hex_to_b64_id(encoded_string: Union[str, bytes]) -> B64MD5:
+def hex_to_b64_id(encoded_string: str | bytes) -> B64MD5:
     if isinstance(encoded_string, bytes):
         encoded_string = encoded_string.decode("utf-8")
     as_str = bytes.fromhex(encoded_string)
@@ -48,15 +51,34 @@ def md5_file_hex(*paths: StrPath) -> HexMD5:
     return HexMD5(_md5_file_hasher(*paths).hexdigest())
 
 
-def _md5_file_hasher(*paths: StrPath) -> "hashlib._Hash":
+_KB: int = 1_024
+_CHUNKSIZE: int = 128 * _KB
+"""Chunk size (in bytes) for iteratively reading from file, if needed."""
+
+
+def _md5_file_hasher(*paths: StrPath) -> _hashlib.HASH:
     md5_hash = _md5()
 
-    for path in sorted(Path(p) for p in paths):
-        with path.open("rb") as f:
-            if os.stat(f.fileno()).st_size <= 1024 * 1024:
-                md5_hash.update(f.read())
-            else:
+    # Note: We use str paths (instead of pathlib.Path objs) for minor perf improvements.
+    for path in sorted(map(str, paths)):
+        with open(path, "rb") as f:
+            try:
                 with mmap.mmap(f.fileno(), length=0, access=mmap.ACCESS_READ) as mview:
                     md5_hash.update(mview)
+            except OSError:
+                # This occurs if the mmap-ed file is on a different/mounted filesystem,
+                # so we'll fall back on a less performant implementation.
+
+                # Note: At the time of implementation, the walrus operator `:=`
+                # is avoided to maintain support for users on python 3.7.
+                # Consider revisiting once 3.7 support is no longer needed.
+                chunk = f.read(_CHUNKSIZE)
+                while chunk:
+                    md5_hash.update(chunk)
+                    chunk = f.read(_CHUNKSIZE)
+            except ValueError:
+                # This occurs when mmap-ing an empty file, which can be skipped.
+                # See: https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Modules/mmapmodule.c#L1589
+                pass
 
     return md5_hash

wandb/sdk/lib/service_connection.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import atexit
+import os
+from typing import Callable
+
+from wandb.proto import wandb_internal_pb2 as pb
+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 service_token
+from wandb.sdk.lib.exit_hooks import ExitHooks
+from wandb.sdk.lib.mailbox import Mailbox
+from wandb.sdk.lib.sock_client import SockClient, SockClientTimeoutError
+from wandb.sdk.service import service
+
+
+class WandbServiceNotOwnedError(Exception):
+    """Raised when the current process does not own the service process."""
+
+
+class WandbServiceConnectionError(Exception):
+    """Raised on failure to connect to the service process."""
+
+
+class WandbAttachFailedError(Exception):
+    """Raised if attaching to a run fails."""
+
+
+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:
+    """Attemps 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()
+
+    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)
+
+
+def _start_and_connect_service(
+    settings: wandb_settings.Settings,
+) -> ServiceConnection:
+    """Starts 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,
+    )
+
+    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: SockClient,
+        proc: service._Service | None,
+        cleanup: Callable[[], None] | None = None,
+    ):
+        """Returns a new ServiceConnection.
+
+        Args:
+            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, mailbox: Mailbox) -> InterfaceBase:
+        """Returns an interface for communicating with the service."""
+        return InterfaceSock(self._client, mailbox)
+
+    def send_record(self, record: pb.Record) -> None:
+        """Sends data to the service."""
+        self._client.send_record_publish(record)
+
+    def inform_init(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends an init request to the service."""
+        request = spb.ServerInformInitRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(inform_init=request)
+
+    def inform_finish(self, run_id: str) -> None:
+        """Sends an finish request to the service."""
+        request = spb.ServerInformFinishRequest()
+        request._info.stream_id = run_id
+        self._client.send(inform_finish=request)
+
+    def inform_attach(
+        self,
+        attach_id: str,
+    ) -> wandb_settings_pb2.Settings:
+        """Sends an attach request to the service.
+
+        Raises a WandbAttachFailedError if attaching is not possible.
+        """
+        request = spb.ServerInformAttachRequest()
+        request._info.stream_id = attach_id
+
+        try:
+            response = self._client.send_and_recv(inform_attach=request)
+            return response.inform_attach_response.settings
+        except SockClientTimeoutError:
+            raise WandbAttachFailedError(
+                "Could not attach because the run does not belong to"
+                " the current service process, or because the service"
+                " process is busy (unlikely)."
+            )
+
+    def inform_start(
+        self,
+        settings: wandb_settings_pb2.Settings,
+        run_id: str,
+    ) -> None:
+        """Sends a start request to the service."""
+        request = spb.ServerInformStartRequest()
+        request.settings.CopyFrom(settings)
+        request._info.stream_id = run_id
+        self._client.send(inform_start=request)
+
+    def teardown(self, exit_code: int) -> int:
+        """Shuts down the service process and returns its exit code.
+
+        This may only be called once.
+
+        Returns:
+            The exit code of the service process.
+
+        Raises:
+            WandbServiceNotOwnedError: If the current process did not start
+                the service process.
+        """
+        if not self._proc:
+            raise WandbServiceNotOwnedError(
+                "Cannot tear down service started by different process",
+            )
+
+        assert not self._torn_down
+        self._torn_down = True
+
+        if self._cleanup:
+            self._cleanup()
+
+        # Clear the service token to prevent new connections from being made.
+        service_token.clear_service_token()
+
+        self._client.send(
+            inform_teardown=spb.ServerInformTeardownRequest(
+                exit_code=exit_code,
+            )
+        )
+
+        return self._proc.join()

wandb/sdk/lib/service_token.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import dataclasses
+import os
+
+from wandb import env
+
+_CURRENT_VERSION = "2"
+_SUPPORTED_TRANSPORTS = "tcp"
+
+
+def get_service_token() -> ServiceToken | None:
+    """Reads 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
+
+    parts = token.split("-")
+    if len(parts) != 5:
+        raise ValueError(f"Invalid token: {token}")
+
+    version, pid_str, transport, host, port_str = parts
+
+    if version != _CURRENT_VERSION:
+        raise ValueError(
+            f"Expected version {_CURRENT_VERSION},"
+            f" but got {version} (token={token})"
+        )
+    if transport not in _SUPPORTED_TRANSPORTS:
+        raise ValueError(
+            f"Unsupported transport: {transport} (token={token})",
+        )
+
+    try:
+        return ServiceToken(
+            version=version,
+            pid=int(pid_str),
+            transport=transport,
+            host=host,
+            port=int(port_str),
+        )
+    except ValueError as e:
+        raise ValueError(f"Invalid token: {token}") from e
+
+
+def set_service_token(parent_pid: int, transport: str, host: str, port: int) -> None:
+    """Stores a service token in an environment variable.
+
+    Args:
+        parent_pid: The process ID of the process that started the service.
+        transport: The transport used to communicate with the service.
+        host: The host part of the internet address on which the service
+            is listening (e.g. localhost).
+        port: The port the service is listening on.
+
+    Raises:
+        ValueError: If given an unsupported transport.
+    """
+    if transport not in _SUPPORTED_TRANSPORTS:
+        raise ValueError(f"Unsupported transport: {transport}")
+
+    os.environ[env.SERVICE] = "-".join(
+        (
+            _CURRENT_VERSION,
+            str(parent_pid),
+            transport,
+            host,
+            str(port),
+        )
+    )
+
+
+def clear_service_token() -> None:
+    """Clears the environment variable storing the service token."""
+    os.environ.pop(env.SERVICE, None)
+
+
+@dataclasses.dataclass(frozen=True)
+class ServiceToken:
+    """An identifier for a running service process."""
+
+    version: str
+    pid: int
+    transport: str
+    host: str
+    port: int

wandb/sdk/lib/sock_client.py

@@ -14,9 +14,11 @@ if TYPE_CHECKING:
 
 
 class SockClientClosedError(Exception):
-    """Socket has been closed."""
+    """Raised on operations on a closed socket."""
 
-    pass
+
+class SockClientTimeoutError(Exception):
+    """Raised if the server didn't respond before the timeout."""
 
 
 class SockBuffer:
@@ -182,8 +184,10 @@ class SockClient:
         # it should be relatively stable.
         # This pass would be solved as part of the fix in https://wandb.atlassian.net/browse/WB-8709
         response = self.read_server_response(timeout=1)
+
         if response is None:
-            raise Exception("No response")
+            raise SockClientTimeoutError("No response after 1 second.")
+
         return response
 
     def send(

wandb/sdk/service/server.py

@@ -20,7 +20,6 @@ class WandbServer:
     _pid: Optional[int]
     _sock_port: Optional[int]
     _debug: bool
-    _serve_sock: bool
     _sock_server: Optional[SocketServer]
     _startup_debug_enabled: bool
 
@@ -31,14 +30,12 @@ class WandbServer:
         address: Optional[str] = None,
         pid: Optional[int] = None,
         debug: bool = True,
-        serve_sock: bool = False,
     ) -> None:
         self._sock_port = sock_port
         self._port_fname = port_fname
         self._address = address
         self._pid = pid
         self._debug = debug
-        self._serve_sock = serve_sock
         self._sock_server = None
         self._startup_debug_enabled = _startup_debug.is_enabled()
 
@@ -97,7 +94,7 @@ class WandbServer:
             pid = str(self._pid or 0)
             transport = "s" if sock_port else "g"
             port = sock_port or 0
-            # this format is similar to wandb_manager token, but it's purely informative now
+            # 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})"
@@ -109,7 +106,7 @@ class WandbServer:
         self._setup_tracelog()
         mux = StreamMux()
         self._startup_debug_print("before_network")
-        sock_port = self._start_sock(mux=mux) if self._serve_sock else None
+        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")

wandb/sdk/service/service.py

@@ -21,8 +21,6 @@ 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
@@ -31,25 +29,18 @@ if TYPE_CHECKING:
 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
 
@@ -65,10 +56,6 @@ class _Service:
 
         _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
@@ -175,10 +162,6 @@ class _Service:
                 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."
@@ -194,7 +177,6 @@ class _Service:
                 "--pid",
                 pid,
             ]
-            service_args.append("--serve-sock")
 
             if os.environ.get("WANDB_SERVICE_PROFILE") == "memray":
                 _ = get_module(
@@ -253,10 +235,6 @@ class _Service:
     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:

wandb/sdk/wandb_init.py

@@ -578,6 +578,7 @@ class _WandbInit:
         ):
             setattr(drun, symbol, lambda *_, **__: None)  # type: ignore
         # attributes
+        drun._backend = None
         drun._step = 0
         drun._attach_id = None
         drun._run_obj = None
@@ -655,9 +656,9 @@ class _WandbInit:
                         f"Successfully finished last run (ID:{latest_run._run_id}). Initializing new run:<br/>"
                     )
         elif isinstance(wandb.run, Run):
-            manager = self._wl._get_manager()
+            service = self._wl.service
             # We shouldn't return a stale global run if we are in a new pid
-            if not manager or os.getpid() == wandb.run._init_pid:
+            if not service or os.getpid() == wandb.run._init_pid:
                 logger.info("wandb.init() called when a run is still active")
                 with telemetry.context() as tel:
                     tel.feature.init_return_run = True
@@ -665,15 +666,20 @@ class _WandbInit:
 
         logger.info("starting backend")
 
-        manager = self._wl._get_manager()
-        if manager:
-            logger.info("setting up manager")
-            manager._inform_init(
-                settings=self.settings.to_proto(), run_id=self.settings.run_id
+        service = self._wl.service
+        if service:
+            logger.info("sending inform_init request")
+            service.inform_init(
+                settings=self.settings.to_proto(),
+                run_id=self.settings.run_id,
             )
 
         mailbox = Mailbox()
-        backend = Backend(settings=self.settings, manager=manager, mailbox=mailbox)
+        backend = Backend(
+            settings=self.settings,
+            service=service,
+            mailbox=mailbox,
+        )
         backend.ensure_launched()
         logger.info("backend started and connected")
         # Make sure we are logged in
@@ -739,7 +745,7 @@ class _WandbInit:
             if os.environ.get(wandb.env._DISABLE_SERVICE):
                 tel.feature.service_disabled = True
 
-            if manager:
+            if service:
                 tel.feature.service = True
             if self.settings._flow_control_disabled:
                 tel.feature.flow_control_disabled = True
@@ -830,7 +836,7 @@ class _WandbInit:
 
         if error is not None:
             logger.error(f"encountered error: {error}")
-            if not manager:
+            if not service:
                 # Shutdown the backend and get rid of the logger
                 # we don't need to do console cleanup at this point
                 backend.cleanup()
@@ -857,9 +863,10 @@ class _WandbInit:
         logger.info("starting run threads in backend")
         # initiate run (stats and metadata probing)
 
-        if manager:
-            manager._inform_start(
-                settings=self.settings.to_proto(), run_id=self.settings.run_id
+        if service:
+            service.inform_start(
+                settings=self.settings.to_proto(),
+                run_id=self.settings.run_id,
             )
 
         assert backend.interface
@@ -934,33 +941,37 @@ def _attach(
     if logger is None:
         raise UsageError("logger is not initialized")
 
-    manager = _wl._get_manager()
-    response = manager._inform_attach(attach_id=attach_id) if manager else None
-    if response is None:
-        raise UsageError(f"Unable to attach to run {attach_id}")
+    service = _wl.service
+    if not service:
+        raise UsageError(f"Unable to attach to run {attach_id} (no service process)")
+
+    try:
+        attach_settings = service.inform_attach(attach_id=attach_id)
+    except Exception as e:
+        raise UsageError(f"Unable to attach to run {attach_id}") from e
 
     settings: Settings = copy.copy(_wl._settings)
 
     settings.update(
         {
             "run_id": attach_id,
-            "_start_time": response._start_time.value,
-            "_start_datetime": response._start_datetime.value,
-            "_offline": response._offline.value,
+            "_start_time": attach_settings._start_time.value,
+            "_start_datetime": attach_settings._start_datetime.value,
+            "_offline": attach_settings._offline.value,
         },
         source=Source.INIT,
     )
 
     # TODO: consolidate this codepath with wandb.init()
     mailbox = Mailbox()
-    backend = Backend(settings=settings, manager=manager, mailbox=mailbox)
+    backend = Backend(settings=settings, service=service, mailbox=mailbox)
     backend.ensure_launched()
     logger.info("attach backend started and connected")
 
     if run is None:
         run = Run(settings=settings)
     else:
-        run._init()
+        run._init(settings=settings)
     run._set_library(_wl)
     run._set_backend(backend)
     backend._hack_set_run(run)