wandb 0.22.1__py3-none-win_arm64.whl

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -0,0 +1,315 @@
+"""Artifact manifest entry."""
+
+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 urllib.parse import urlparse
+
+from wandb.proto.wandb_deprecated import Deprecated
+from wandb.sdk.lib.deprecate import deprecate
+from wandb.sdk.lib.filesystem import copy_or_overwrite_changed
+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"
+
+
+def _checksum_cache_path(file_path: str) -> str:
+    """Get path for checksum in central cache directory."""
+    from wandb.sdk.artifacts.artifact_file_cache import artifacts_cache_dir
+
+    # Create a unique cache key based on the file's absolute path
+    abs_path = os.path.abspath(file_path)
+    path_hash = hashlib.sha256(abs_path.encode()).hexdigest()
+
+    # Store in wandb cache directory under checksums subdirectory
+    cache_dir = artifacts_cache_dir() / "checksums"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    return str(cache_dir / f"{path_hash}.checksum")
+
+
+def _read_cached_checksum(file_path: str) -> str | None:
+    """Read checksum from cache if it exists and is valid."""
+    checksum_path = _checksum_cache_path(file_path)
+
+    try:
+        with open(file_path) as f, open(checksum_path) as f_checksum:
+            if os.path.getmtime(f_checksum.name) < os.path.getmtime(f.name):
+                # File was modified after checksum was written
+                return None
+            # Read and return the cached checksum
+            return f_checksum.read().strip()
+    except OSError:
+        # File doesn't exist or couldn't be opened
+        return None
+
+
+def _write_cached_checksum(file_path: str, checksum: str) -> None:
+    """Write checksum to cache directory."""
+    checksum_path = _checksum_cache_path(file_path)
+    try:
+        with open(checksum_path, "w") as f:
+            f.write(checksum)
+    except OSError:
+        # Non-critical failure, just log it
+        logger.debug(f"Failed to write checksum cache for {file_path!r}")
+
+
+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.
+
+        # Fast integrity check using cached checksum from persistent cache
+        with suppress(OSError):
+            if self.digest == _read_cached_checksum(dest_path):
+                return FilePathStr(dest_path)
+
+        # Fallback to computing/caching the checksum hash
+        try:
+            md5_hash = md5_file_b64(dest_path)
+        except (FileNotFoundError, IsADirectoryError):
+            logger.debug(f"unable to find {dest_path!r}, skip searching for file")
+        else:
+            _write_cached_checksum(dest_path, md5_hash)
+            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,
+            )
+
+        # Determine the final path
+        final_path = (
+            dest_path
+            if skip_cache
+            else copy_or_overwrite_changed(cache_path, dest_path)
+        )
+
+        # Cache the checksum for future downloads
+        _write_cached_checksum(str(final_path), self.digest)
+
+        return FilePathStr(final_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)

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -0,0 +1,94 @@
+"""Artifact manifest v1."""
+
+from __future__ import annotations
+
+from operator import itemgetter
+from typing import Any, Mapping
+
+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
+
+
+class ArtifactManifestV1(ArtifactManifest):
+    @classmethod
+    def version(cls) -> int:
+        return 1
+
+    @classmethod
+    def from_manifest_json(
+        cls, manifest_json: dict, 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()
+        }
+
+        return cls(
+            storage_policy_cls.from_config(storage_policy_config, api=api), entries
+        )
+
+    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.
+
+        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 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
+        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")
+        # sort by key (path)
+        for name, entry in sorted(self.entries.items(), key=itemgetter(0)):
+            hasher.update(f"{name}:{entry.digest}\n".encode())
+        return HexMD5(hasher.hexdigest())

wandb/sdk/artifacts/artifact_saver.py

