wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/artifacts/artifact_download_logger.py

@@ -0,0 +1,42 @@
+"""Artifact download logger."""
+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_manifest.py

@@ -0,0 +1,67 @@
+"""Artifact manifest."""
+from typing import TYPE_CHECKING, Dict, List, Mapping, Optional
+
+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,
+    )
+
+
+class ArtifactManifest:
+    entries: Dict[str, "ArtifactManifestEntry"]
+
+    @classmethod
+    def from_manifest_json(cls, manifest_json: Dict) -> "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)
+        raise ValueError("Invalid manifest version.")
+
+    @classmethod
+    def version(cls) -> int:
+        raise NotImplementedError
+
+    def __init__(
+        self,
+        storage_policy: "WandbStoragePolicy",
+        entries: Optional[Mapping[str, "ArtifactManifestEntry"]] = None,
+    ) -> None:
+        self.storage_policy = storage_policy
+        self.entries = dict(entries) if entries else {}
+
+    def to_manifest_json(self) -> Dict:
+        raise NotImplementedError
+
+    def digest(self) -> HexMD5:
+        raise NotImplementedError
+
+    def add_entry(self, entry: "ArtifactManifestEntry") -> None:
+        if (
+            entry.path in self.entries
+            and entry.digest != self.entries[entry.path].digest
+        ):
+            raise ValueError("Cannot add the same path twice: %s" % entry.path)
+        self.entries[entry.path] = entry
+
+    def remove_entry(self, entry: "ArtifactManifestEntry") -> None:
+        if entry.path not in self.entries:
+            raise FileNotFoundError(f"Cannot remove missing entry: '{entry.path}'")
+        del self.entries[entry.path]
+
+    def get_entry_by_path(self, path: str) -> Optional["ArtifactManifestEntry"]:
+        return self.entries.get(path)
+
+    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
+        ]

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -0,0 +1,159 @@
+"""Artifact manifest entry."""
+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 (
+    B64MD5,
+    ETag,
+    b64_to_hex_id,
+    hex_to_b64_id,
+    md5_file_b64,
+)
+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
+
+
+class ArtifactManifestEntry:
+    """A single entry in an artifact manifest."""
+
+    path: LogicalPath
+    digest: Union[B64MD5, URIStr, FilePathStr, ETag]
+    ref: Optional[Union[FilePathStr, URIStr]]
+    birth_artifact_id: Optional[str]
+    size: Optional[int]
+    extra: Dict
+    local_path: Optional[str]
+
+    _parent_artifact: Optional["Artifact"] = None
+    _download_url: Optional[str] = None
+
+    def __init__(
+        self,
+        path: StrPath,
+        digest: Union[B64MD5, URIStr, FilePathStr, ETag],
+        ref: Optional[Union[FilePathStr, URIStr]] = None,
+        birth_artifact_id: Optional[str] = None,
+        size: Optional[int] = None,
+        extra: Optional[Dict] = None,
+        local_path: Optional[StrPath] = 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
+
+    @property
+    def name(self) -> LogicalPath:
+        # TODO(hugh): add telemetry to see if anyone is still using this.
+        termwarn("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: Optional[str] = None) -> FilePathStr:
+        """Download this artifact entry to the specified root path.
+
+        Arguments:
+            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)
+        dest_path = os.path.join(root, self.path)
+
+        # Skip checking the cache (and possibly downloading) if the file already exists
+        # and has the digest we're expecting.
+        if os.path.exists(dest_path) and self.digest == md5_file_b64(dest_path):
+            return FilePathStr(dest_path)
+
+        if self.ref is not None:
+            cache_path = self._parent_artifact.manifest.storage_policy.load_reference(
+                self, local=True
+            )
+        else:
+            cache_path = self._parent_artifact.manifest.storage_policy.load_file(
+                self._parent_artifact, self
+            )
+        return FilePathStr(
+            str(filesystem.copy_or_overwrite_changed(cache_path, dest_path))
+        )
+
+    def ref_target(self) -> Union[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_path('file.txt').ref_url()
+            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
+        )
+
+    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

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -0,0 +1,91 @@
+"""Artifact manifest v1."""
+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
+
+
+class ArtifactManifestV1(ArtifactManifest):
+    @classmethod
+    def version(cls) -> int:
+        return 1
+
+    @classmethod
+    def from_manifest_json(cls, manifest_json: Dict) -> "ArtifactManifestV1":
+        if manifest_json["version"] != cls.version():
+            raise ValueError(
+                "Expected manifest version 1, got %s" % 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)
+        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 = {
+            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"),
+            )
+            for name, val in manifest_json["contents"].items()
+        }
+
+        return cls(storage_policy_cls.from_config(storage_policy_config), entries)
+
+    def __init__(
+        self,
+        storage_policy: "WandbStoragePolicy",
+        entries: Optional[Mapping[str, ArtifactManifestEntry]] = 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.
+
+        If include_local is True we also include the local paths to files. This is
+        used to represent an artifact that's waiting to be saved on the current
+        system. We don't need to include the local paths in the artifact manifest
+        contents.
+        """
+        contents = {}
+        for entry in sorted(self.entries.values(), key=lambda k: k.path):
+            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[entry.path] = json_entry
+        return {
+            "version": self.__class__.version(),
+            "storagePolicy": self.storage_policy.name(),
+            "storagePolicyConfig": self.storage_policy.config() or {},
+            "contents": contents,
+        }
+
+    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]):
+            hasher.update(f"{name}:{entry.digest}\n".encode())
+        return HexMD5(hasher.hexdigest())

