wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/sdk/service/service.py

@@ -1,9 +1,10 @@
 """Reliably launch and connect to backend server process (wandb service).
 
-Backend server process can be connected to using tcp sockets or grpc transport.
+Backend server process can be connected to using tcp sockets transport.
 """
-
+import datetime
 import os
+import pathlib
 import platform
 import shutil
 import subprocess
@@ -12,7 +13,8 @@ import tempfile
 import time
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from wandb import _sentry
+from wandb import _minimum_nexus_version, _sentry, termlog
+from wandb.env import error_reporting_enabled
 from wandb.errors import Error
 from wandb.util import get_module
 
@@ -42,13 +44,22 @@ class ServiceStartPortError(Error):
     pass
 
 
+def _check_nexus_version_compatibility(nexus_version: str) -> None:
+    """Checks if the installed nexus version is compatible with the wandb version."""
+    from pkg_resources import parse_version
+
+    if parse_version(nexus_version) < parse_version(_minimum_nexus_version):
+        raise ImportError(
+            f"Requires wandb-core version {_minimum_nexus_version} or later, "
+            f"but you have {nexus_version}. Run `pip install --upgrade wandb[nexus]` to upgrade."
+        )
+
+
 class _Service:
     _settings: "Settings"
-    _grpc_port: Optional[int]
     _sock_port: Optional[int]
     _service_interface: ServiceInterface
     _internal_proc: Optional[subprocess.Popen]
-    _use_grpc: bool
     _startup_debug_enabled: bool
 
     def __init__(
@@ -57,25 +68,15 @@ class _Service:
     ) -> None:
         self._settings = settings
         self._stub = None
-        self._grpc_port = 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")
 
-        # Temporary setting to allow use of grpc so that we can keep
-        # that code from rotting during the transition
-        self._use_grpc = self._settings._service_transport == "grpc"
-
-        # current code only supports grpc or socket server implementation, in the
+        # current code only supports socket server implementation, in the
         # future we might be able to support both
-        if self._use_grpc:
-            from .service_grpc import ServiceGrpcInterface
-
-            self._service_interface = ServiceGrpcInterface()
-        else:
-            self._service_interface = ServiceSockInterface()
+        self._service_interface = ServiceSockInterface()
 
     def _startup_debug_print(self, message: str) -> None:
         if not self._startup_debug_enabled:
@@ -130,7 +131,6 @@ class _Service:
                 if not pf.is_valid:
                     time.sleep(0.2)
                     continue
-                self._grpc_port = pf.grpc_port
                 self._sock_port = pf.sock_port
             except Exception as e:
                 # todo: point at the docs. this could be due to a number of reasons,
@@ -186,8 +186,11 @@ class _Service:
                         "wandb_core",
                         required="The nexus experiment requires the wandb_core module.",
                     )
+                    _check_nexus_version_compatibility(wandb_nexus.__version__)
                     nexus_path = wandb_nexus.get_nexus_path()
                 service_args.extend([nexus_path])
+                if not error_reporting_enabled():
+                    service_args.append("--no-observability")
                 exec_cmd_list = []
             else:
                 service_args.extend(["wandb", "service"])
@@ -199,10 +202,40 @@ class _Service:
                 pid,
                 "--debug",
             ]
-            if self._use_grpc:
-                service_args.append("--serve-grpc")
-            else:
-                service_args.append("--serve-sock")
+            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,
@@ -211,6 +244,7 @@ class _Service:
                 )
             except Exception as e:
                 _sentry.reraise(e)
+
             self._startup_debug_print("wait_ports")
             try:
                 self._wait_for_ports(fname, proc=internal_proc)
@@ -223,10 +257,6 @@ class _Service:
     def start(self) -> None:
         self._launch_server()
 
-    @property
-    def grpc_port(self) -> Optional[int]:
-        return self._grpc_port
-
     @property
     def sock_port(self) -> Optional[int]:
         return self._sock_port

wandb/sdk/service/service_base.py

@@ -1,6 +1,6 @@
 """Base service abstract class.
 
-Derived classes for grpc and socket service interfaces classes should implement
+Derived classes for socket service interfaces classes should implement
 abstract methods.
 """
 

wandb/sdk/service/streams.py

@@ -2,7 +2,7 @@
 
 StreamThread: Thread that runs internal.wandb_internal()
 StreamRecord: All the external state for the internal thread (queues, etc)
