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

wandb/sdk/artifacts/artifact_cache.py

@@ -0,0 +1,14 @@
+"""Recent Artifact storage.
+
+Artifacts are registered in the cache to ensure they won't be immediately garbage
+collected and can be retrieved by their ID.
+"""
+from typing import TYPE_CHECKING, Dict
+
+from wandb.sdk.lib.capped_dict import CappedDict
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
+# There is nothing special about the artifact cache, it's just a global capped dict.
+artifact_cache: Dict[str, "Artifact"] = CappedDict(100)

wandb/sdk/artifacts/artifact_manifest.py

@@ -5,9 +5,7 @@ from wandb.sdk.lib.hashutil import HexMD5
 
 if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-    from wandb.sdk.artifacts.storage_policies.wandb_storage_policy import (
-        WandbStoragePolicy,
-    )
+    from wandb.sdk.artifacts.storage_policy import StoragePolicy
 
 
 class ArtifactManifest:
@@ -29,12 +27,15 @@ class ArtifactManifest:
 
     def __init__(
         self,
-        storage_policy: "WandbStoragePolicy",
+        storage_policy: "StoragePolicy",
         entries: Optional[Mapping[str, "ArtifactManifestEntry"]] = None,
     ) -> None:
         self.storage_policy = storage_policy
         self.entries = dict(entries) if entries else {}
 
+    def __len__(self) -> int:
+        return len(self.entries)
+
     def to_manifest_json(self) -> Dict:
         raise NotImplementedError
 

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -1,11 +1,10 @@
 """Artifact manifest entry."""
+import json
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, Optional, Union
 from urllib.parse import urlparse
 
