wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/sdk/artifacts/artifact_download_logger.py

@@ -0,0 +1,45 @@
+"""Artifact download logger."""
+
+from __future__ import annotations
+
+import multiprocessing.dummy
+import time
+from typing import Callable
+
+from wandb.errors.term import termlog
+
+
+class ArtifactDownloadLogger:
+    def __init__(
+        self,
+        nfiles: int,
+        clock_for_testing: Callable[[], float] = time.monotonic,
+        termlog_for_testing: Callable[..., None] = termlog,
+    ) -> None:
+        self._nfiles = nfiles
+        self._clock = clock_for_testing
+        self._termlog = termlog_for_testing
+
+        self._n_files_downloaded = 0
+        self._spinner_index = 0
+        self._last_log_time = self._clock()
+        self._lock = multiprocessing.dummy.Lock()
+
+    def notify_downloaded(self) -> None:
+        with self._lock:
+            self._n_files_downloaded += 1
+            if self._n_files_downloaded == self._nfiles:
+                self._termlog(
+                    f"  {self._nfiles} of {self._nfiles} files downloaded.  ",
+                    # ^ trailing spaces to wipe out ellipsis from previous logs
+                    newline=True,
+                )
+                self._last_log_time = self._clock()
+            elif self._clock() - self._last_log_time > 0.1:
+                self._spinner_index += 1
+                spinner = r"-\|/"[self._spinner_index % 4]
+                self._termlog(
+                    f"{spinner} {self._n_files_downloaded} of {self._nfiles} files downloaded...\r",
+                    newline=False,
+                )
+                self._last_log_time = self._clock()

wandb/sdk/artifacts/artifact_file_cache.py