-StreamAction: Lightweight record for stream ops for thread safety with grpc
+StreamAction: Lightweight record for stream ops for thread safety
 StreamMux: Container for dictionary of stream threads per runid
 """
 import functools
@@ -321,6 +321,11 @@ class StreamMux:
             server_info_handle = stream.interface.deliver_request_server_info()
             final_summary_handle = stream.interface.deliver_get_summary()
             sampled_history_handle = stream.interface.deliver_request_sampled_history()
+            internal_messages_handle = stream.interface.deliver_internal_messages()
+
+            result = internal_messages_handle.wait(timeout=-1)
+            assert result
+            internal_messages_response = result.response.internal_messages_response
             job_info_handle = stream.interface.deliver_request_job_info()
 
             # wait for them, it's ok to do this serially but this can be improved
@@ -345,10 +350,11 @@ class StreamMux:
             job_info = result.response.job_info_response
 
             Run._footer(
-                sampled_history,
-                final_summary,
-                poll_exit_response,
-                server_info_response,
+                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,
                 job_info=job_info,
                 settings=stream._settings,  # type: ignore
                 printer=printer,

wandb/sdk/verify/verify.py

@@ -314,7 +314,7 @@ def check_artifacts() -> bool:
     # test checksum
     sing_art_dir = "./verify_sing_art"
     alias = "sing_art1"
-    name = "sing-artys"
+    name = nice_id("sing-artys")
     singular_art = artifact_with_path_or_paths(name, singular=True)
     cont_test, download_artifact, failed_test_strings = log_use_download_artifact(
         singular_art, alias, name, sing_art_dir, failed_test_strings, False
@@ -332,7 +332,7 @@ def check_artifacts() -> bool:
     # test manifest and digest
     multi_art_dir = "./verify_art"
     alias = "art1"
-    name = "my-artys"
+    name = nice_id("my-artys")
     art1 = artifact_with_path_or_paths(name, "./verify_art_dir", singular=False)
     cont_test, download_artifact, failed_test_strings = log_use_download_artifact(
         art1, alias, name, multi_art_dir, failed_test_strings, True

wandb/sdk/wandb_init.py

@@ -136,6 +136,8 @@ class _WandbInit:
 
         self.deprecated_features_used: Dict[str, str] = dict()
 
+        self._require_nexus = os.environ.get("WANDB_REQUIRE_NEXUS") == "True"
+
     def _setup_printer(self, settings: Settings) -> None:
         if self.printer:
             return
@@ -427,7 +429,8 @@ class _WandbInit:
             return None
 
         # Attempt to save the code on every execution
-        if self.notebook.save_ipynb():  # type: ignore
+        # todo(nexus): remove nexus check once incremental artifact is supported
+        if not self._require_nexus and self.notebook.save_ipynb():  # type: ignore
             assert self.run is not None
             res = self.run.log_code(root=None)
             logger.info("saved code: %s", res)  # type: ignore
@@ -445,7 +448,8 @@ class _WandbInit:
         assert self.notebook
         ipython = self.notebook.shell
         self.notebook.save_history()
-        if self.notebook.save_ipynb():
+        # todo(nexus): remove nexus check once incremental artifact is supported
+        if not self._require_nexus and self.notebook.save_ipynb():
             assert self.run is not None
             res = self.run.log_code(root=None)
             logger.info("saved code and history: %s", res)  # type: ignore
@@ -571,6 +575,7 @@ class _WandbInit:
         logger.info(
             f"wandb.init called with sweep_config: {self.sweep_config}\nconfig: {self.config}"
         )
+
         if self.settings._noop:
             return self._make_run_disabled()
         if self.settings.reinit or (
@@ -636,6 +641,7 @@ class _WandbInit:
         with telemetry.context(run=run, obj=self._init_telemetry_obj) as tel:
             tel.cli_version = wandb.__version__
             tel.python_version = platform.python_version()
+            tel.platform = f"{platform.system()}-{platform.machine()}".lower()
             hf_version = _huggingface_version()
             if hf_version:
                 tel.huggingface_version = hf_version

wandb/sdk/wandb_manager.py

@@ -1,6 +1,6 @@
 """Manage wandb processes.
 
-Create a grpc manager channel.
+Create a manager channel.
 """
 
 import atexit
@@ -35,7 +35,7 @@ class ManagerConnectionRefusedError(ManagerConnectionError):
 
 class _ManagerToken:
     _version = "2"
-    _supported_transports = {"grpc", "tcp"}
+    _supported_transports = {"tcp"}
     _token_str: str
     _pid: int
     _transport: str
@@ -126,7 +126,6 @@ class _Manager:
             raise ManagerConnectionError(f"Connection to wandb service failed: {e}")
 
     def __init__(self, settings: "Settings") -> None:
-        # TODO: warn if user doesn't have grpc installed
         from wandb.sdk.service import service
 
         self._settings = settings
@@ -134,21 +133,12 @@ class _Manager:
         self._hooks = None
 
         self._service = service._Service(settings=self._settings)
-
-        # Temporary setting to allow use of grpc so that we can keep
-        # that code from rotting during the transition
-        use_grpc = self._settings._service_transport == "grpc"
-
         token = _ManagerToken.from_environment()
         if not token:
             self._service.start()
             host = "localhost"
-            if use_grpc:
-                transport = "grpc"
-                port = self._service.grpc_port
-            else:
-                transport = "tcp"
-                port = self._service.sock_port
+            transport = "tcp"
+            port = self._service.sock_port
             assert port
             token = _ManagerToken.from_params(transport=transport, host=host, port=port)
             token.set_environment()