-import wandb
-from wandb import util
 from wandb.errors.term import termwarn
 from wandb.sdk.lib import filesystem
 from wandb.sdk.lib.hashutil import (
@@ -18,7 +17,6 @@ from wandb.sdk.lib.hashutil import (
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
 
 if TYPE_CHECKING:
-    from wandb.apis.public import RetryingClient
     from wandb.sdk.artifacts.artifact import Artifact
 
 
@@ -56,6 +54,38 @@ class ArtifactManifestEntry:
         if self.local_path and self.size is None:
             self.size = Path(self.local_path).stat().st_size
 
+    def __repr__(self) -> str:
+        cls = self.__class__.__name__
+        ref = f", ref={self.ref!r}" if self.ref is not None else ""
+        birth_artifact_id = (
+            f", birth_artifact_id={self.birth_artifact_id!r}"
+            if self.birth_artifact_id is not None
+            else ""
+        )
+        size = f", size={self.size}" if self.size is not None else ""
+        extra = f", extra={json.dumps(self.extra)}" if self.extra else ""
+        local_path = f", local_path={self.local_path!r}" if self.local_path else ""
+        others = ref + birth_artifact_id + size + extra + local_path
+        return f"{cls}(path={self.path!r}, digest={self.digest!r}{others})"
+
+    def __eq__(self, other: object) -> bool:
+        """Strict equality, comparing all public fields.
+
+        ArtifactManifestEntries for the same file may not compare equal if they were
+        added in different ways or created for different parent artifacts.
+        """
+        if not isinstance(other, ArtifactManifestEntry):
+            return False
+        return (
+            self.path == other.path
+            and self.digest == other.digest
+            and self.ref == other.ref
+            and self.birth_artifact_id == other.birth_artifact_id
+            and self.size == other.size
+            and self.extra == other.extra
+            and self.local_path == other.local_path
+        )
+
     @property
     def name(self) -> LogicalPath:
         # TODO(hugh): add telemetry to see if anyone is still using this.
@@ -151,9 +181,7 @@ class ArtifactManifestEntry:
     def _is_artifact_reference(self) -> bool:
         return self.ref is not None and urlparse(self.ref).scheme == "wandb-artifact"
 
-    def _get_referenced_artifact(self, client: "RetryingClient") -> "Artifact":
-        artifact: Artifact = wandb.Artifact._from_id(
-            hex_to_b64_id(util.host_from_path(self.ref)), client
-        )
-        assert artifact is not None
-        return artifact
+    def _referenced_artifact_id(self) -> Optional[str]:
+        if not self._is_artifact_reference():
+            return None
+        return hex_to_b64_id(urlparse(self.ref).netloc)

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -3,7 +3,6 @@ from typing import Any, Dict, Mapping, Optional
 
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
 from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-from wandb.sdk.artifacts.storage_policies.wandb_storage_policy import WandbStoragePolicy
 from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.lib.hashutil import HexMD5, _md5
 
@@ -23,13 +22,6 @@ class ArtifactManifestV1(ArtifactManifest):
         storage_policy_name = manifest_json["storagePolicy"]
         storage_policy_config = manifest_json.get("storagePolicyConfig", {})
         storage_policy_cls = StoragePolicy.lookup_by_name(storage_policy_name)
-        if storage_policy_cls is None:
-            raise ValueError('Failed to find storage policy "%s"' % storage_policy_name)
-        if not issubclass(storage_policy_cls, WandbStoragePolicy):
-            raise ValueError(
-                "No handler found for storage handler of type '%s'"
-                % storage_policy_name
-            )
 
         entries: Mapping[str, ArtifactManifestEntry]
         entries = {
@@ -49,7 +41,7 @@ class ArtifactManifestV1(ArtifactManifest):
 
     def __init__(
         self,
-        storage_policy: "WandbStoragePolicy",
+        storage_policy: "StoragePolicy",
         entries: Optional[Mapping[str, ArtifactManifestEntry]] = None,
     ) -> None:
         super().__init__(storage_policy, entries=entries)

wandb/sdk/artifacts/artifact_saver.py

@@ -4,15 +4,15 @@ import json
 import os
 import sys
 import tempfile
-from typing import TYPE_CHECKING, Awaitable, Dict, List, Optional, Sequence
+from typing import TYPE_CHECKING, Awaitable, Dict, Optional, Sequence
 
 import wandb
 import wandb.filesync.step_prepare
-from wandb import env, util
+from wandb import util
 from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
-from wandb.sdk.lib.filesystem import mkdir_exists_ok
+from wandb.sdk.artifacts.staging import get_staging_dir
 from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64
-from wandb.sdk.lib.paths import FilePathStr, URIStr
+from wandb.sdk.lib.paths import URIStr
 
 if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
@@ -65,9 +65,9 @@ class ArtifactSaver:
         distributed_id: Optional[str] = None,
         finalize: bool = True,
         metadata: Optional[Dict] = None,
+        ttl_duration_seconds: Optional[int] = None,
         description: Optional[str] = None,
         aliases: Optional[Sequence[str]] = None,
-        labels: Optional[List[str]] = None,
         use_after_commit: bool = False,
         incremental: bool = False,
         history_step: Optional[int] = None,
@@ -82,9 +82,9 @@ class ArtifactSaver:
                 distributed_id,
                 finalize,
                 metadata,
+                ttl_duration_seconds,
                 description,
                 aliases,
-                labels,
                 use_after_commit,
                 incremental,
                 history_step,
@@ -102,33 +102,17 @@ class ArtifactSaver:
         distributed_id: Optional[str] = None,
         finalize: bool = True,
         metadata: Optional[Dict] = None,
+        ttl_duration_seconds: Optional[int] = None,
         description: Optional[str] = None,
         aliases: Optional[Sequence[str]] = None,
-        labels: Optional[List[str]] = None,
         use_after_commit: bool = False,
         incremental: bool = False,
         history_step: Optional[int] = None,
         base_id: Optional[str] = None,
     ) -> Optional[Dict]:
-        aliases = aliases or []
         alias_specs = []
-        for alias in aliases:
-            if ":" in alias:
-                # Users can explicitly alias this artifact to names
-                # other than the primary one passed in by using the
-                # 'secondaryName:alias' notation.
-                idx = alias.index(":")
-                artifact_collection_name = alias[: idx - 1]
-                tag = alias[idx + 1 :]
-            else:
-                artifact_collection_name = name
-                tag = alias
-            alias_specs.append(
-                {
-                    "artifactCollectionName": artifact_collection_name,
-                    "alias": tag,
-                }
-            )
+        for alias in aliases or []:
+            alias_specs.append({"artifactCollectionName": name, "alias": alias})
 
         """Returns the server artifact."""
         self._server_artifact, latest = self._api.create_artifact(
@@ -136,35 +120,27 @@ class ArtifactSaver:
             name,
             self._digest,
             metadata=metadata,
+            ttl_duration_seconds=ttl_duration_seconds,
             aliases=alias_specs,
-            labels=labels,
             description=description,
             is_user_created=self._is_user_created,
             distributed_id=distributed_id,
             client_id=client_id,
             sequence_client_id=sequence_client_id,
-            enable_digest_deduplication=use_after_commit,  # Reuse logical duplicates in the `use_artifact` flow
             history_step=history_step,
         )
 
-        # TODO(artifacts):
-        #   if it's committed, all is good. If it's committing, just moving ahead isn't necessarily
-        #   correct. It may be better to poll until it's committed or failed, and then decided what to
-        #   do
         assert self._server_artifact is not None  # mypy optionality unwrapper
         artifact_id = self._server_artifact["id"]
         if base_id is None and latest:
             base_id = latest["id"]
-        if (
-            self._server_artifact["state"] == "COMMITTED"
-            or self._server_artifact["state"] == "COMMITTING"
-        ):
-            # TODO: update aliases, labels, description etc?
+        if self._server_artifact["state"] == "COMMITTED":
             if use_after_commit:
                 self._api.use_artifact(artifact_id)
             return self._server_artifact
-        elif (
+        if (
             self._server_artifact["state"] != "PENDING"
+            # For old servers, see https://github.com/wandb/wandb/pull/6190
             and self._server_artifact["state"] != "DELETED"
         ):
             raise Exception(
@@ -306,16 +282,3 @@ class ArtifactSaver:
                     os.remove(entry.local_path)
                 except OSError:
                     pass
-
-
-def get_staging_dir() -> FilePathStr:
-    path = os.path.join(env.get_data_dir(), "artifacts", "staging")
-    try:
-        mkdir_exists_ok(path)
-    except OSError as e:
-        raise PermissionError(
-            f"Unable to write staging files to {path}. To fix this problem, please set "
-            f"{env.DATA_DIR} to a directory where you have the necessary write access."
-        ) from e
-
-    return FilePathStr(os.path.abspath(os.path.expanduser(path)))

wandb/sdk/artifacts/artifact_ttl.py

@@ -0,0 +1,6 @@
+"""Artifact TTL."""
+from enum import Enum
+
+
+class ArtifactTTL(Enum):
+    INHERIT = 0

wandb/sdk/artifacts/artifacts_cache.py

@@ -1,23 +1,23 @@
 """Artifact cache."""
 import contextlib
+import errno
 import hashlib
 import os
-import secrets
-from typing import IO, TYPE_CHECKING, ContextManager, Dict, Generator, Optional, Tuple
+import shutil
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import IO, TYPE_CHECKING, ContextManager, Generator, Optional, Tuple
 
 import wandb
 from wandb import env, util
-from wandb.sdk.artifacts.exceptions import ArtifactNotLoggedError
-from wandb.sdk.lib.capped_dict import CappedDict
-from wandb.sdk.lib.filesystem import mkdir_exists_ok
+from wandb.errors import term
+from wandb.sdk.lib.filesystem import files_in
 from wandb.sdk.lib.hashutil import B64MD5, ETag, b64_to_hex_id
 from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 
 if TYPE_CHECKING:
     import sys
 
-    from wandb.sdk.artifacts.artifact import Artifact
-
     if sys.version_info >= (3, 8):
         from typing import Protocol
     else:
@@ -29,26 +29,18 @@ if TYPE_CHECKING:
 
 
 class ArtifactsCache:
-    _TMP_PREFIX = "tmp"
-
     def __init__(self, cache_dir: StrPath) -> None:
-        self._cache_dir = cache_dir
-        mkdir_exists_ok(self._cache_dir)
-        self._md5_obj_dir = os.path.join(self._cache_dir, "obj", "md5")
-        self._etag_obj_dir = os.path.join(self._cache_dir, "obj", "etag")
-        self._artifacts_by_id: Dict[str, Artifact] = CappedDict()
-        self._artifacts_by_client_id: Dict[str, Artifact] = CappedDict()
+        self._cache_dir = Path(cache_dir)
+        self._obj_dir = self._cache_dir / "obj"
+        self._temp_dir = self._cache_dir / "tmp"
+        self._temp_dir.mkdir(parents=True, exist_ok=True)
 
     def check_md5_obj_path(
         self, b64_md5: B64MD5, size: int
     ) -> Tuple[FilePathStr, bool, "Opener"]:
         hex_md5 = b64_to_hex_id(b64_md5)
-        path = os.path.join(self._cache_dir, "obj", "md5", hex_md5[:2], hex_md5[2:])
-        opener = self._cache_opener(path)
-        if os.path.isfile(path) and os.path.getsize(path) == size:
-            return FilePathStr(path), True, opener
-        mkdir_exists_ok(os.path.dirname(path))
-        return FilePathStr(path), False, opener
+        path = self._obj_dir / "md5" / hex_md5[:2] / hex_md5[2:]
+        return self._check_or_create(path, size)
 
     # TODO(spencerpearson): this method at least needs its signature changed.
     # An ETag is not (necessarily) a checksum.
@@ -62,103 +54,138 @@ class ArtifactsCache:
             hashlib.sha256(url.encode("utf-8")).digest()
             + hashlib.sha256(etag.encode("utf-8")).digest()
         ).hexdigest()
-        path = os.path.join(self._cache_dir, "obj", "etag", hexhash[:2], hexhash[2:])
-        opener = self._cache_opener(path)
-        if os.path.isfile(path) and os.path.getsize(path) == size:
-            return FilePathStr(path), True, opener
-        mkdir_exists_ok(os.path.dirname(path))
-        return FilePathStr(path), False, opener
-
-    def get_artifact(self, artifact_id: str) -> Optional["Artifact"]:
-        return self._artifacts_by_id.get(artifact_id)
+        path = self._obj_dir / "etag" / hexhash[:2] / hexhash[2:]
+        return self._check_or_create(path, size)
 
-    def store_artifact(self, artifact: "Artifact") -> None:
-        if not artifact.id:
-            raise ArtifactNotLoggedError(artifact, "store_artifact")
-        self._artifacts_by_id[artifact.id] = artifact
-
-    def get_client_artifact(self, client_id: str) -> Optional["Artifact"]:
-        return self._artifacts_by_client_id.get(client_id)
+    def _check_or_create(
+        self, path: Path, size: int
+    ) -> Tuple[FilePathStr, bool, "Opener"]:
+        opener = self._cache_opener(path, size)
+        hit = path.is_file() and path.stat().st_size == size
+        return FilePathStr(str(path)), hit, opener
 
-    def store_client_artifact(self, artifact: "Artifact") -> None:
-        self._artifacts_by_client_id[artifact._client_id] = artifact
+    def cleanup(
+        self,
+        target_size: Optional[int] = None,
+        target_fraction: Optional[float] = None,
+        remove_temp: bool = False,
+    ) -> int:
+        """Clean up the cache, removing the least recently used files first.
+
+        Args:
+            target_size: The target size of the cache in bytes. If the cache is larger
+                than this, we will remove the least recently used files until the cache
+                is smaller than this size.
+            target_fraction: The target fraction of the cache to reclaim. If the cache
+                is larger than this, we will remove the least recently used files until
+                the cache is smaller than this fraction of its current size. It is an
+                error to specify both target_size and target_fraction.
+            remove_temp: Whether to remove temporary files. Temporary files are files
+                that are currently being written to the cache. If remove_temp is True,
+                all temp files will be removed, regardless of the target_size or
+                target_fraction.
+
+        Returns:
+            The number of bytes reclaimed.
+        """
+        if target_size is None and target_fraction is None:
+            # Default to clearing the entire cache.
+            target_size = 0
+        if target_size is not None and target_fraction is not None:
+            raise ValueError("Cannot specify both target_size and target_fraction")
+        if target_size and target_size < 0:
+            raise ValueError("target_size must be non-negative")
+        if target_fraction and (target_fraction < 0 or target_fraction > 1):
+            raise ValueError("target_fraction must be between 0 and 1")
 
-    def cleanup(self, target_size: int, remove_temp: bool = False) -> int:
         bytes_reclaimed = 0
-        paths = {}
         total_size = 0
         temp_size = 0
-        for root, _, files in os.walk(self._cache_dir):
-            for file in files:
+
+        # Remove all temporary files if requested. Otherwise sum their size.
+        for entry in files_in(self._temp_dir):
+            size = entry.stat().st_size
+            total_size += size
+            if remove_temp:
                 try:
-                    path = os.path.join(root, file)
-                    stat = os.stat(path)
-
-                    if file.startswith(ArtifactsCache._TMP_PREFIX):
-                        if remove_temp:
-                            os.remove(path)
-                            bytes_reclaimed += stat.st_size
-                        else:
-                            temp_size += stat.st_size
-                        continue
+                    os.remove(entry.path)
+                    bytes_reclaimed += size
                 except OSError:
-                    continue
-                paths[path] = stat
-                total_size += stat.st_size
-
+                    pass
+            else:
+                temp_size += size
         if temp_size:
             wandb.termwarn(
                 f"Cache contains {util.to_human_size(temp_size)} of temporary files. "
                 "Run `wandb artifact cleanup --remove-temp` to remove them."
             )
 
-        sorted_paths = sorted(paths.items(), key=lambda x: x[1].st_atime)
-        for path, stat in sorted_paths:
-            if total_size < target_size:
-                return bytes_reclaimed
+        entries = []
+        for file_entry in files_in(self._obj_dir):
+            total_size += file_entry.stat().st_size
+            entries.append(file_entry)
+
+        if target_fraction is not None:
+            target_size = int(total_size * target_fraction)
+        assert target_size is not None
 
+        for entry in sorted(entries, key=lambda x: x.stat().st_atime):
+            if total_size <= target_size:
+                return bytes_reclaimed
             try:
-                os.remove(path)
+                os.remove(entry.path)
             except OSError:
                 pass
+            total_size -= entry.stat().st_size
+            bytes_reclaimed += entry.stat().st_size
+
+        if total_size > target_size:
+            wandb.termerror(
+                f"Failed to reclaim enough space in {self._cache_dir}. Try running"
+                " `wandb artifact cleanup --remove-temp` to remove temporary files."
+            )
 
-            total_size -= stat.st_size
-            bytes_reclaimed += stat.st_size
         return bytes_reclaimed
 
-    def _cache_opener(self, path: StrPath) -> "Opener":
+    def _free_space(self) -> int:
+        """Return the number of bytes of free space in the cache directory."""
+        return shutil.disk_usage(self._cache_dir)[2]
+
+    def _reserve_space(self, size: int) -> None:
+        """If a `size` write would exceed disk space, remove cached items to make space.
+
+        Raises:
+            OSError: If there is not enough space to write `size` bytes, even after
+                removing cached items.
+        """
+        if size <= self._free_space():
+            return
+
+        term.termwarn("Cache size exceeded. Attempting to reclaim space...")
+        self.cleanup(target_fraction=0.5)
+        if size <= self._free_space():
+            return
+
+        self.cleanup(target_size=0)
+        if size > self._free_space():
+            raise OSError(errno.ENOSPC, f"Insufficient free space in {self._cache_dir}")
+
+    def _cache_opener(self, path: Path, size: int) -> "Opener":
         @contextlib.contextmanager
         def helper(mode: str = "w") -> Generator[IO, None, None]:
             if "a" in mode:
                 raise ValueError("Appending to cache files is not supported")
 
-            dirname = os.path.dirname(path)
-            tmp_file = os.path.join(
-                dirname, f"{ArtifactsCache._TMP_PREFIX}_{secrets.token_hex(8)}"
-            )
-            with util.fsync_open(tmp_file, mode=mode) as f:
-                yield f
-
+            self._reserve_space(size)
+            temp_file = NamedTemporaryFile(dir=self._temp_dir, mode=mode, delete=False)
             try:
-                # Use replace where we can, as it implements an atomic
-                # move on most platforms. If it doesn't exist, we have
-                # to use rename which isn't atomic in all cases but there
-                # isn't a better option.
-                #
-                # The atomic replace is important in the event multiple processes
-                # attempt to write to / read from the cache at the same time. Each
-                # writer firsts stages its writes to a temporary file in the cache.
-                # Once it is finished, we issue an atomic replace operation to update
-                # the cache. Although this can result in redundant downloads, this
-                # guarantees that readers can NEVER read incomplete files from the
-                # cache.
-                #
-                # IMPORTANT: Replace is NOT atomic across different filesystems. This why
-                # it is critical that the temporary files sit directly in the cache --
-                # they need to be on the same filesystem!
-                os.replace(tmp_file, path)
-            except AttributeError:
-                os.rename(tmp_file, path)
+                yield temp_file
+                temp_file.close()
+                path.parent.mkdir(parents=True, exist_ok=True)
+                os.replace(temp_file.name, path)
+            except Exception:
+                os.remove(temp_file.name)
+                raise
 
         return helper
 
@@ -169,6 +196,5 @@ _artifacts_cache = None
 def get_artifacts_cache() -> ArtifactsCache:
     global _artifacts_cache
     if _artifacts_cache is None:
-        cache_dir = os.path.join(env.get_cache_dir(), "artifacts")
-        _artifacts_cache = ArtifactsCache(cache_dir)
+        _artifacts_cache = ArtifactsCache(env.get_cache_dir() / "artifacts")
     return _artifacts_cache

wandb/sdk/artifacts/staging.py

@@ -0,0 +1,25 @@
+"""Manages artifact file staging.
+
+Artifact files are copied to the staging area as soon as they are added to an artifact
+in order to avoid file changes corrupting the artifact. Once the upload is complete, the
+file should be moved to the artifact cache.
+"""
+
+import os
+
+from wandb import env
+from wandb.sdk.lib.filesystem import mkdir_exists_ok
+from wandb.sdk.lib.paths import FilePathStr
+
+
+def get_staging_dir() -> FilePathStr:
+    path = os.path.join(env.get_data_dir(), "artifacts", "staging")
+    try:
+        mkdir_exists_ok(path)
+    except OSError as e:
+        raise PermissionError(
+            f"Unable to write staging files to {path}. To fix this problem, please set "
+            f"{env.DATA_DIR} to a directory where you have the necessary write access."
+        ) from e
+
+    return FilePathStr(os.path.abspath(os.path.expanduser(path)))

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -204,7 +204,9 @@ class S3Handler(StorageHandler):
             )
         return entries
 
-    def _size_from_obj(self, obj: "boto3.s3.Object") -> int:
+    def _size_from_obj(
+        self, obj: Union["boto3.s3.Object", "boto3.s3.ObjectSummary"]
+    ) -> int:
         # ObjectSummary has size, Object has content_length
         size: int
         if hasattr(obj, "size"):
@@ -215,7 +217,7 @@ class S3Handler(StorageHandler):
 
     def _entry_from_obj(
         self,
-        obj: "boto3.s3.Object",
+        obj: Union["boto3.s3.Object", "boto3.s3.ObjectSummary"],
         path: str,
         name: Optional[StrPath] = None,
         prefix: str = "",
@@ -261,18 +263,21 @@ class S3Handler(StorageHandler):
         )
 
     @staticmethod
-    def _etag_from_obj(obj: "boto3.s3.Object") -> ETag:
+    def _etag_from_obj(obj: Union["boto3.s3.Object", "boto3.s3.ObjectSummary"]) -> ETag:
         etag: ETag
         etag = obj.e_tag[1:-1]  # escape leading and trailing quote
         return etag
 
-    @staticmethod
-    def _extra_from_obj(obj: "boto3.s3.Object") -> Dict[str, str]:
+    def _extra_from_obj(
+        self, obj: Union["boto3.s3.Object", "boto3.s3.ObjectSummary"]
+    ) -> Dict[str, str]:
         extra = {
             "etag": obj.e_tag[1:-1],  # escape leading and trailing quote
         }
-        # ObjectSummary will never have version_id
-        if hasattr(obj, "version_id") and obj.version_id != "null":
+        if not hasattr(obj, "version_id"):
+            # Convert ObjectSummary to Object to get the version_id.
+            obj = self._s3.Object(obj.bucket_name, obj.key)  # type: ignore[union-attr]
+        if hasattr(obj, "version_id") and obj.version_id and obj.version_id != "null":
             extra["versionID"] = obj.version_id
         return extra
 

wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py

@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Optional, Sequence, Union
 
 import wandb
 from wandb import util
+from wandb.sdk.artifacts.artifact_cache import artifact_cache
 from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
 from wandb.sdk.artifacts.storage_handler import StorageHandler
 from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 
@@ -20,7 +20,6 @@ class WBLocalArtifactHandler(StorageHandler):
 
     def __init__(self) -> None:
         self._scheme = "wandb-client-artifact"
-        self._cache = get_artifacts_cache()
 
     def can_handle(self, parsed_url: "ParseResult") -> bool:
         return parsed_url.scheme == self._scheme
@@ -54,7 +53,7 @@ class WBLocalArtifactHandler(StorageHandler):
         """
         client_id = util.host_from_path(path)
         target_path = util.uri_from_path(path)
-        target_artifact = self._cache.get_client_artifact(client_id)
+        target_artifact = artifact_cache.get(client_id)
         if not isinstance(target_artifact, wandb.Artifact):
             raise RuntimeError("Local Artifact not found - invalid reference")
         target_entry = target_artifact._manifest.entries[target_path]