wandb 0.22.0__py3-none-macosx_12_0_arm64.whl → 0.22.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/data_types/object_3d.py

@@ -138,7 +138,7 @@ def box3d(
     label: "Optional[str]" = None,
     score: "Optional[numeric]" = None,
 ) -> "Box3D":
-    """Returns a Box3D.
+    """A 3D bounding box. The box is specified by its center, size and orientation.
 
     Args:
         center: The center point of the box as a length-3 ndarray.
@@ -149,7 +149,72 @@ def box3d(
             r + xi + yj + zk.
         color: The box's color as an (r, g, b) tuple with 0 <= r,g,b <= 1.
         label: An optional label for the box.
-        score: An optional score for the box.
+        score: An optional score for the box. Typically used to indicate
+            the confidence of a detection.
+
+    Returns:
+        A Box3D object.
+
+    Example:
+    The following example creates a point cloud with 60 boxes rotating
+    around the X, Y and Z axes.
+
+    ```python
+    import wandb
+
+    import math
+    import numpy as np
+    from scipy.spatial.transform import Rotation
+
+
+    with wandb.init() as run:
+        run.log(
+            {
+                "points": wandb.Object3D.from_point_cloud(
+                    points=np.random.uniform(-5, 5, size=(100, 3)),
+                    boxes=[
+                        wandb.box3d(
+                            center=(0.3 * t - 3, 0, 0),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [t * math.pi / 10, 0, 0]
+                            ).as_quat(),
+                            color=(0.5 + t / 40, 0.5, 0.5),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ]
+                    + [
+                        wandb.box3d(
+                            center=(0, 0.3 * t - 3, 0.3),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [0, t * math.pi / 10, 0]
+                            ).as_quat(),
+                            color=(0.5, 0.5 + t / 40, 0.5),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ]
+                    + [
+                        wandb.box3d(
+                            center=(0.3, 0.3, 0.3 * t - 3),
+                            size=(0.1, 0.1, 0.1),
+                            orientation=Rotation.from_euler(
+                                "xyz", [0, 0, t * math.pi / 10]
+                            ).as_quat(),
+                            color=(0.5, 0.5, 0.5 + t / 40),
+                            label=f"box {t}",
+                            score=0.9,
+                        )
+                        for t in range(20)
+                    ],
+                ),
+            }
+        )
+    ```
     """
     try:
         import numpy as np

wandb/sdk/interface/interface.py

@@ -87,11 +87,38 @@ def file_enum_to_policy(enum: "pb.FilesItem.PolicyType.V") -> "PolicyName":
 
 
 class InterfaceBase:
+    """Methods for sending different types of Records to the service.
+
+    None of the methods may be called from an asyncio context other than
+    deliver_async().
+    """
+
     _drop: bool
 
     def __init__(self) -> None:
         self._drop = False
 
+    @abstractmethod
+    async def deliver_async(
+        self,
+        record: pb.Record,
+    ) -> MailboxHandle[pb.Result]:
+        """Send a record and create a handle to wait for the response.
+
+        The synchronous publish and deliver methods on this class cannot be
+        called in the asyncio thread because they block. Instead of having
+        an async copy of every method, this is a general method for sending
+        any kind of record in the asyncio thread.
+
+        Args:
+            record: The record to send. This method takes ownership of the
+                record and it must not be used afterward.
+
+        Returns:
+            A handle to wait for a response to the record.
+        """
+        raise NotImplementedError
+
     def publish_header(self) -> None:
         header = pb.HeaderRecord()
         self._publish_header(header)
@@ -392,9 +419,13 @@ class InterfaceBase:
             proto_manifest.manifest_file_path = path
             return proto_manifest
 
+        # Set storage policy on storageLayout (always V2) and storageRegion, only allow coreweave-us on wandb.ai for now.
+        # NOTE: the decode logic is NewManifestFromProto in core/pkg/artifacts/manifest.go
+        # The creation logic is in artifacts/_factories.py make_storage_policy
         for k, v in artifact_manifest.storage_policy.config().items() or {}.items():
             cfg = proto_manifest.storage_policy_config.add()
             cfg.key = k
+            # TODO: Why json.dumps when existing values are plain string? We want to send complex structure without defining the proto?
             cfg.value_json = json.dumps(v)
 
         for entry in sorted(artifact_manifest.entries.values(), key=lambda k: k.path):
@@ -1020,10 +1051,6 @@ class InterfaceBase:
     ) -> MailboxHandle[pb.Result]:
         raise NotImplementedError
 
-    @abstractmethod
-    def deliver_operation_stats(self) -> MailboxHandle[pb.Result]:
-        raise NotImplementedError
-
     def deliver_poll_exit(self) -> MailboxHandle[pb.Result]:
         poll_exit = pb.PollExitRequest()
         return self._deliver_poll_exit(poll_exit)

wandb/sdk/interface/interface_queue.py

@@ -8,12 +8,15 @@ import logging
 from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Optional
 
+from typing_extensions import override
+
 from .interface_shared import InterfaceShared
 
 if TYPE_CHECKING:
     from queue import Queue
 
     from wandb.proto import wandb_internal_pb2 as pb
+    from wandb.sdk.mailbox.mailbox_handle import MailboxHandle
 
 
 logger = logging.getLogger("wandb")
@@ -31,6 +34,13 @@ class InterfaceQueue(InterfaceShared):
         self._process = process
         super().__init__()
 
+    @override
+    async def deliver_async(
+        self,
+        record: "pb.Record",
+    ) -> "MailboxHandle[pb.Result]":
+        raise NotImplementedError
+
     def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
         if self._process and not self._process.is_alive():
             raise Exception("The wandb backend process has shutdown")

wandb/sdk/interface/interface_shared.py

@@ -87,7 +87,6 @@ class InterfaceShared(InterfaceBase):
         stop_status: Optional[pb.StopStatusRequest] = None,
         internal_messages: Optional[pb.InternalMessagesRequest] = None,
         network_status: Optional[pb.NetworkStatusRequest] = None,
-        operation_stats: Optional[pb.OperationStatsRequest] = None,
         poll_exit: Optional[pb.PollExitRequest] = None,
         partial_history: Optional[pb.PartialHistoryRequest] = None,
         sampled_history: Optional[pb.SampledHistoryRequest] = None,
@@ -129,8 +128,6 @@ class InterfaceShared(InterfaceBase):
             request.internal_messages.CopyFrom(internal_messages)
         elif network_status:
             request.network_status.CopyFrom(network_status)
-        elif operation_stats:
-            request.operations.CopyFrom(operation_stats)
         elif poll_exit:
             request.poll_exit.CopyFrom(poll_exit)
         elif partial_history:
@@ -424,10 +421,6 @@ class InterfaceShared(InterfaceBase):
         record = self._make_record(exit=exit_data)
         return self._deliver(record)
 
-    def deliver_operation_stats(self):
-        record = self._make_request(operation_stats=pb.OperationStatsRequest())
-        return self._deliver(record)
-
     def _deliver_poll_exit(
         self,
         poll_exit: pb.PollExitRequest,

wandb/sdk/interface/interface_sock.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
 from typing_extensions import override
 
 from wandb.proto import wandb_server_pb2 as spb
+from wandb.sdk.lib import asyncio_manager
 
 from .interface_shared import InterfaceShared
 
@@ -21,10 +22,12 @@ logger = logging.getLogger("wandb")
 class InterfaceSock(InterfaceShared):
     def __init__(
         self,
+        asyncer: asyncio_manager.AsyncioManager,
         client: ServiceClient,
         stream_id: str,
     ) -> None:
         super().__init__()
+        self._asyncer = asyncer
         self._client = client
         self._stream_id = stream_id
 
@@ -37,13 +40,16 @@ class InterfaceSock(InterfaceShared):
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
-        self._client.publish(request)
+        self._asyncer.run(lambda: self._client.publish(request))
 
-    @override
     def _deliver(self, record: pb.Record) -> MailboxHandle[pb.Result]:
+        return self._asyncer.run(lambda: self.deliver_async(record))
+
+    @override
+    async def deliver_async(self, record: pb.Record) -> MailboxHandle[pb.Result]:
         self._assign(record)
         request = spb.ServerRequest()
         request.record_publish.CopyFrom(record)
 
-        handle = self._client.deliver(request)
+        handle = await self._client.deliver(request)
         return handle.map(lambda response: response.result_communicate)

wandb/sdk/internal/_generated/__init__.py

@@ -1,15 +1,5 @@
 # Generated by ariadne-codegen
 
+__all__ = ["SERVER_FEATURES_QUERY_GQL", "ServerFeaturesQuery"]
 from .operations import SERVER_FEATURES_QUERY_GQL
-from .server_features_query import (
-    ServerFeaturesQuery,
-    ServerFeaturesQueryServerInfo,
-    ServerFeaturesQueryServerInfoFeatures,
-)
-
-__all__ = [
-    "SERVER_FEATURES_QUERY_GQL",
-    "ServerFeaturesQuery",
-    "ServerFeaturesQueryServerInfo",
-    "ServerFeaturesQueryServerInfoFeatures",
-]
+from .server_features_query import ServerFeaturesQuery

wandb/sdk/internal/job_builder.py

@@ -132,6 +132,7 @@ def get_min_supported_for_source_dict(
 
 class JobBuilder:
     _settings: SettingsStatic
+    _files_dir: str
     _metadatafile_path: Optional[str]
     _requirements_path: Optional[str]
     _config: Optional[Dict[str, Any]]
@@ -146,8 +147,26 @@ class JobBuilder:
     _verbose: bool
     _services: Dict[str, str]
 
-    def __init__(self, settings: SettingsStatic, verbose: bool = False):
+    def __init__(
+        self,
+        settings: SettingsStatic,
+        verbose: bool = False,
+        *,
+        files_dir: str,
+    ):
+        """Instantiate a JobBuilder.
+
+        Args:
+            settings: Parameters for the job builder.
+                In a run, this is the run's settings.
+                Otherwise, this is a set of undocumented parameters,
+                all of which should be made explicit like files_dir.
+            files_dir: The directory where to write files.
+                In a run, this should be the run's files directory.
+        """
         self._settings = settings
+        self._files_dir = files_dir
+
         self._metadatafile_path = None
         self._requirements_path = None
         self._config = None
@@ -460,9 +479,7 @@ class JobBuilder:
             )
             return None
 
-        if not os.path.exists(
-            os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME)
-        ):
+        if not os.path.exists(os.path.join(self._files_dir, REQUIREMENTS_FNAME)):
             self._log_if_verbose(
                 "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job",
                 "warn",
@@ -471,7 +488,7 @@ class JobBuilder:
         metadata = self._handle_metadata_file()
         if metadata is None:
             self._log_if_verbose(
-                f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
+                f"Ensure read and write access to run files dir: {self._files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables",
                 "warn",
             )
             return None
@@ -560,15 +577,15 @@ class JobBuilder:
             f.write(json.dumps(source_info, indent=4))
 
         artifact.add_file(
-            os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME),
+            os.path.join(self._files_dir, REQUIREMENTS_FNAME),
             name=FROZEN_REQUIREMENTS_FNAME,
         )
 
         if source_type == "repo":
             # add diff
-            if os.path.exists(os.path.join(self._settings.files_dir, DIFF_FNAME)):
+            if os.path.exists(os.path.join(self._files_dir, DIFF_FNAME)):
                 artifact.add_file(
-                    os.path.join(self._settings.files_dir, DIFF_FNAME),
+                    os.path.join(self._files_dir, DIFF_FNAME),
                     name=DIFF_FNAME,
                 )
 
@@ -619,8 +636,8 @@ class JobBuilder:
     def _handle_metadata_file(
         self,
     ) -> Optional[Dict]:
-        if os.path.exists(os.path.join(self._settings.files_dir, METADATA_FNAME)):
-            with open(os.path.join(self._settings.files_dir, METADATA_FNAME)) as f:
+        if os.path.exists(os.path.join(self._files_dir, METADATA_FNAME)):
+            with open(os.path.join(self._files_dir, METADATA_FNAME)) as f:
                 metadata: Dict = json.load(f)
             return metadata
 

wandb/sdk/internal/sender.py

@@ -311,7 +311,10 @@ class SendManager:
         self._output_raw_file = None
 
         # job builder
-        self._job_builder = JobBuilder(settings)
+        self._job_builder = JobBuilder(
+            settings,
+            files_dir=settings.files_dir,
+        )
 
         time_now = time.monotonic()
         self._debounce_config_time = time_now
@@ -343,7 +346,7 @@ class SendManager:
         publish_interface = InterfaceQueue(record_q=record_q)
         context_keeper = context.ContextKeeper()
         return SendManager(
-            settings=SettingsStatic(settings.to_proto()),
+            settings=SettingsStatic(dict(settings)),
             record_q=record_q,
             result_q=result_q,
             interface=publish_interface,

wandb/sdk/internal/settings_static.py

@@ -2,9 +2,7 @@ from __future__ import annotations
 
 from typing import Any, Iterable
 
-from wandb.proto import wandb_settings_pb2
-from wandb.sdk.lib import RunMoment
-from wandb.sdk.wandb_settings import CLIENT_ONLY_SETTINGS, Settings
+from wandb.sdk.wandb_settings import Settings
 
 
 class SettingsStatic(Settings):
@@ -14,87 +12,9 @@ class SettingsStatic(Settings):
     attributes or items.
     """
 
-    def __init__(self, proto: wandb_settings_pb2.Settings) -> None:
-        data = self._proto_to_dict(proto)
+    def __init__(self, data: dict[str, Any]) -> None:
         super().__init__(**data)
 
-    def _proto_to_dict(self, proto: wandb_settings_pb2.Settings) -> dict:
-        data = {}
-
-        exclude_fields = {
-            "model_config",
-            "model_fields",
-            "model_fields_set",
-            "__fields__",
-            "__model_fields_set",
-            "__pydantic_self__",
-            "__pydantic_initialised__",
-        }
-
-        fields = (
-            Settings.model_fields
-            if hasattr(Settings, "model_fields")
-            else Settings.__fields__
-        )  # type: ignore [attr-defined]
-
-        fields = {k: v for k, v in fields.items() if k not in exclude_fields}  # type: ignore [union-attr]
-
-        forks_specified: list[str] = []
-        for key in fields:
-            if key in CLIENT_ONLY_SETTINGS:
-                continue
-
-            value: Any = None
-
-            field_info = fields[key]
-            annotation = str(field_info.annotation)
-
-            if key == "_stats_open_metrics_filters":
-                # todo: it's an underscored field, refactor into
-                #  something more elegant?
-                # I'm really about this. It's ugly, but it works.
-                # Do not try to repeat this at home.
-                value_type = getattr(proto, key).WhichOneof("value")
-                if value_type == "sequence":
-                    value = list(getattr(proto, key).sequence.value)
-                elif value_type == "mapping":
-                    unpacked_mapping = {}
-                    for outer_key, outer_value in getattr(
-                        proto, key
-                    ).mapping.value.items():
-                        unpacked_inner = {}
-                        for inner_key, inner_value in outer_value.value.items():
-                            unpacked_inner[inner_key] = inner_value
-                        unpacked_mapping[outer_key] = unpacked_inner
-                    value = unpacked_mapping
-            elif key == "fork_from" or key == "resume_from":
-                value = getattr(proto, key)
-                if value.run:
-                    value = RunMoment(
-                        run=value.run, value=value.value, metric=value.metric
-                    )
-                    forks_specified.append(key)
-                else:
-                    value = None
-            else:
-                if proto.HasField(key):  # type: ignore [arg-type]
-                    value = getattr(proto, key).value
-                    # Convert to list if the field is a sequence
-                    if any(t in annotation for t in ("tuple", "Sequence", "list")):
-                        value = list(value)
-                else:
-                    value = None
-
-            if value is not None:
-                data[key] = value
-
-        if len(forks_specified) > 1:
-            raise ValueError(
-                "Only one of fork_from or resume_from can be specified, not both"
-            )
-
-        return data
-
     def __setattr__(self, name: str, value: object) -> None:
         raise AttributeError("Error: SettingsStatic is a readonly object")
 

wandb/sdk/launch/create_job.py

@@ -417,10 +417,11 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
     if job_source == "code":
         job_source = "artifact"
 
-    settings = wandb.Settings(x_files_dir=tmpdir, job_source=job_source)
+    settings = wandb.Settings(job_source=job_source)
     job_builder = JobBuilder(
         settings=settings,  # type: ignore
         verbose=True,
+        files_dir=tmpdir,
     )
     job_builder._partial = True
     # never allow notebook runs

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -27,7 +27,11 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
-from wandb.sdk.launch.utils import recursive_macro_sub
+from wandb.sdk.launch.utils import (
+    recursive_macro_sub,
+    sanitize_identifiers_for_k8s,
+    yield_containers,
+)
 from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
@@ -39,6 +43,7 @@ from ..utils import (
     MAX_ENV_LENGTHS,
     PROJECT_SYNCHRONOUS,
     get_kube_context_and_api_client,
+    make_k8s_label_safe,
     make_name_dns_safe,
 )
 from .abstract import AbstractRun, AbstractRunner
@@ -708,6 +713,7 @@ class KubernetesRunner(AbstractRunner):
         api_key_secret: Optional["V1Secret"] = None,
         wait_for_ready: bool = True,
         wait_timeout: int = 300,
+        auxiliary_resource_label_value: Optional[str] = None,
     ) -> None:
         """Prepare a service for launch.
 
@@ -725,6 +731,10 @@ class KubernetesRunner(AbstractRunner):
         config["metadata"].setdefault("labels", {})
         config["metadata"]["labels"][WANDB_K8S_RUN_ID] = run_id
         config["metadata"]["labels"]["wandb.ai/created-by"] = "launch-agent"
+        if auxiliary_resource_label_value:
+            config["metadata"]["labels"][WANDB_K8S_LABEL_AUXILIARY_RESOURCE] = (
+                auxiliary_resource_label_value
+            )
 
         env_vars = launch_project.get_env_vars_dict(
             self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
@@ -734,6 +744,13 @@ class KubernetesRunner(AbstractRunner):
         }
         add_wandb_env(config, wandb_config_env)
 
+        if auxiliary_resource_label_value:
+            add_label_to_pods(
+                config,
+                WANDB_K8S_LABEL_AUXILIARY_RESOURCE,
+                auxiliary_resource_label_value,
+            )
+
         if api_key_secret:
             for cont in yield_containers(config):
                 env = cont.setdefault("env", [])
@@ -751,6 +768,8 @@ class KubernetesRunner(AbstractRunner):
                 cont["env"] = env
 
         try:
+            sanitize_identifiers_for_k8s(config)
+
             await kubernetes_asyncio.utils.create_from_dict(
                 api_client, config, namespace=namespace
             )
@@ -924,6 +943,9 @@ class KubernetesRunner(AbstractRunner):
         additional_services: List[Dict[str, Any]] = recursive_macro_sub(
             launch_project.launch_spec.get("additional_services", []), update_dict
         )
+        auxiliary_resource_label_value = make_k8s_label_safe(
+            f"aux-{launch_project.target_entity}-{launch_project.target_project}-{launch_project.run_id}"
+        )
         if additional_services:
             wandb.termlog(
                 f"{LOG_PREFIX}Creating additional services: {additional_services}"
@@ -943,6 +965,7 @@ class KubernetesRunner(AbstractRunner):
                         secret,
                         wait_for_ready,
                         wait_timeout,
+                        auxiliary_resource_label_value,
                     )
                     for resource in additional_services
                     if resource.get("config", {})
@@ -980,7 +1003,7 @@ class KubernetesRunner(AbstractRunner):
             job_name,
             namespace,
             secret,
-            f"aux-{launch_project.target_entity}-{launch_project.target_project}-{launch_project.run_id}",
+            auxiliary_resource_label_value,
         )
         if self.backend_config[PROJECT_SYNCHRONOUS]:
             await submitted_job.wait()
@@ -1131,24 +1154,6 @@ async def maybe_create_imagepull_secret(
         raise LaunchError(f"Exception when creating Kubernetes secret: {str(e)}\n")
 
 
-def yield_containers(root: Any) -> Iterator[dict]:
-    """Yield all container specs in a manifest.
-
-    Recursively traverses the manifest and yields all container specs. Container
-    specs are identified by the presence of a "containers" key in the value.
-    """
-    if isinstance(root, dict):
-        for k, v in root.items():
-            if k == "containers":
-                if isinstance(v, list):
-                    yield from v
-            elif isinstance(v, (dict, list)):
-                yield from yield_containers(v)
-    elif isinstance(root, list):
-        for item in root:
-            yield from yield_containers(item)
-
-
 def add_wandb_env(root: Union[dict, list], env_vars: Dict[str, str]) -> None:
     """Injects wandb environment variables into specs.