wandb 0.18.0rc1__py3-none-win_amd64.whl → 0.18.2__py3-none-win_amd64.whl

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

@@ -15,18 +15,12 @@ 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.env import core_debug, 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
@@ -35,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
 
@@ -69,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
@@ -162,9 +145,6 @@ class _Service:
 
             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 = []
 
@@ -176,16 +156,12 @@ class _Service:
 
                 service_args.extend([core_path])
 
-                if not core_error_reporting_enabled(default="True"):
+                if not error_reporting_enabled():
                     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."
@@ -201,7 +177,6 @@ class _Service:
                 "--pid",
                 pid,
             ]
-            service_args.append("--serve-sock")
 
             if os.environ.get("WANDB_SERVICE_PROFILE") == "memray":
                 _ = get_module(
@@ -260,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/service/streams.py

@@ -317,7 +317,6 @@ class StreamMux:
         # These could be done in parallel in the future
         for _sid, stream in started_streams.items():
             # dispatch all our final requests
-            server_info_handle = stream.interface.deliver_request_server_info()
             poll_exit_handle = stream.interface.deliver_poll_exit()
             final_summary_handle = stream.interface.deliver_get_summary()
             sampled_history_handle = stream.interface.deliver_request_sampled_history()
@@ -327,11 +326,6 @@ class StreamMux:
             assert result
             internal_messages_response = result.response.internal_messages_response
 
-            # wait for them, it's ok to do this serially but this can be improved
-            result = server_info_handle.wait(timeout=-1)
-            assert result
-            server_info_response = result.response.server_info_response
-
             result = poll_exit_handle.wait(timeout=-1)
             assert result
             poll_exit_response = result.response.poll_exit_response
@@ -348,7 +342,6 @@ class StreamMux:
                 sampled_history=sampled_history,
                 final_summary=final_summary,
                 poll_exit_response=poll_exit_response,
-                server_info_response=server_info_response,
                 internal_messages_response=internal_messages_response,
                 settings=stream._settings,  # type: ignore
                 printer=printer,

wandb/sdk/wandb_init.py

@@ -174,12 +174,19 @@ class _WandbInit:
 
         # we add this logic to be backward compatible with the old behavior of disable
         # where it would disable the service if the mode was set to disabled
+        # TODO: use the regular settins object to handle this
         mode = kwargs.get("mode")
         settings_mode = (kwargs.get("settings") or {}).get("mode") or os.environ.get(
-            "WANDB_MODE"
+            wandb.env.MODE
         )
-        _disable_service = mode == "disabled" or settings_mode == "disabled"
-        setup_settings = {"_disable_service": _disable_service}
+        settings__disable_service = (kwargs.get("settings") or {}).get(
+            "_disable_service"
+        ) or os.environ.get(wandb.env._DISABLE_SERVICE)
+
+        setup_settings = {
+            "mode": mode or settings_mode,
+            "_disable_service": settings__disable_service,
+        }
 
         self._wl = wandb_setup.setup(settings=setup_settings)
         # Make sure we have a logger setup (might be an early logger)
@@ -648,9 +655,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
@@ -658,15 +665,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
@@ -732,7 +744,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
@@ -823,7 +835,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()
@@ -850,9 +862,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
@@ -927,33 +940,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)