wandb 0.22.2__py3-none-musllinux_1_2_aarch64.whl → 0.22.3__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/artifacts/artifact_file_cache.py

@@ -147,7 +147,7 @@ class ArtifactFileCache:
         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."
+                "Run `wandb artifact cache cleanup --remove-temp` to remove them."
             )
 
         entries = []

wandb/sdk/artifacts/artifact_manifest.py

@@ -2,75 +2,80 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Mapping
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict
+
+from pydantic import Field
+from typing_extensions import Annotated
 
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib.hashutil import HexMD5
 
+from ._models.base_model import ArtifactsBase
+
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-    from wandb.sdk.artifacts.storage_policy import StoragePolicy
+    from .artifact_manifest_entry import ArtifactManifestEntry
+    from .storage_policy import StoragePolicy
 
 
-class ArtifactManifest:
-    entries: dict[str, ArtifactManifestEntry]
+class ArtifactManifest(ArtifactsBase, ABC):
+    # Note: this can't be named "version" since it conflicts with the prior `version()` classmethod.
+    manifest_version: Annotated[Any, Field(repr=False)]
+    entries: Dict[str, ArtifactManifestEntry] = Field(default_factory=dict)  # noqa: UP006
+
+    storage_policy: Annotated[StoragePolicy, Field(exclude=True, repr=False)]
 
     @classmethod
+    def version(cls) -> int:
+        return cls.model_fields["manifest_version"].default
+
+    @classmethod
+    @abstractmethod
     def from_manifest_json(
-        cls, manifest_json: dict, api: InternalApi | None = None
+        cls, manifest_json: dict[str, Any], api: InternalApi | None = None
     ) -> ArtifactManifest:
-        if "version" not in manifest_json:
+        if (version := manifest_json.get("version")) is None:
             raise ValueError("Invalid manifest format. Must contain version field.")
-        version = manifest_json["version"]
+
         for sub in cls.__subclasses__():
             if sub.version() == version:
                 return sub.from_manifest_json(manifest_json, api=api)
         raise ValueError("Invalid manifest version.")
 
