wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/artifacts/artifact_manifests/artifact_manifest_v1.py

@@ -0,0 +1,92 @@
+"""Artifact manifest v1."""
+
+from __future__ import annotations
+
+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 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/artifacts/artifact_saver.py

@@ -0,0 +1,265 @@
+"""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,
+        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(
+            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,
+        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)
+            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)
+
+        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(
+                        "wandb-artifact://{}/{}".format(
+                            b64_to_hex_id(B64MD5(artifact_id)), artifact_file_path
+                        )
+                    )

wandb/sdk/artifacts/artifact_state.py

@@ -0,0 +1,11 @@
+"""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/artifacts/artifact_ttl.py

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

wandb/sdk/artifacts/exceptions.py

@@ -0,0 +1,57 @@
+"""Artifact exceptions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from wandb import errors
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    ArtifactT = TypeVar("ArtifactT", bound=Artifact)
+
+
+class ArtifactStatusError(AttributeError):
+    """Raised when an artifact is in an invalid state for the requested operation."""
+
+    def __init__(
+        self,
+        msg: str = "Artifact is in an invalid state for the requested operation.",
+        name: str | None = None,
+        obj: ArtifactT | None = None,
+    ):
+        # Follow the same pattern as AttributeError in python 3.10+ by `name/obj` attributes
+        # See: https://docs.python.org/3/library/exceptions.html#AttributeError
+        try:
+            super().__init__(msg, name=name, obj=obj)
+        except TypeError:
+            # The `name`/`obj` keyword args and attributes were only added in python >= 3.10
+            super().__init__(msg)
+            self.name = name or ""
+            self.obj = obj
+
+
+class ArtifactNotLoggedError(ArtifactStatusError):
+    """Raised for Artifact methods or attributes only available after logging."""
+
+    def __init__(self, fullname: str, obj: ArtifactT):
+        *_, name = fullname.split(".")
+        msg = (
+            f"{fullname!r} used prior to logging artifact or while in offline mode. "
+            f"Call {type(obj).wait.__qualname__}() before accessing logged artifact properties."
+        )
+        super().__init__(msg=msg, name=name, obj=obj)
+
+
+class ArtifactFinalizedError(ArtifactStatusError):
+    """Raised for Artifact methods or attributes that can't be changed after logging."""
+
+    def __init__(self, fullname: str, obj: ArtifactT):
+        *_, name = fullname.split(".")
+        msg = f"{fullname!r} used on logged artifact. Can't modify finalized artifact."
+        super().__init__(msg=msg, name=name, obj=obj)
+
+
+class WaitTimeoutError(errors.Error):
+    """Raised when wait() timeout occurs before process is finished."""

wandb/sdk/artifacts/staging.py

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

wandb/sdk/artifacts/storage_handler.py

@@ -0,0 +1,62 @@
+"""Storage handler."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+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 = 10**7
+
+
+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,
+    ) -> 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: URIStr | FilePathStr,
+        name: str | None = None,
+        checksum: bool = True,
+        max_objects: int | None = 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