wandb 0.21.1__py3-none-win32.whl → 0.21.2__py3-none-win32.whl

wandb/sdk/artifacts/artifact_file_cache.py

@@ -9,6 +9,7 @@ import os
 import shutil
 import subprocess
 import sys
+from functools import lru_cache
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 from typing import IO, ContextManager, Iterator, Protocol
@@ -236,12 +237,15 @@ class ArtifactFileCache:
             ) from e
 
 
-_artifact_file_cache: ArtifactFileCache | None = None
+# Memo `ArtifactFileCache` instances while avoiding reliance on global
+# variable(s).  Notes:
+# - @lru_cache should be thread-safe.
+# - We don't memoize `get_artifact_file_cache` directly, as the cache_dir
+#   may change at runtime.  This is likely rare in practice, though.
+@lru_cache(maxsize=1)
+def _build_artifact_file_cache(cache_dir: StrPath) -> ArtifactFileCache:
+    return ArtifactFileCache(cache_dir)
 
 
 def get_artifact_file_cache() -> ArtifactFileCache:
-    global _artifact_file_cache
-    cache_dir = env.get_cache_dir() / "artifacts"
-    if _artifact_file_cache is None or _artifact_file_cache._cache_dir != cache_dir:
-        _artifact_file_cache = ArtifactFileCache(cache_dir)
-    return _artifact_file_cache
+    return _build_artifact_file_cache(env.get_cache_dir() / "artifacts")

wandb/sdk/artifacts/artifact_manifest.py

@@ -50,10 +50,12 @@ class ArtifactManifest:
 
     def add_entry(self, entry: ArtifactManifestEntry, overwrite: bool = False) -> None:
         path = entry.path
-        if not overwrite:
-            prev_entry = self.entries.get(path)
-            if prev_entry and (entry.digest != prev_entry.digest):
-                raise ValueError(f"Cannot add the same path twice: {path!r}")
+        if (
+            (not overwrite)
+            and (old_entry := self.entries.get(path))
+            and (entry.digest != old_entry.digest)
+        ):
+            raise ValueError(f"Cannot add the same path twice: {path!r}")
         self.entries[path] = entry
 
     def remove_entry(self, entry: ArtifactManifestEntry) -> None:
@@ -67,9 +69,8 @@ class ArtifactManifest:
 
     def get_entries_in_directory(self, directory: str) -> list[ArtifactManifestEntry]:
         return [
-            self.entries[entry_key]
-            for entry_key in self.entries
-            if entry_key.startswith(
-                directory + "/"
-            )  # entries use forward slash even for windows
+            entry
+            for key, entry in self.entries.items()
+            # entry keys (paths) use forward slash even for windows
+            if key.startswith(f"{directory}/")
         ]

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -40,6 +40,9 @@ if TYPE_CHECKING:
         local_path: str
 
 
+_WB_ARTIFACT_SCHEME = "wandb-artifact"
+
+
 class ArtifactManifestEntry:
     """A single entry in an artifact manifest."""
 
