wandb 0.22.1__py3-none-win_arm64.whl

wandb/sdk/artifacts/artifact_ttl.py

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

wandb/sdk/artifacts/exceptions.py

@@ -0,0 +1,72 @@
+"""Artifact exceptions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from wandb import errors
+from wandb._strutils import nameof
+
+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 {nameof(obj.wait)}() 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."""
+
+
+class TooFewItemsError(ValueError):
+    """Raised when there are fewer items than expected in a collection.
+
+    Intended for internal use only.
+    """
+
+
+class TooManyItemsError(ValueError):
+    """Raised when there are more items than expected in a collection.
+
+    Intended for internal use only.
+    """

wandb/sdk/artifacts/staging.py

@@ -0,0 +1,27 @@
+"""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.
+"""
+
+from __future__ import annotations
+
+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

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -0,0 +1,214 @@
+"""Azure storage handler."""
+
+from __future__ import annotations
+
+from pathlib import PurePosixPath
+from types import ModuleType
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import ParseResult, parse_qsl, urlparse
+
+import wandb
+from wandb import util
+from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler
+from wandb.sdk.lib.hashutil import ETag
+from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
+
+if TYPE_CHECKING:
+    import azure.identity  # type: ignore
+    import azure.storage.blob  # type: ignore
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class AzureHandler(StorageHandler):
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == "https" and parsed_url.netloc.endswith(
+            ".blob.core.windows.net"
+        )
+
+    def __init__(self, scheme: str | None = None) -> None:
+        self._cache = get_artifact_file_cache()
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        assert manifest_entry.ref is not None
+        if not local:
+            return manifest_entry.ref
+
+        path, hit, cache_open = self._cache.check_etag_obj_path(
+            URIStr(manifest_entry.ref),
+            ETag(manifest_entry.digest),
+            manifest_entry.size or 0,
+        )
+        if hit:
+            return path
+
+        account_url, container_name, blob_name, query = self._parse_uri(
+            manifest_entry.ref
+        )
+        version_id = manifest_entry.extra.get("versionID")
+        blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient(
+            account_url, credential=self._get_credential(account_url)
+        )
+        blob_client = blob_service_client.get_blob_client(
+            container=container_name, blob=blob_name
+        )
+        if version_id is None:
+            # Try current version, then all versions.
+            try:
+                downloader = blob_client.download_blob(
+                    etag=manifest_entry.digest,
+                    match_condition=self._get_module(
+                        "azure.core"
+                    ).MatchConditions.IfNotModified,
+                )
+            except self._get_module("azure.core.exceptions").ResourceModifiedError:
+                container_client = blob_service_client.get_container_client(
+                    container_name
+                )
+                for blob_properties in container_client.walk_blobs(
+                    name_starts_with=blob_name, include=["versions"]
+                ):
+                    if (
+                        blob_properties.name == blob_name
+                        and blob_properties.etag == manifest_entry.digest
+                        and blob_properties.version_id is not None
+                    ):
+                        downloader = blob_client.download_blob(
+                            version_id=blob_properties.version_id
+                        )
+                        break
+                else:  # didn't break
+                    raise ValueError(
+                        f"Couldn't find blob version for {manifest_entry.ref} matching "
+                        f"etag {manifest_entry.digest}."
+                    )
+        else:
+            downloader = blob_client.download_blob(version_id=version_id)
+        with cache_open(mode="wb") as f:
+            downloader.readinto(f)
+        return path
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: StrPath | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        account_url, container_name, blob_name, query = self._parse_uri(path)
+        path = URIStr(f"{account_url}/{container_name}/{blob_name}")
+
+        if not checksum:
+            return [
+                ArtifactManifestEntry(path=name or blob_name, digest=path, ref=path)
+            ]
+
+        blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient(
+            account_url, credential=self._get_credential(account_url)
+        )
+        blob_client = blob_service_client.get_blob_client(
+            container=container_name, blob=blob_name
+        )
+        if blob_client.exists(version_id=query.get("versionId")):
+            blob_properties = blob_client.get_blob_properties(
+                version_id=query.get("versionId")
+            )
+
+            if not self._is_directory_stub(blob_properties):
+                return [
+                    self._create_entry(
+                        blob_properties,
+                        path=name or PurePosixPath(blob_name).name,
+                        ref=URIStr(
+                            f"{account_url}/{container_name}/{blob_properties.name}"
+                        ),
+                    )
+                ]
+
+        entries: list[ArtifactManifestEntry] = []
+        container_client = blob_service_client.get_container_client(container_name)
+        max_objects = max_objects or DEFAULT_MAX_OBJECTS
+        for blob_properties in container_client.list_blobs(
+            name_starts_with=f"{blob_name}/"
+        ):
+            if len(entries) >= max_objects:
+                wandb.termwarn(
+                    f"Found more than {max_objects} objects under path, limiting upload "
+                    f"to {max_objects} objects. Increase max_objects to upload more"
+                )
+                break
+            if not self._is_directory_stub(blob_properties):
+                suffix = PurePosixPath(blob_properties.name).relative_to(blob_name)
+                entries.append(
+                    self._create_entry(
+                        blob_properties,
+                        path=LogicalPath(name) / suffix if name else suffix,
+                        ref=URIStr(
+                            f"{account_url}/{container_name}/{blob_properties.name}"
+                        ),
+                    )
+                )
+
+        return entries
+
+    def _get_module(self, name: str) -> ModuleType:
+        module = util.get_module(
+            name,
+            lazy=False,
+            required="Azure references require the azure library, run "
+            "pip install wandb[azure]",
+        )
+        assert isinstance(module, ModuleType)
+        return module
+
+    def _get_credential(
+        self, account_url: str
+    ) -> azure.identity.DefaultAzureCredential | str:
+        # NOTE: Always returns default credential for reinit="create_new" runs.
+        if (
+            wandb.run
+            and wandb.run.settings.azure_account_url_to_access_key is not None
+            and account_url in wandb.run.settings.azure_account_url_to_access_key
+        ):
+            return wandb.run.settings.azure_account_url_to_access_key[account_url]
+        return self._get_module("azure.identity").DefaultAzureCredential()
+
+    def _parse_uri(self, uri: str) -> tuple[str, str, str, dict[str, str]]:
+        parsed_url = urlparse(uri)
+        query = dict(parse_qsl(parsed_url.query))
+        account_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+        _, container_name, blob_name = parsed_url.path.split("/", 2)
+        return account_url, container_name, blob_name, query
+
+    def _create_entry(
+        self,
+        blob_properties: azure.storage.blob.BlobProperties,
+        path: StrPath,
+        ref: URIStr,
+    ) -> ArtifactManifestEntry:
+        extra = {"etag": blob_properties.etag.strip('"')}
+        if blob_properties.version_id:
+            extra["versionID"] = blob_properties.version_id
+        return ArtifactManifestEntry(
+            path=path,
+            ref=ref,
+            digest=blob_properties.etag.strip('"'),
+            size=blob_properties.size,
+            extra=extra,
+        )
+
+    def _is_directory_stub(
+        self, blob_properties: azure.storage.blob.BlobProperties
+    ) -> bool:
+        return (
+            blob_properties.has_key("metadata")
+            and "hdi_isfolder" in blob_properties.metadata
+            and blob_properties.metadata["hdi_isfolder"] == "true"
+        )

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -0,0 +1,224 @@
+"""GCS storage handler."""
+
+from __future__ import annotations
+
+import time
+from pathlib import PurePosixPath
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import ParseResult, urlparse
+
+from wandb import util
+from wandb.errors.term import termlog
+from wandb.sdk.artifacts.artifact_file_cache import get_artifact_file_cache
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler
+from wandb.sdk.lib.hashutil import ETag
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    import google.cloud.storage as gcs_module  # type: ignore
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class _GCSIsADirectoryError(Exception):
+    """Raised when we try to download a GCS folder."""
+
+
+class GCSHandler(StorageHandler):
+    _client: gcs_module.client.Client | None
+
+    def __init__(self, scheme: str | None = None) -> None:
+        self._scheme = scheme or "gs"
+        self._client = None
+        self._cache = get_artifact_file_cache()
+
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == self._scheme
+
+    def init_gcs(self) -> gcs_module.client.Client:
+        if self._client is not None:
+            return self._client
+        storage = util.get_module(
+            "google.cloud.storage",
+            required="gs:// references requires the google-cloud-storage library, run pip install wandb[gcp]",
+        )
+        self._client = storage.Client()
+        return self._client
+
+    def _parse_uri(self, uri: str) -> tuple[str, str, str | None]:
+        url = urlparse(uri)
+        bucket = url.netloc
+        key = url.path[1:]
+        version = url.fragment if url.fragment else None
+        return bucket, key, version
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        assert manifest_entry.ref is not None
+        if not local:
+            return manifest_entry.ref
+
+        path, hit, cache_open = self._cache.check_etag_obj_path(
+            url=URIStr(manifest_entry.ref),
+            etag=ETag(manifest_entry.digest),
+            size=manifest_entry.size or 0,
+        )
+        if hit:
+            return path
+
+        self.init_gcs()
+        assert self._client is not None  # mypy: unwraps optionality
+        assert manifest_entry.ref is not None
+        bucket, key, _ = self._parse_uri(manifest_entry.ref)
+        version = manifest_entry.extra.get("versionID")
+
+        if self._is_dir(manifest_entry):
+            raise _GCSIsADirectoryError(
+                f"Unable to download GCS folder {manifest_entry.ref!r}, skipping"
+            )
+
+        obj = None
+        # First attempt to get the generation specified, this will return None if versioning is not enabled
+        if version is not None:
+            obj = self._client.bucket(bucket).get_blob(key, generation=version)
+
+        if obj is None:
+            # Object versioning is disabled on the bucket, so just get
+            # the latest version and make sure the MD5 matches.
+            obj = self._client.bucket(bucket).get_blob(key)
+            if obj is None:
+                raise ValueError(
+                    f"Unable to download object {manifest_entry.ref} with generation {version}"
+                )
+            if obj.etag != manifest_entry.digest:
+                raise ValueError(
+                    f"Digest mismatch for object {manifest_entry.ref}: "
+                    f"expected {manifest_entry.digest} but found {obj.etag}"
+                )
+
+        with cache_open(mode="wb") as f:
+            obj.download_to_file(f)
+        return path
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: StrPath | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        self.init_gcs()
+        assert self._client is not None  # mypy: unwraps optionality
+
+        # After parsing any query params / fragments for additional context,
+        # such as version identifiers, pare down the path to just the bucket
+        # and key.
+        bucket, key, version = self._parse_uri(path)
+        path = URIStr(f"{self._scheme}://{bucket}/{key}")
+        max_objects = max_objects or DEFAULT_MAX_OBJECTS
+
+        if not checksum:
+            return [ArtifactManifestEntry(path=name or key, ref=path, digest=path)]
+
+        start_time = None
+        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
+        if multi:
+            start_time = time.time()
+            termlog(
+                f'Generating checksum for up to {max_objects} objects with prefix "{key}"... ',
+                newline=False,
+            )
+            objects = self._client.bucket(bucket).list_blobs(
+                prefix=key, max_results=max_objects
+            )
+        else:
+            objects = [obj]
+
+        entries = [
+            self._entry_from_obj(obj, path, name, prefix=key, multi=multi)
+            for obj in objects
+            if not obj.name.endswith("/")
+        ]
+        if start_time is not None:
+            termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+        if len(entries) > max_objects:
+            raise ValueError(
+                f"Exceeded {max_objects} objects tracked, pass max_objects to add_reference"
+            )
+        return entries
+
+    def _entry_from_obj(
+        self,
+        obj: gcs_module.blob.Blob,
+        path: str,
+        name: StrPath | None = None,
+        prefix: str = "",
+        multi: bool = False,
+    ) -> ArtifactManifestEntry:
+        """Create an ArtifactManifestEntry from a GCS object.
+
+        Args:
+            obj: The GCS object
+            path: The GCS-style path (e.g.: "gs://bucket/file.txt")
+            name: The user assigned name, or None if not specified
+            prefix: The prefix to add (will be the same as `path` for directories)
+            multi: Whether or not this is a multi-object add.
+        """
+        bucket, key, _ = self._parse_uri(path)
+
+        # Always use posix paths, since that's what S3 uses.
+        posix_key = PurePosixPath(obj.name)  # the bucket key
+        posix_path = PurePosixPath(bucket) / PurePosixPath(
+            key
+        )  # the path, with the scheme stripped
+        posix_prefix = PurePosixPath(prefix)  # the prefix, if adding a prefix
+        posix_name = PurePosixPath(name or "")
+        posix_ref = posix_path
+
+        if name is None:
+            # We're adding a directory (prefix), so calculate a relative path.
+            if str(posix_prefix) in str(posix_key) and posix_prefix != posix_key:
+                posix_name = posix_key.relative_to(posix_prefix)
+                posix_ref = posix_path / posix_name
+            else:
+                posix_name = PurePosixPath(posix_key.name)
+                posix_ref = posix_path
+        elif multi:
+            # We're adding a directory with a name override.
+            relpath = posix_key.relative_to(posix_prefix)
+            posix_name = posix_name / relpath
+            posix_ref = posix_path / relpath
+        return ArtifactManifestEntry(
+            path=posix_name,
+            ref=URIStr(f"{self._scheme}://{posix_ref}"),
+            digest=obj.etag,
+            size=obj.size,
+            extra={"versionID": obj.generation},
+        )
+
+    def _is_dir(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+    ) -> bool:
+        assert self._client is not None
+        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
+        return key.endswith("/") or (
+            not (manifest_entry.size or PurePosixPath(key).suffix)
+            and bucket_obj.get_blob(key) is None
+            and bucket_obj.get_blob(f"{key}/") is not None
+        )