@@ -0,0 +1,251 @@
+"""Artifact cache."""
+
+from __future__ import annotations
+
+import contextlib
+import errno
+import hashlib
+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
+
+import wandb
+from wandb import env, util
+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
+
+
+class Opener(Protocol):
+    def __call__(self, mode: str = ...) -> ContextManager[IO]:
+        pass
+
+
+def _get_sys_umask_threadsafe() -> int:
+    # Workaround to get the current system umask, since
+    # - `os.umask()` isn't thread-safe
+    # - we don't want to inadvertently change the umask of the current process
+    # See: https://stackoverflow.com/questions/53227072/reading-umask-thread-safe
+    umask_cmd = (sys.executable, "-c", "import os; print(os.umask(22))")
+    return int(subprocess.check_output(umask_cmd))
+
+
+class ArtifactFileCache:
+    def __init__(self, cache_dir: StrPath) -> None:
+        self._cache_dir = Path(cache_dir)
+        self._obj_dir = self._cache_dir / "obj"
+        self._temp_dir = self._cache_dir / "tmp"
+        self._ensure_write_permissions()
+
+        # NamedTemporaryFile sets the file mode to 600 [1], we reset to the default.
+        # [1] https://stackoverflow.com/questions/10541760/can-i-set-the-umask-for-tempfile-namedtemporaryfile-in-python
+        self._sys_umask = _get_sys_umask_threadsafe()
+
+        self._override_cache_path: StrPath | None = None
+
+    def check_md5_obj_path(
+        self, b64_md5: B64MD5, size: int
+    ) -> tuple[FilePathStr, bool, Opener]:
+        # Check if we're using vs skipping the cache
+        if self._override_cache_path is not None:
+            skip_cache = True
+            path = Path(self._override_cache_path)
+        else:
+            skip_cache = False
+            hex_md5 = b64_to_hex_id(b64_md5)
+            path = self._obj_dir / "md5" / hex_md5[:2] / hex_md5[2:]
+        return self._check_or_create(path, size, skip_cache=skip_cache)
+
+    # TODO(spencerpearson): this method at least needs its signature changed.
+    # An ETag is not (necessarily) a checksum.
+    def check_etag_obj_path(
+        self,
+        url: URIStr,
+        etag: ETag,
+        size: int,
+    ) -> tuple[FilePathStr, bool, Opener]:
+        # Check if we're using vs skipping the cache
+        if self._override_cache_path is not None:
+            skip_cache = True
+            path = Path(self._override_cache_path)
+        else:
+            skip_cache = False
+            hexhash = hashlib.sha256(
+                hashlib.sha256(url.encode("utf-8")).digest()
+                + hashlib.sha256(etag.encode("utf-8")).digest()
+            ).hexdigest()
+            path = self._obj_dir / "etag" / hexhash[:2] / hexhash[2:]
+        return self._check_or_create(path, size, skip_cache=skip_cache)
+
+    def _check_or_create(
+        self, path: Path, size: int, skip_cache: bool = False
+    ) -> tuple[FilePathStr, bool, Opener]:
+        opener = self._opener(path, size, skip_cache=skip_cache)
+        hit = path.is_file() and path.stat().st_size == size
+        return FilePathStr(str(path)), hit, opener
+
+    def cleanup(
+        self,
+        target_size: int | None = None,
+        remove_temp: bool = False,
+        target_fraction: float | None = None,
+    ) -> 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.
+            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.
+            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.
+
+        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 is not None and target_size < 0:
+            raise ValueError("target_size must be non-negative")
+        if target_fraction is not None and (target_fraction < 0 or target_fraction > 1):
+            raise ValueError("target_fraction must be between 0 and 1")
+
+        bytes_reclaimed = 0
+        total_size = 0
+        temp_size = 0
+
+        # 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:
+                    os.remove(entry.path)
+                    bytes_reclaimed += size
+                except OSError:
+                    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."
+            )
+
+        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(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 cache cleanup --remove-temp` to remove temporary files."
+            )
+
+        return bytes_reclaimed
+
+    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
+
+        wandb.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 _opener(self, path: Path, size: int, skip_cache: bool = False) -> Opener:
+        @contextlib.contextmanager
+        def atomic_open(mode: str = "w") -> Iterator[IO]:
+            if "a" in mode:
+                raise ValueError("Appending to cache files is not supported")
+
+            if skip_cache:
+                # We skip the cache, but we'll still need an intermediate, temporary file to ensure atomicity.
+                # Put the temp file in the same root as the destination file in an attempt to avoid moving/copying
+                # across filesystems.
+                temp_dir = path.parent
+            else:
+                self._reserve_space(size)
+                temp_dir = self._temp_dir
+
+            temp_dir.mkdir(parents=True, exist_ok=True)
+            temp_file = NamedTemporaryFile(dir=temp_dir, mode=mode, delete=False)
+            try:
+                yield temp_file
+                temp_file.close()
+                os.chmod(temp_file.name, 0o666 & ~self._sys_umask)
+                path.parent.mkdir(parents=True, exist_ok=True)
+                os.replace(temp_file.name, path)
+            except Exception:
+                os.remove(temp_file.name)
+                raise
+
+        return atomic_open
+
+    def _ensure_write_permissions(self) -> None:
+        """Raise an error if we cannot write to the cache directory."""
+        try:
+            self._temp_dir.mkdir(parents=True, exist_ok=True)
+            with NamedTemporaryFile(dir=self._temp_dir) as f:
+                f.write(b"wandb")
+        except PermissionError as e:
+            raise PermissionError(
+                f"Unable to write to {self._cache_dir}. "
+                "Ensure that the current user has write permissions."
+            ) from e
+
+
+# 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:
+    return _build_artifact_file_cache(env.get_cache_dir() / "artifacts")

wandb/sdk/artifacts/artifact_instance_cache.py

@@ -0,0 +1,17 @@
+"""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 __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+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_instance_cache: dict[str, Artifact] = CappedDict(100)

wandb/sdk/artifacts/artifact_manifest.py

@@ -0,0 +1,76 @@
+"""Artifact manifest."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Mapping
+
+from wandb.sdk.internal.internal_api import Api as InternalApi
+from wandb.sdk.lib.hashutil import HexMD5
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+    from wandb.sdk.artifacts.storage_policy import StoragePolicy
+
+
+class ArtifactManifest:
+    entries: dict[str, ArtifactManifestEntry]
+
+    @classmethod
+    def from_manifest_json(
+        cls, manifest_json: dict, api: InternalApi | None = None
+    ) -> ArtifactManifest:
+        if "version" not in manifest_json:
+            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:
+        raise NotImplementedError
+
+    def digest(self) -> HexMD5:
+        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 (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:
+        try:
+            del self.entries[entry.path]
+        except LookupError:
+            raise FileNotFoundError(f"Cannot remove missing entry: '{entry.path}'")
+
+    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}/")
+        ]

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -0,0 +1,258 @@
+"""Artifact manifest entry."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import json
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+
+from wandb.proto.wandb_deprecated import Deprecated
+from wandb.sdk.lib import filesystem
+from wandb.sdk.lib.deprecate import deprecate
+from wandb.sdk.lib.hashutil import (
+    B64MD5,
+    ETag,
+    b64_to_hex_id,
+    hex_to_b64_id,
+    md5_file_b64,
+)
+from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from typing_extensions import TypedDict
+
+    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
+
+
+_WB_ARTIFACT_SCHEME = "wandb-artifact"
+
+
+class ArtifactManifestEntry:
+    """A single entry in an artifact manifest."""
+
+    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.
+        """
+        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
+        )
+
+    @property
+    def name(self) -> LogicalPath:
+        """Deprecated; use `path` instead."""
+        deprecate(
+            field_name=Deprecated.artifactmanifestentry__name,
+            warning_message="ArtifactManifestEntry.name is deprecated, use .path instead.",
+        )
+        return self.path
+
+    def parent_artifact(self) -> Artifact:
+        """Get the artifact to which this artifact entry belongs.
+
+        Returns:
+            (PublicArtifact): The parent artifact
+        """
+        if self._parent_artifact is None:
+            raise NotImplementedError
+        return self._parent_artifact
+
+    def download(
+        self,
+        root: str | None = None,
+        skip_cache: bool | None = None,
+        executor: concurrent.futures.Executor | None = None,
+        multipart: bool | None = None,
+    ) -> FilePathStr:
+        """Download this artifact entry to the specified root path.
+
+        Args:
+            root: (str, optional) The root path in which to download this
+                artifact entry. Defaults to the artifact's root.
+
+        Returns:
+            (str): The path of the downloaded artifact entry.
+        """
+        if self._parent_artifact is None:
+            raise NotImplementedError
+
+        root = root or self._parent_artifact._default_root()
+        self._parent_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
+
+        # Skip checking the cache (and possibly downloading) if the file already exists
+        # and has the digest we're expecting.
+        try:
+            md5_hash = md5_file_b64(dest_path)
+        except (FileNotFoundError, IsADirectoryError):
+            logger.debug(f"unable to find {dest_path}, skip searching for file")
+        else:
+            if self.digest == md5_hash:
+                return FilePathStr(dest_path)
+
+        if self.ref is not None:
+            cache_path = self._parent_artifact.manifest.storage_policy.load_reference(
+                self, local=True, dest_path=override_cache_path
+            )
+        else:
+            cache_path = self._parent_artifact.manifest.storage_policy.load_file(
+                self._parent_artifact,
+                self,
+                dest_path=override_cache_path,
+                executor=executor,
+                multipart=multipart,
+            )
+
+        if skip_cache:
+            return FilePathStr(dest_path)
+        else:
+            return FilePathStr(
+                str(filesystem.copy_or_overwrite_changed(cache_path, dest_path))
+            )
+
+    def ref_target(self) -> FilePathStr | URIStr:
+        """Get the reference URL that is targeted by this artifact entry.
+
+        Returns:
+            (str): The reference URL of this artifact entry.
+
+        Raises:
+            ValueError: If this artifact entry was not a reference.
+        """
+        if self.ref is None:
+            raise ValueError("Only reference entries support ref_target().")
+        if 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
+        )
+
+    def ref_url(self) -> str:
+        """Get a URL to this artifact entry.
+
+        These URLs can be referenced by another artifact.
+
+        Returns:
+            (str): A URL representing this artifact entry.
+
+        Examples:
+            Basic usage
+            ```
+            ref_url = source_artifact.get_entry("file.txt").ref_url()
+            derived_artifact.add_reference(ref_url)
+            ```
+        """
+        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 = {
+            "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
+
+    def _is_artifact_reference(self) -> bool:
+        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():
+            return None
+        return hex_to_b64_id(urlparse(self.ref).netloc)