-    @classmethod
-    def version(cls) -> int:
-        raise NotImplementedError
-
-    def __init__(
-        self,
-        storage_policy: StoragePolicy,
-        entries: Mapping[str, ArtifactManifestEntry] | None = 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:
+    @abstractmethod
+    def to_manifest_json(self) -> dict[str, Any]:
         raise NotImplementedError
 
+    @abstractmethod
     def digest(self) -> HexMD5:
         raise NotImplementedError
 
+    @abstractmethod
+    def size(self) -> int:
+        raise NotImplementedError
+
     def add_entry(self, entry: ArtifactManifestEntry, overwrite: bool = False) -> None:
-        path = entry.path
         if (
             (not overwrite)
-            and (old_entry := self.entries.get(path))
+            and (old_entry := self.entries.get(entry.path))
             and (entry.digest != old_entry.digest)
         ):
-            raise ValueError(f"Cannot add the same path twice: {path!r}")
-        self.entries[path] = entry
+            raise ValueError(f"Cannot add the same path twice: {entry.path!r}")
+        self.entries[entry.path] = entry
 
     def remove_entry(self, entry: ArtifactManifestEntry) -> None:
         try:
             del self.entries[entry.path]
         except LookupError:
-            raise FileNotFoundError(f"Cannot remove missing entry: '{entry.path}'")
+            raise FileNotFoundError(f"Cannot remove missing entry: {entry.path!r}")
 
     def get_entry_by_path(self, path: str) -> ArtifactManifestEntry | None:
         return self.entries.get(path)
 
     def get_entries_in_directory(self, directory: str) -> list[ArtifactManifestEntry]:
-        return [
-            entry
-            for key, entry in self.entries.items()
-            # entry keys (paths) use forward slash even for windows
-            if key.startswith(f"{directory}/")
-        ]
+        # entry keys (paths) use forward slash even for windows
+        dir_prefix = f"{directory}/"
+        return [obj for key, obj in self.entries.items() if key.startswith(dir_prefix)]

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -1,17 +1,24 @@
 """Artifact manifest entry."""
 
+# Older-style type annotations required for Pydantic v1 / python 3.8 compatibility.
+# ruff: noqa: UP006, UP007, UP045
+
 from __future__ import annotations
 
 import concurrent.futures
 import hashlib
-import json
 import logging
 import os
 from contextlib import suppress
-from pathlib import Path
-from typing import TYPE_CHECKING
+from os.path import getsize
+from typing import TYPE_CHECKING, Any, Dict, Final, Optional, Union
 from urllib.parse import urlparse
 
+from pydantic import Field, NonNegativeInt
+from typing_extensions import Annotated, Self
+
+from wandb._pydantic import field_validator, model_validator
+from wandb._strutils import nameof
 from wandb.proto.wandb_deprecated import Deprecated
 from wandb.sdk.lib.deprecate import deprecate
 from wandb.sdk.lib.filesystem import copy_or_overwrite_changed
@@ -22,27 +29,18 @@ from wandb.sdk.lib.hashutil import (
     hex_to_b64_id,
     md5_file_b64,
 )
-from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
+from wandb.sdk.lib.paths import FilePathStr, LogicalPath, URIStr
 
-logger = logging.getLogger(__name__)
+from ._models.base_model import ArtifactsBase
 
 if TYPE_CHECKING:
-    from typing_extensions import TypedDict
+    from .artifact import Artifact
 
-    from wandb.sdk.artifacts.artifact import Artifact
 
-    class ArtifactManifestEntryDict(TypedDict, total=False):
-        path: str
-        digest: str
-        skip_cache: bool
-        ref: str
-        birthArtifactID: str
-        size: int
-        extra: dict
-        local_path: str
+logger = logging.getLogger(__name__)
 
 
-_WB_ARTIFACT_SCHEME = "wandb-artifact"
+_WB_ARTIFACT_SCHEME: Final[str] = "wandb-artifact"
 
 
 def _checksum_cache_path(file_path: str) -> str:
@@ -87,76 +85,54 @@ def _write_cached_checksum(file_path: str, checksum: str) -> None:
         logger.debug(f"Failed to write checksum cache for {file_path!r}")
 
 
-class ArtifactManifestEntry:
-    """A single entry in an artifact manifest."""
+class ArtifactManifestEntry(ArtifactsBase):
+    """A single entry in an artifact manifest.
+
+    External code should avoid instantiating this class directly.
+    """
 
     path: LogicalPath
-    digest: B64MD5 | URIStr | FilePathStr | ETag
-    skip_cache: bool
-    ref: FilePathStr | URIStr | None
-    birth_artifact_id: str | None
-    size: int | None
-    extra: dict
-    local_path: str | None
-
-    _parent_artifact: Artifact | None = None
-    _download_url: str | None = None
-
-    def __init__(
-        self,
-        path: StrPath,
-        digest: B64MD5 | URIStr | FilePathStr | ETag,
-        skip_cache: bool | None = False,
-        ref: FilePathStr | URIStr | None = None,
-        birth_artifact_id: str | None = None,
-        size: int | None = None,
-        extra: dict | None = None,
-        local_path: StrPath | None = None,
-    ) -> None:
-        self.path = LogicalPath(path)
-        self.digest = digest
-        self.ref = ref
-        self.birth_artifact_id = birth_artifact_id
-        self.size = size
-        self.extra = extra or {}
-        self.local_path = str(local_path) if local_path else None
-        if self.local_path and self.size is None:
-            self.size = Path(self.local_path).stat().st_size
-        self.skip_cache = skip_cache or False
 
-    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 ""
-        skip_cache = f", skip_cache={self.skip_cache}"
-        others = ref + birth_artifact_id + size + extra + local_path + skip_cache
-        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.
+    digest: Union[B64MD5, ETag, URIStr, FilePathStr]
+    ref: Union[URIStr, FilePathStr, None] = None
+    birth_artifact_id: Annotated[Optional[str], Field(alias="birthArtifactID")] = None
+    size: Optional[NonNegativeInt] = None
+    extra: Dict[str, Any] = Field(default_factory=dict)
+    local_path: Optional[str] = None
+
+    skip_cache: bool = False
+
+    # Note: Pydantic considers these private attributes, omitting them from validation and comparison logic.
+    _parent_artifact: Optional[Artifact] = None
+    _download_url: Optional[str] = None
+
+    @field_validator("path", mode="before")
+    def _validate_path(cls, v: Any) -> LogicalPath:
+        """Coerce `path` to a LogicalPath.
+
+        LogicalPath doesn't implement its own pydantic validator, and implementing one for
+        both pydantic V1 _and_ V2 would add too much boilerplate. Until we drop V1 support,
+        just coerce to LogicalPath in the field validator here.
         """
-        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
-            and self.skip_cache == other.skip_cache
-        )
+        return LogicalPath(v)
+
+    @field_validator("local_path", mode="before")
+    def _validate_local_path(cls, v: Any) -> str | None:
+        """Coerce `local_path` to a str. Necessary if the input is a `PosixPath`."""
+        return str(v) if v else None
+
+    @model_validator(mode="after")
+    def _infer_size_from_local_path(self) -> Self:
+        """If `size` isn't set, try to infer it from `local_path`."""
+        if (self.size is None) and self.local_path:
+            self.size = getsize(self.local_path)
+        return self
+
+    def __repr__(self) -> str:
+        # For compatibility with prior behavior, don't display `extra` if it's empty
+        exclude = None if self.extra else {"extra"}
+        repr_dict = self.model_dump(by_alias=False, exclude_none=True, exclude=exclude)
+        return f"{nameof(type(self))}({', '.join(f'{k}={v!r}' for k, v in repr_dict.items())})"
 
     @property
     def name(self) -> LogicalPath:
@@ -193,16 +169,8 @@ class ArtifactManifestEntry:
             (str): The path of the downloaded artifact entry.
         """
         artifact = self.parent_artifact()
-
-        root = root or artifact._default_root()
-        artifact._add_download_root(root)
-        path = str(Path(self.path))
-        dest_path = os.path.join(root, path)
-
-        if skip_cache:
-            override_cache_path = dest_path
-        else:
-            override_cache_path = None
+        rootdir = artifact._add_download_root(root)
+        dest_path = os.path.join(rootdir, self.path)
 
         # Skip checking the cache (and possibly downloading) if the file already exists
         # and has the digest we're expecting.
@@ -222,26 +190,28 @@ class ArtifactManifestEntry:
             if self.digest == md5_hash:
                 return FilePathStr(dest_path)
 
+        # Override the target cache path IF we're skipping the cache.
+        # Note that `override_cache_path is None` <=> `skip_cache is False`.
+        override_cache_path = FilePathStr(dest_path) if skip_cache else None
+        storage_policy = artifact.manifest.storage_policy
         if self.ref is not None:
-            cache_path = artifact.manifest.storage_policy.load_reference(
+            cache_path = storage_policy.load_reference(
                 self, local=True, dest_path=override_cache_path
             )
         else:
-            cache_path = artifact.manifest.storage_policy.load_file(
+            cache_path = storage_policy.load_file(
                 artifact, self, dest_path=override_cache_path, executor=executor
             )
 
         # Determine the final path
-        final_path = (
-            dest_path
-            if skip_cache
-            else copy_or_overwrite_changed(cache_path, dest_path)
+        final_path = FilePathStr(
+            override_cache_path or copy_or_overwrite_changed(cache_path, dest_path)
         )
 
         # Cache the checksum for future downloads
-        _write_cached_checksum(str(final_path), self.digest)
+        _write_cached_checksum(final_path, self.digest)
 
-        return FilePathStr(final_path)
+        return final_path
 
     def ref_target(self) -> FilePathStr | URIStr:
         """Get the reference URL that is targeted by this artifact entry.
@@ -254,11 +224,9 @@ class ArtifactManifestEntry:
         """
         if self.ref is None:
             raise ValueError("Only reference entries support ref_target().")
-        if self._parent_artifact is None:
+        if (parent_artifact := self._parent_artifact) is None:
             return self.ref
-        return self._parent_artifact.manifest.storage_policy.load_reference(
-            self._parent_artifact.manifest.entries[self.path], local=False
-        )
+        return parent_artifact.manifest.storage_policy.load_reference(self, local=False)
 
     def ref_url(self) -> str:
         """Get a URL to this artifact entry.