@@ -0,0 +1,277 @@
+"""Artifact saver."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import json
+import os
+import tempfile
+from typing import TYPE_CHECKING, Awaitable, Sequence
+
+import wandb
+import wandb.filesync.step_prepare
+from wandb import util
+from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
+from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64
+from wandb.sdk.lib.paths import URIStr
+
+if TYPE_CHECKING:
+    from typing import Protocol
+
+    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
+
+    class SaveFn(Protocol):
+        def __call__(
+            self, entry: ArtifactManifestEntry, progress_callback: ProgressFn
+        ) -> bool:
+            pass
+
+    class SaveFnAsync(Protocol):
+        def __call__(
+            self, entry: ArtifactManifestEntry, progress_callback: ProgressFn
+        ) -> Awaitable[bool]:
+            pass
+
+
+class ArtifactSaver:
+    _server_artifact: dict | None  # TODO better define this dict
+
+    def __init__(
+        self,
+        api: InternalApi,
+        digest: str,
+        manifest_json: dict,
+        file_pusher: FilePusher,
+        is_user_created: bool = False,
+    ) -> None:
+        self._api = api
+        self._file_pusher = file_pusher
+        self._digest = digest
+        self._manifest = ArtifactManifest.from_manifest_json(
+            manifest_json,
+            api=self._api,
+        )
+        self._is_user_created = is_user_created
+        self._server_artifact = None
+
+    def save(
+        self,
+        entity: str,
+        project: str,
+        type: str,
+        name: str,
+        client_id: str,
+        sequence_client_id: str,
+        distributed_id: str | None = None,
+        finalize: bool = True,
+        metadata: dict | None = None,
+        ttl_duration_seconds: int | None = None,
+        description: str | None = None,
+        aliases: Sequence[str] | None = None,
+        tags: Sequence[str] | None = None,
+        use_after_commit: bool = False,
+        incremental: bool = False,
+        history_step: int | None = None,
+        base_id: str | None = None,
+    ) -> dict | None:
+        return self._save_internal(
+            entity,
+            project,
+            type,
+            name,
+            client_id,
+            sequence_client_id,
+            distributed_id,
+            finalize,
+            metadata,
+            ttl_duration_seconds,
+            description,
+            aliases,
+            tags,
+            use_after_commit,
+            incremental,
+            history_step,
+            base_id,
+        )
+
+    def _save_internal(
+        self,
+        entity: str,
+        project: str,
+        type: str,
+        name: str,
+        client_id: str,
+        sequence_client_id: str,
+        distributed_id: str | None = None,
+        finalize: bool = True,
+        metadata: dict | None = None,
+        ttl_duration_seconds: int | None = None,
+        description: str | None = None,
+        aliases: Sequence[str] | None = None,
+        tags: Sequence[str] | None = None,
+        use_after_commit: bool = False,
+        incremental: bool = False,
+        history_step: int | None = None,
+        base_id: str | None = None,
+    ) -> dict | None:
+        alias_specs = []
+        for alias in aliases or []:
+            alias_specs.append({"artifactCollectionName": name, "alias": alias})
+
+        tag_specs = [{"tagName": tag} for tag in tags or []]
+
+        """Returns the server artifact."""
+        self._server_artifact, latest = self._api.create_artifact(
+            type,
+            name,
+            self._digest,
+            metadata=metadata,
+            ttl_duration_seconds=ttl_duration_seconds,
+            aliases=alias_specs,
+            tags=tag_specs,
+            description=description,
+            is_user_created=self._is_user_created,
+            distributed_id=distributed_id,
+            client_id=client_id,
+            sequence_client_id=sequence_client_id,
+            history_step=history_step,
+        )
+
+        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":
+            if use_after_commit:
+                self._api.use_artifact(
+                    artifact_id,
+                    artifact_entity_name=entity,
+                    artifact_project_name=project,
+                )
+            return self._server_artifact
+        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(
+                'Unknown artifact state "{}"'.format(self._server_artifact["state"])
+            )
+
+        manifest_type = "FULL"
+        manifest_filename = "wandb_manifest.json"
+        if incremental:
+            manifest_type = "INCREMENTAL"
+            manifest_filename = "wandb_manifest.incremental.json"
+        elif distributed_id:
+            manifest_type = "PATCH"
+            manifest_filename = "wandb_manifest.patch.json"
+        artifact_manifest_id, _ = self._api.create_artifact_manifest(
+            manifest_filename,
+            "",
+            artifact_id,
+            base_artifact_id=base_id,
+            include_upload=False,
+            type=manifest_type,
+        )
+
+        step_prepare = wandb.filesync.step_prepare.StepPrepare(
+            self._api, 0.1, 0.01, 1000
+        )  # TODO: params
+        step_prepare.start()
+
+        # Upload Artifact "L1" files, the actual artifact contents
+        self._file_pusher.store_manifest_files(
+            self._manifest,
+            artifact_id,
+            lambda entry, progress_callback: self._manifest.storage_policy.store_file(
+                artifact_id,
+                artifact_manifest_id,
+                entry,
+                step_prepare,
+                progress_callback=progress_callback,
+            ),
+        )
+
+        def before_commit() -> None:
+            self._resolve_client_id_manifest_references()
+            with tempfile.NamedTemporaryFile("w+", suffix=".json", delete=False) as fp:
+                path = os.path.abspath(fp.name)
+                json.dump(self._manifest.to_manifest_json(), fp, indent=4)
+            digest = md5_file_b64(path)
+            if distributed_id or incremental:
+                # If we're in the distributed flow, we want to update the
+                # patch manifest we created with our finalized digest.
+                _, resp = self._api.update_artifact_manifest(
+                    artifact_manifest_id,
+                    digest=digest,
+                )
+            else:
+                # In the regular flow, we can recreate the full manifest with the
+                # updated digest.
+                #
+                # NOTE: We do this for backwards compatibility with older backends
+                # that don't support the 'updateArtifactManifest' API.
+                _, resp = self._api.create_artifact_manifest(
+                    manifest_filename,
+                    digest,
+                    artifact_id,
+                    base_artifact_id=base_id,
+                )
+
+            # We're duplicating the file upload logic a little, which isn't great.
+            upload_url = resp["uploadUrl"]
+            upload_headers = resp["uploadHeaders"]
+            extra_headers = {}
+            for upload_header in upload_headers:
+                key, val = upload_header.split(":", 1)
+                extra_headers[key] = val
+            with open(path, "rb") as fp2:
+                self._api.upload_file_retry(
+                    upload_url,
+                    fp2,
+                    extra_headers=extra_headers,
+                )
+
+        commit_result: concurrent.futures.Future[None] = concurrent.futures.Future()
+
+        # This will queue the commit. It will only happen after all the file uploads are done
+        self._file_pusher.commit_artifact(
+            artifact_id,
+            finalize=finalize,
+            before_commit=before_commit,
+            result_future=commit_result,
+        )
+
+        # Block until all artifact files are uploaded and the
+        # artifact is committed.
+        try:
+            commit_result.result()
+        finally:
+            step_prepare.shutdown()
+
+        if finalize and use_after_commit:
+            self._api.use_artifact(
+                artifact_id,
+                artifact_entity_name=entity,
+                artifact_project_name=project,
+            )
+
+        return self._server_artifact
+
+    def _resolve_client_id_manifest_references(self) -> None:
+        for entry_path in self._manifest.entries:
+            entry = self._manifest.entries[entry_path]
+            if entry.ref is not None:
+                if entry.ref.startswith("wandb-client-artifact:"):
+                    client_id = util.host_from_path(entry.ref)
+                    artifact_file_path = util.uri_from_path(entry.ref)
+                    artifact_id = self._api._resolve_client_id(client_id)
+                    if artifact_id is None:
+                        raise RuntimeError(f"Could not resolve client id {client_id}")
+                    entry.ref = URIStr(
+                        f"wandb-artifact://{b64_to_hex_id(B64MD5(artifact_id))}/{artifact_file_path}"
+                    )

wandb/sdk/artifacts/artifact_state.py

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