wandb/sdk/{internal → artifacts}/artifact_saver.py

@@ -1,3 +1,4 @@
+"""Artifact saver."""
 import concurrent.futures
 import json
 import os
@@ -8,17 +9,17 @@ from typing import TYPE_CHECKING, Awaitable, Dict, List, Optional, Sequence
 import wandb
 import wandb.filesync.step_prepare
 from wandb import env, util
-from wandb.sdk.interface.artifacts import ArtifactManifest, ArtifactManifestEntry
+from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
 from wandb.sdk.lib.filesystem import mkdir_exists_ok
 from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64
 from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+    from wandb.sdk.internal.file_pusher import FilePusher
     from wandb.sdk.internal.internal_api import Api as InternalApi
     from wandb.sdk.internal.progress import ProgressFn
 
-    from .file_pusher import FilePusher
-
     if sys.version_info >= (3, 8):
         from typing import Protocol
     else:
@@ -26,13 +27,13 @@ if TYPE_CHECKING:
 
     class SaveFn(Protocol):
         def __call__(
-            self, entry: ArtifactManifestEntry, progress_callback: "ProgressFn"
+            self, entry: "ArtifactManifestEntry", progress_callback: "ProgressFn"
         ) -> bool:
             pass
 
     class SaveFnAsync(Protocol):
         def __call__(
-            self, entry: ArtifactManifestEntry, progress_callback: "ProgressFn"
+            self, entry: "ArtifactManifestEntry", progress_callback: "ProgressFn"
         ) -> Awaitable[bool]:
             pass
 

wandb/sdk/artifacts/artifact_state.py

@@ -0,0 +1,10 @@
+"""Artifact state."""
+from enum import Enum
+
+
+class ArtifactState(Enum):
+    PENDING = "PENDING"
+    COMMITTED = "COMMITTED"
+    DELETED = "DELETED"
+    GARBAGE_COLLECTED = "GARBAGE_COLLECTED"
+    PENDING_DELETION = "PENDING_DELETION"

wandb/sdk/{interface/artifacts/artifact_cache.py → artifacts/artifacts_cache.py}

@@ -1,11 +1,13 @@
+"""Artifact cache."""
 import contextlib
 import hashlib
 import os
 import secrets
 from typing import IO, TYPE_CHECKING, ContextManager, Dict, Generator, Optional, Tuple
 
-from wandb import env, util
-from wandb.sdk.interface.artifacts import Artifact, ArtifactNotLoggedError
+from wandb import env, termwarn, 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.sdk.lib.hashutil import B64MD5, ETag, b64_to_hex_id
 from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
@@ -13,7 +15,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     import sys
 
-    from wandb.sdk import wandb_artifacts
+    from wandb.sdk.artifacts.artifact import Artifact
 
     if sys.version_info >= (3, 8):
         from typing import Protocol
@@ -33,8 +35,8 @@ class ArtifactsCache:
         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] = {}