@@ -279,26 +247,13 @@ class ArtifactManifestEntry:
             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 = {
-            "path": self.path,
-            "digest": self.digest,
-        }
-        if self.size is not None:
-            contents["size"] = self.size
-        if self.ref:
-            contents["ref"] = self.ref
-        if self.birth_artifact_id:
-            contents["birthArtifactID"] = self.birth_artifact_id
-        if self.local_path:
-            contents["local_path"] = self.local_path
-        if self.skip_cache:
-            contents["skip_cache"] = self.skip_cache
-        if self.extra:
-            contents["extra"] = self.extra
-        return contents
+        return f"{_WB_ARTIFACT_SCHEME}://{b64_to_hex_id(parent_id)}/{self.path}"
+
+    def to_json(self) -> dict[str, Any]:
+        # NOTE: The method name `to_json` is a bit misleading, as this returns a
+        # python dict, NOT a JSON string. The historical name is kept for continuity,
+        # but consider deprecating this in favor of `BaseModel.model_dump()`.
+        return self.model_dump(exclude_none=True)  # type: ignore[return-value]
 
     def _is_artifact_reference(self) -> bool:
         return self.ref is not None and urlparse(self.ref).scheme == _WB_ARTIFACT_SCHEME

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -1,61 +1,48 @@
 """Artifact manifest v1."""
 
+# Older-style type annotations required for Pydantic v1 / python 3.8 compatibility.
+# ruff: noqa: UP006
+
 from __future__ import annotations
 
 from operator import itemgetter
-from typing import Any, Mapping
+from typing import Any, ClassVar, Dict, Literal, final
+
+from pydantic import Field
+from typing_extensions import Annotated
 
-from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
-from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.internal.internal_api import Api as InternalApi
 from wandb.sdk.lib.hashutil import HexMD5, _md5
 
+from .._factories import make_storage_policy
+from .._models.manifest import ArtifactManifestV1Data
+from ..artifact_manifest import ArtifactManifest
+from ..artifact_manifest_entry import ArtifactManifestEntry
+from ..storage_policy import StoragePolicy
+
 
+@final
 class ArtifactManifestV1(ArtifactManifest):
-    @classmethod
-    def version(cls) -> int:
-        return 1
+    manifest_version: Annotated[Literal[1], Field(repr=False)] = 1
+    entries: Dict[str, ArtifactManifestEntry] = Field(default_factory=dict)
+
+    storage_policy: StoragePolicy = Field(
+        default_factory=make_storage_policy, exclude=True, repr=False
+    )
 
     @classmethod
     def from_manifest_json(
-        cls, manifest_json: dict, api: InternalApi | None = None
+        cls, manifest_json: dict[str, Any], api: InternalApi | None = None
     ) -> ArtifactManifestV1:
-        if manifest_json["version"] != cls.version():
-            raise ValueError(
-                "Expected manifest version 1, got {}".format(manifest_json["version"])
-            )
-
-        storage_policy_name = manifest_json["storagePolicy"]
-        storage_policy_config = manifest_json.get("storagePolicyConfig", {})
-        storage_policy_cls = StoragePolicy.lookup_by_name(storage_policy_name)
-
-        entries: Mapping[str, ArtifactManifestEntry]
-        entries = {
-            name: ArtifactManifestEntry(
-                path=name,
-                digest=val["digest"],
-                birth_artifact_id=val.get("birthArtifactID"),
-                ref=val.get("ref"),
-                size=val.get("size"),
-                extra=val.get("extra"),
-                local_path=val.get("local_path"),
-                skip_cache=val.get("skip_cache"),
-            )
-            for name, val in manifest_json["contents"].items()
-        }
+        data = ArtifactManifestV1Data(**manifest_json)
 
+        policy_name = data.storage_policy
+        policy_cfg = data.storage_policy_config
+        policy = StoragePolicy.lookup_by_name(policy_name).from_config(policy_cfg, api)
         return cls(
-            storage_policy_cls.from_config(storage_policy_config, api=api), entries
+            manifest_version=data.version, entries=data.contents, storage_policy=policy
         )
 
-    def __init__(
-        self,
-        storage_policy: StoragePolicy,
-        entries: Mapping[str, ArtifactManifestEntry] | None = None,
-    ) -> None:
-        super().__init__(storage_policy, entries=entries)
-
     def to_manifest_json(self) -> dict:
         """This is the JSON that's stored in wandb_manifest.json.
 
@@ -64,31 +51,26 @@ class ArtifactManifestV1(ArtifactManifest):
         system. We don't need to include the local paths in the artifact manifest
         contents.
         """