@@ -221,15 +224,11 @@ class ArtifactManifestEntry:
             derived_artifact.add_reference(ref_url)
             ```
         """
-        if self._parent_artifact is None:
-            raise NotImplementedError
-        assert self._parent_artifact.id is not None
-        return (
-            "wandb-artifact://"
-            + b64_to_hex_id(B64MD5(self._parent_artifact.id))
-            + "/"
-            + self.path
-        )
+        if (parent_artifact := self.parent_artifact()) is None:
+            raise ValueError("Parent artifact is not set")
+        elif (parent_id := parent_artifact.id) is None:
+            raise ValueError("Parent artifact ID is not set")
+        return f"{_WB_ARTIFACT_SCHEME}://{b64_to_hex_id(B64MD5(parent_id))}/{self.path}"
 
     def to_json(self) -> ArtifactManifestEntryDict:
         contents: ArtifactManifestEntryDict = {
@@ -251,7 +250,7 @@ class ArtifactManifestEntry:
         return contents
 
     def _is_artifact_reference(self) -> bool:
-        return self.ref is not None and urlparse(self.ref).scheme == "wandb-artifact"
+        return self.ref is not None and urlparse(self.ref).scheme == _WB_ARTIFACT_SCHEME
 
     def _referenced_artifact_id(self) -> str | None:
         if not self._is_artifact_reference():

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from operator import itemgetter
 from typing import Any, Mapping
 
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
@@ -64,7 +65,7 @@ class ArtifactManifestV1(ArtifactManifest):
         contents.
         """
         contents = {}
-        for entry in sorted(self.entries.values(), key=lambda k: k.path):
+        for name, entry in sorted(self.entries.items(), key=itemgetter(0)):
             json_entry: dict[str, Any] = {
                 "digest": entry.digest,
             }
@@ -76,7 +77,7 @@ class ArtifactManifestV1(ArtifactManifest):
                 json_entry["extra"] = entry.extra
             if entry.size is not None:
                 json_entry["size"] = entry.size
-            contents[entry.path] = json_entry
+            contents[name] = json_entry
         return {
             "version": self.__class__.version(),
             "storagePolicy": self.storage_policy.name(),
@@ -87,6 +88,7 @@ class ArtifactManifestV1(ArtifactManifest):
     def digest(self) -> HexMD5:
         hasher = _md5()
         hasher.update(b"wandb-artifact-manifest-v1\n")
-        for name, entry in sorted(self.entries.items(), key=lambda kv: kv[0]):
+        # sort by key (path)
+        for name, entry in sorted(self.entries.items(), key=itemgetter(0)):
             hasher.update(f"{name}:{entry.digest}\n".encode())
         return HexMD5(hasher.hexdigest())

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -42,7 +42,7 @@ class HTTPHandler(StorageHandler):
 
         path, hit, cache_open = self._cache.check_etag_obj_path(
             URIStr(manifest_entry.ref),
-            ETag(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            ETag(manifest_entry.digest),
             manifest_entry.size if manifest_entry.size is not None else 0,
         )
         if hit:

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -94,7 +94,7 @@ class S3Handler(StorageHandler):
 
         path, hit, cache_open = self._cache.check_etag_obj_path(
             URIStr(manifest_entry.ref),
-            ETag(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            ETag(manifest_entry.digest),
             manifest_entry.size if manifest_entry.size is not None else 0,
         )
         if hit:

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -167,7 +167,7 @@ class WandbStoragePolicy(StoragePolicy):
             self._cache._override_cache_path = dest_path
 
         path, hit, cache_open = self._cache.check_md5_obj_path(
-            B64MD5(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            B64MD5(manifest_entry.digest),
             manifest_entry.size if manifest_entry.size is not None else 0,
         )
         if hit:

wandb/sdk/data_types/video.py

@@ -25,8 +25,8 @@ if TYPE_CHECKING:  # pragma: no cover
 
 
 def _should_print_spinner() -> bool:
-    singleton = wandb_setup.singleton_if_setup()
-    if singleton and (singleton.settings.quiet or singleton.settings.silent):
+    settings = wandb_setup.singleton().settings_if_loaded
+    if settings and (settings.quiet or settings.silent):
         return False
 
     return not env.is_quiet() and not env.is_silent()

wandb/sdk/interface/interface_queue.py

@@ -8,8 +8,6 @@ import logging
 from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Optional
 
-from wandb.sdk.mailbox import Mailbox
-
 from .interface_shared import InterfaceShared
 
 if TYPE_CHECKING:
@@ -27,12 +25,11 @@ class InterfaceQueue(InterfaceShared):
         record_q: Optional["Queue[pb.Record]"] = None,
         result_q: Optional["Queue[pb.Result]"] = None,
         process: Optional[BaseProcess] = None,
-        mailbox: Optional[Mailbox] = None,
     ) -> None:
         self.record_q = record_q
         self.result_q = result_q
         self._process = process
-        super().__init__(mailbox=mailbox)
+        super().__init__()
 
     def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
         if self._process and not self._process.is_alive():

wandb/sdk/interface/interface_shared.py

@@ -10,7 +10,7 @@ from typing import Any, Optional, cast
 
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto import wandb_telemetry_pb2 as tpb
-from wandb.sdk.mailbox import Mailbox, MailboxHandle
+from wandb.sdk.mailbox import MailboxHandle
 from wandb.util import json_dumps_safer, json_friendly
 
 from .interface import InterfaceBase
@@ -19,9 +19,8 @@ logger = logging.getLogger("wandb")
 
 
 class InterfaceShared(InterfaceBase):
-    def __init__(self, mailbox: Optional[Mailbox] = None) -> None:
+    def __init__(self) -> None:
         super().__init__()
-        self._mailbox = mailbox
 
     def _publish_output(self, outdata: pb.OutputRecord) -> None:
         rec = pb.Record()
@@ -67,7 +66,7 @@ class InterfaceShared(InterfaceBase):
         self, job_input: pb.JobInputRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(job_input=job_input)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _make_stats(self, stats_dict: dict) -> pb.StatsRecord:
         stats = pb.StatsRecord()
@@ -263,6 +262,9 @@ class InterfaceShared(InterfaceBase):
     def _publish(self, record: pb.Record, local: Optional[bool] = None) -> None:
         raise NotImplementedError
 
+    def _deliver(self, record: pb.Record) -> "MailboxHandle[pb.Result]":
+        raise NotImplementedError
+
     def _publish_defer(self, state: "pb.DeferRequest.DeferState.V") -> None:
         defer = pb.DeferRequest(state=state)
         rec = self._make_request(defer=defer)
@@ -333,19 +335,19 @@ class InterfaceShared(InterfaceBase):
         log_artifact: pb.LogArtifactRequest,
     ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(log_artifact=log_artifact)
-        return self._deliver_record(rec)
+        return self._deliver(rec)
 
     def _deliver_download_artifact(
         self, download_artifact: pb.DownloadArtifactRequest
     ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(download_artifact=download_artifact)
-        return self._deliver_record(rec)
+        return self._deliver(rec)
 
     def _deliver_link_artifact(
         self, link_artifact: pb.LinkArtifactRequest
     ) -> MailboxHandle[pb.Result]:
         rec = self._make_request(link_artifact=link_artifact)
-        return self._deliver_record(rec)
+        return self._deliver(rec)
 
     def _publish_artifact(self, proto_artifact: pb.ArtifactRecord) -> None:
         rec = self._make_record(artifact=proto_artifact)
@@ -360,7 +362,7 @@ class InterfaceShared(InterfaceBase):
         status: pb.StatusRequest,
     ) -> MailboxHandle[pb.Result]:
         req = self._make_request(status=status)
-        return self._deliver_record(req)
+        return self._deliver(req)
 
     def _publish_exit(self, exit_data: pb.RunExitRecord) -> None:
         rec = self._make_record(exit=exit_data)
@@ -373,110 +375,97 @@ class InterfaceShared(InterfaceBase):
     def _deliver_shutdown(self) -> MailboxHandle[pb.Result]:
         request = pb.Request(shutdown=pb.ShutdownRequest())
         record = self._make_record(request=request)
-        return self._deliver_record(record)
-
-    def _get_mailbox(self) -> Mailbox:
-        mailbox = self._mailbox
-        assert mailbox
-        return mailbox
-
-    def _deliver_record(self, record: pb.Record) -> MailboxHandle[pb.Result]:
-        mailbox = self._get_mailbox()
-
-        handle = mailbox.require_response(record)
-        self._publish(record)
-
-        return handle.map(lambda resp: resp.result_communicate)
+        return self._deliver(record)
 
     def _deliver_run(self, run: pb.RunRecord) -> MailboxHandle[pb.Result]:
         record = self._make_record(run=run)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_finish_sync(
         self,
         sync_finish: pb.SyncFinishRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(sync_finish=sync_finish)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_run_start(
         self,
         run_start: pb.RunStartRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_start=run_start)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_get_summary(
         self,
         get_summary: pb.GetSummaryRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(get_summary=get_summary)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_get_system_metrics(
         self, get_system_metrics: pb.GetSystemMetricsRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(get_system_metrics=get_system_metrics)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_exit(
         self,
         exit_data: pb.RunExitRecord,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_record(exit=exit_data)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def deliver_operation_stats(self):
         record = self._make_request(operation_stats=pb.OperationStatsRequest())
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_poll_exit(
         self,
         poll_exit: pb.PollExitRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(poll_exit=poll_exit)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_finish_without_exit(
         self, run_finish_without_exit: pb.RunFinishWithoutExitRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_finish_without_exit=run_finish_without_exit)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_stop_status(
         self,
         stop_status: pb.StopStatusRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(stop_status=stop_status)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_attach(
         self,
         attach: pb.AttachRequest,
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(attach=attach)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_network_status(
         self, network_status: pb.NetworkStatusRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(network_status=network_status)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_internal_messages(
         self, internal_message: pb.InternalMessagesRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(internal_messages=internal_message)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_request_sampled_history(
         self, sampled_history: pb.SampledHistoryRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(sampled_history=sampled_history)
-        return self._deliver_record(record)
+        return self._deliver(record)
 
     def _deliver_request_run_status(
         self, run_status: pb.RunStatusRequest
     ) -> MailboxHandle[pb.Result]:
         record = self._make_request(run_status=run_status)
-        return self._deliver_record(record)
+        return self._deliver(record)

wandb/sdk/interface/interface_sock.py

@@ -1,19 +1,18 @@
-"""InterfaceSock - Derived from InterfaceShared using a socket to send to internal thread.
-
-See interface.py for how interface classes relate to each other.
-
-"""
+from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
+
+from typing_extensions import override
 
-from wandb.sdk.mailbox import Mailbox
+from wandb.proto import wandb_server_pb2 as spb
 
-from ..lib.sock_client import SockClient
 from .interface_shared import InterfaceShared
 
 if TYPE_CHECKING:
     from wandb.proto import wandb_internal_pb2 as pb
+    from wandb.sdk.lib.service.service_client import ServiceClient
+    from wandb.sdk.mailbox import MailboxHandle
 
 
 logger = logging.getLogger("wandb")
@@ -22,18 +21,29 @@ logger = logging.getLogger("wandb")
 class InterfaceSock(InterfaceShared):
     def __init__(
         self,
-        sock_client: SockClient,
-        mailbox: Mailbox,
+        client: ServiceClient,
         stream_id: str,
     ) -> None:
-        super().__init__(mailbox=mailbox)
-        self._sock_client = sock_client
+        super().__init__()
+        self._client = client
         self._stream_id = stream_id
 
     def _assign(self, record: Any) -> None:
         assert self._stream_id
         record._info.stream_id = self._stream_id
 
-    def _publish(self, record: "pb.Record", local: Optional[bool] = None) -> None:
+    @override
+    def _publish(self, record: pb.Record, local: bool | None = None) -> None:
         self._assign(record)
-        self._sock_client.send_record_publish(record)
+        request = spb.ServerRequest()
+        request.record_publish.CopyFrom(record)
+        self._client.publish(request)
+
+    @override
+    def _deliver(self, record: pb.Record) -> MailboxHandle[pb.Result]:
+        self._assign(record)
+        request = spb.ServerRequest()
+        request.record_publish.CopyFrom(record)
+
+        handle = self._client.deliver(request)
+        return handle.map(lambda response: response.result_communicate)

wandb/sdk/internal/settings_static.py

@@ -4,7 +4,7 @@ from typing import Any, Iterable
 
 from wandb.proto import wandb_settings_pb2
 from wandb.sdk.lib import RunMoment
-from wandb.sdk.wandb_settings import Settings
+from wandb.sdk.wandb_settings import CLIENT_ONLY_SETTINGS, Settings
 
 
 class SettingsStatic(Settings):
@@ -41,8 +41,7 @@ class SettingsStatic(Settings):
 
         forks_specified: list[str] = []
         for key in fields:
-            # Skip Python-only keys that do not exist on the proto.
-            if key in ("reinit",):
+            if key in CLIENT_ONLY_SETTINGS:
                 continue
 
             value: Any = None

wandb/sdk/launch/create_job.py

@@ -11,6 +11,7 @@ from wandb.apis.internal import Api
 from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal.job_builder import JobBuilder
 from wandb.sdk.launch.git_reference import GitReference
+from wandb.sdk.launch.inputs.internal import _validate_schema
 from wandb.sdk.launch.utils import (
     _is_git_uri,
     get_current_python_version,
@@ -116,6 +117,7 @@ def _create_job(
     dockerfile: Optional[str] = None,
     base_image: Optional[str] = None,
     services: Optional[Dict[str, str]] = None,
+    schema: Optional[Dict[str, Any]] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -206,6 +208,15 @@ def _create_job(
     if "latest" not in aliases:
         aliases += ["latest"]
 
+    metadata = {"_partial": True}
+    if schema:
+        _validate_schema(schema)
+        metadata = {
+            "input_schemas": {
+                "@wandb.config": schema,
+            }
+        }
+
     res, _ = api.create_artifact(
         artifact_type_name="job",
         artifact_collection_name=name,
@@ -216,7 +227,7 @@ def _create_job(
         project_name=project,
         run_name=run.id,  # type: ignore # run will be deleted after creation
         description=description,
-        metadata={"_partial": True},
+        metadata=metadata,
         is_user_created=True,
         aliases=[{"artifactCollectionName": name, "alias": a} for a in aliases],
     )