-        self._artifacts_by_client_id: Dict[str, "wandb_artifacts.Artifact"] = {}
+        self._artifacts_by_id: Dict[str, "Artifact"] = CappedDict()
+        self._artifacts_by_client_id: Dict[str, "Artifact"] = CappedDict()
 
     def check_md5_obj_path(
         self, b64_md5: B64MD5, size: int
@@ -74,18 +76,17 @@ class ArtifactsCache:
             raise ArtifactNotLoggedError(artifact, "store_artifact")
         self._artifacts_by_id[artifact.id] = artifact
 
-    def get_client_artifact(
-        self, client_id: str
-    ) -> Optional["wandb_artifacts.Artifact"]:
+    def get_client_artifact(self, client_id: str) -> Optional["Artifact"]:
         return self._artifacts_by_client_id.get(client_id)
 
-    def store_client_artifact(self, artifact: "wandb_artifacts.Artifact") -> None:
+    def store_client_artifact(self, artifact: "Artifact") -> None:
         self._artifacts_by_client_id[artifact._client_id] = artifact
 
-    def cleanup(self, target_size: int) -> int:
+    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:
                 try:
@@ -93,14 +94,23 @@ class ArtifactsCache:
                     stat = os.stat(path)
 
                     if file.startswith(ArtifactsCache._TMP_PREFIX):
-                        os.remove(path)
-                        bytes_reclaimed += stat.st_size
+                        if remove_temp:
+                            os.remove(path)
+                            bytes_reclaimed += stat.st_size
+                        else:
+                            temp_size += stat.st_size
                         continue
                 except OSError:
                     continue
                 paths[path] = stat
                 total_size += stat.st_size
 
+        if temp_size:
+            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:

wandb/sdk/artifacts/exceptions.py

@@ -0,0 +1,55 @@
+"""Artifact exceptions."""
+from typing import TYPE_CHECKING, Optional
+
+from wandb import errors
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class ArtifactStatusError(AttributeError):
+    """Raised when an artifact is in an invalid state for the requested operation."""
+
+    def __init__(
+        self,
+        artifact: Optional["Artifact"] = None,
+        attr: Optional[str] = None,
+        msg: str = "Artifact is in an invalid state for the requested operation.",
+    ):
+        object_name = artifact.__class__.__name__ if artifact else "Artifact"
+        method_id = f"{object_name}.{attr}" if attr else object_name
+        super().__init__(msg.format(artifact=artifact, attr=attr, method_id=method_id))
+        # Follow the same pattern as AttributeError.
+        self.obj = artifact
+        self.name = attr or ""
+
+
+class ArtifactNotLoggedError(ArtifactStatusError):
+    """Raised for Artifact methods or attributes only available after logging."""
+
+    def __init__(
+        self, artifact: Optional["Artifact"] = None, attr: Optional[str] = None
+    ):
+        super().__init__(
+            artifact,
+            attr,
+            "'{method_id}' used prior to logging artifact or while in offline mode. "
+            "Call wait() before accessing logged artifact properties.",
+        )
+
+
+class ArtifactFinalizedError(ArtifactStatusError):
+    """Raised for Artifact methods or attributes that can't be changed after logging."""
+
+    def __init__(
+        self, artifact: Optional["Artifact"] = None, attr: Optional[str] = None
+    ):
+        super().__init__(
+            artifact,
+            attr,
+            "'{method_id}' used on logged artifact. Can't modify finalized artifact.",
+        )
+
+
+class WaitTimeoutError(errors.Error):
+    """Raised when wait() timeout occurs before process is finished."""

wandb/sdk/artifacts/storage_handler.py

@@ -0,0 +1,59 @@
+"""Storage handler."""
+from typing import TYPE_CHECKING, Optional, Sequence, Union
+
+from wandb.sdk.lib.paths import FilePathStr, URIStr
+
+if TYPE_CHECKING:
+    from urllib.parse import ParseResult
+
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+
+DEFAULT_MAX_OBJECTS = 10000
+
+
+class StorageHandler:
+    def can_handle(self, parsed_url: "ParseResult") -> bool:
+        """Checks whether this handler can handle the given url.
+
+        Returns:
+            Whether this handler can handle the given url.
+        """
+        raise NotImplementedError
+
+    def load_path(
+        self,
+        manifest_entry: "ArtifactManifestEntry",
+        local: bool = False,
+    ) -> Union[URIStr, FilePathStr]:
+        """Load a file or directory given the corresponding index entry.
+
+        Args:
+            manifest_entry: The index entry to load
+            local: Whether to load the file locally or not
+
+        Returns:
+            A path to the file represented by `index_entry`
+        """
+        raise NotImplementedError
+
+    def store_path(
+        self,
+        artifact: "Artifact",
+        path: Union[URIStr, FilePathStr],
+        name: Optional[str] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = None,
+    ) -> Sequence["ArtifactManifestEntry"]:
+        """Store the file or directory at the given path to the specified artifact.
+
+        Args:
+            path: The path to store
+            name: If specified, the logical name that should map to `path`
+            checksum: Whether to compute the checksum of the file
+            max_objects: The maximum number of objects to store
+
+        Returns:
+            A list of manifest entries to store within the artifact
+        """
+        raise NotImplementedError