-        contents = {}
-        for name, entry in sorted(self.entries.items(), key=itemgetter(0)):
-            json_entry: dict[str, Any] = {
-                "digest": entry.digest,
-            }
-            if entry.birth_artifact_id:
-                json_entry["birthArtifactID"] = entry.birth_artifact_id
-            if entry.ref:
-                json_entry["ref"] = entry.ref
-            if entry.extra:
-                json_entry["extra"] = entry.extra
-            if entry.size is not None:
-                json_entry["size"] = entry.size
-            contents[name] = json_entry
+        omit_entry_fields = {"path", "local_path", "skip_cache"}
         return {
-            "version": self.__class__.version(),
+            "version": self.manifest_version,
             "storagePolicy": self.storage_policy.name(),
-            "storagePolicyConfig": self.storage_policy.config() or {},
-            "contents": contents,
+            "storagePolicyConfig": self.storage_policy.config(),
+            "contents": {
+                path: entry.model_dump(exclude=omit_entry_fields, exclude_defaults=True)
+                for path, entry in self.entries.items()
+            },
         }
 
+    _DIGEST_HEADER: ClassVar[bytes] = b"wandb-artifact-manifest-v1\n"
+    """Encoded prefix/header for the ArtifactManifest digest."""
+
     def digest(self) -> HexMD5:
-        hasher = _md5()
-        hasher.update(b"wandb-artifact-manifest-v1\n")
+        hasher = _md5(self._DIGEST_HEADER)
         # 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())
+        for path, entry in sorted(self.entries.items(), key=itemgetter(0)):
+            hasher.update(f"{path}:{entry.digest}\n".encode())
+        return hasher.hexdigest()
+
+    def size(self) -> int:
+        return sum(entry.size for entry in self.entries.values() if entry.size)

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -79,6 +79,7 @@ class GCSHandler(StorageHandler):
         bucket, key, _ = self._parse_uri(manifest_entry.ref)
         version = manifest_entry.extra.get("versionID")
 
+        # Skip downloading an entry that corresponds to a folder
         if self._is_dir(manifest_entry):
             raise _GCSIsADirectoryError(
                 f"Unable to download GCS folder {manifest_entry.ref!r}, skipping"
@@ -132,7 +133,8 @@ class GCSHandler(StorageHandler):
         obj = self._client.bucket(bucket).get_blob(key, generation=version)
         if obj is None and version is not None:
             raise ValueError(f"Object does not exist: {path}#{version}")
-        multi = obj is None
+        # HNS buckets have blobs for directories, so we also check the blob name to see if its a directory
+        multi = obj is None or obj.name.endswith("/")
         if multi:
             start_time = time.monotonic()
             termlog(
@@ -215,11 +217,11 @@ class GCSHandler(StorageHandler):
         assert manifest_entry.ref is not None
         bucket, key, _ = self._parse_uri(manifest_entry.ref)
         bucket_obj = self._client.bucket(bucket)
-        # A gcs bucket key should end with a forward slash on gcloud, but
-        # we save these refs without the forward slash in the manifest entry
-        # so we check the size and extension, make sure its not referring to
-        # an actual file with this reference, and that the ref with the slash
-        # exists on gcloud
+        # A gcs folder key should end with a forward slash on gcloud, but
+        # we previously saved these refs without the forward slash in the manifest entry
+        # To check whether the entry corresponds to a folder, we check the size and extension,
+        # make sure there is no file with this reference, and that the ref with the slash
+        # exists on gcloud as a folder
         return key.endswith("/") or (
             not (manifest_entry.size or PurePosixPath(key).suffix)
             and bucket_obj.get_blob(key) is None

wandb/sdk/data_types/image.py

@@ -182,8 +182,8 @@ class Image(BatchableMedia):
                 unless `normalize` is set to `False`.
             - pytorch tensor should be in the format (channel, height, width)
             - NumPy array should be in the format (height, width, channel)
-            mode: The PIL mode for an image. Most common are "L", "RGB",
-                "RGBA". Full explanation at https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
+            mode: The PIL mode for an image. Most common are "L", "RGB", "RGBA".
+                Full Pillow docs for more information https://pillow.readthedocs.io/en/stable/handbook/concepts.html#modes
             caption: Label for display of image.
             grouping: The grouping number for the image.
             classes: A list of class information for the image,