wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/artifacts/storage_handlers/multi_handler.py

@@ -0,0 +1,56 @@
+"""Multi storage handler."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import urlparse
+
+from wandb.sdk.artifacts.storage_handler import StorageHandler
+from wandb.sdk.lib.paths import FilePathStr, URIStr
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+    from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+
+
+class MultiHandler(StorageHandler):
+    _handlers: list[StorageHandler]
+
+    def __init__(
+        self,
+        handlers: list[StorageHandler] | None = None,
+        default_handler: StorageHandler | None = None,
+    ) -> None:
+        self._handlers = handlers or []
+        self._default_handler = default_handler
+
+    def _get_handler(self, url: FilePathStr | URIStr) -> StorageHandler:
+        parsed_url = urlparse(url)
+        for handler in self._handlers:
+            if handler.can_handle(parsed_url):
+                return handler
+        if self._default_handler is not None:
+            return self._default_handler
+        raise ValueError('No storage handler registered for url "{}"'.format(str(url)))
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        assert manifest_entry.ref is not None
+        handler = self._get_handler(manifest_entry.ref)
+        return handler.load_path(manifest_entry, local=local)
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: str | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        handler = self._get_handler(path)
+        return handler.store_path(
+            artifact, path, name=name, checksum=checksum, max_objects=max_objects
+        )

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -0,0 +1,300 @@
+"""S3 storage handler."""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import PurePosixPath
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import parse_qsl, urlparse
+
+from wandb import util
+from wandb.errors import CommError
+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:
+    from urllib.parse import ParseResult
+
+    # We could probably use https://pypi.org/project/boto3-stubs/ or something
+    # instead of `type:ignore`ing these boto imports, but it's nontrivial:
+    # for some reason, despite being actively maintained as of 2022-09-30,
+    # the latest release of boto3-stubs doesn't include all the features we use.
+    import boto3  # type: ignore
+    import boto3.resources.base  # type: ignore
+    import boto3.s3  # type: ignore
+    import boto3.session  # type: ignore
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class S3Handler(StorageHandler):
+    _s3: boto3.resources.base.ServiceResource | None
+    _scheme: str
+
+    def __init__(self, scheme: str | None = None) -> None:
+        self._scheme = scheme or "s3"
+        self._s3 = None
+        self._cache = get_artifact_file_cache()
+
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == self._scheme
+
+    def init_boto(self) -> boto3.resources.base.ServiceResource:
+        if self._s3 is not None:
+            return self._s3
+        boto: boto3 = util.get_module(
+            "boto3",
+            required="s3:// references requires the boto3 library, run pip install wandb[aws]",
+            lazy=False,
+        )
+        self._s3 = boto.session.Session().resource(
+            "s3",
+            endpoint_url=os.getenv("AWS_S3_ENDPOINT_URL"),
+            region_name=os.getenv("AWS_REGION"),
+        )
+        self._botocore = util.get_module("botocore")
+        return self._s3
+
+    def _parse_uri(self, uri: str) -> tuple[str, str, str | None]:
+        url = urlparse(uri)
+        query = dict(parse_qsl(url.query))
+
+        bucket = url.netloc
+        key = url.path[1:]  # strip leading slash
+        version = query.get("versionId")
+
+        return bucket, key, version
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        if not local:
+            assert manifest_entry.ref is not None
+            return manifest_entry.ref
+
+        assert manifest_entry.ref is not None
+
+        path, hit, cache_open = self._cache.check_etag_obj_path(
+            URIStr(manifest_entry.ref),
+            ETag(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            manifest_entry.size if manifest_entry.size is not None else 0,
+        )
+        if hit:
+            return path
+
+        self.init_boto()
+        assert self._s3 is not None  # mypy: unwraps optionality
+        bucket, key, _ = self._parse_uri(manifest_entry.ref)
+        version = manifest_entry.extra.get("versionID")
+
+        extra_args = {}
+        if version:
+            obj_version = self._s3.ObjectVersion(bucket, key, version)
+            extra_args["VersionId"] = version
+            obj = obj_version.Object()
+        else:
+            obj = self._s3.Object(bucket, key)
+
+        try:
+            etag = (
+                obj_version.head()["ETag"][1:-1]  # escape leading and trailing
+                if version
+                else self._etag_from_obj(obj)
+            )
+        except self._botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                raise FileNotFoundError(
+                    f"Unable to find {manifest_entry.path} at s3://{bucket}/{key}"
+                ) from e
+            raise
+
+        if etag != manifest_entry.digest:
+            # Try to match the etag with some other version.
+            if version:
+                raise ValueError(
+                    f"Digest mismatch for object {manifest_entry.ref} with version {version}: expected {manifest_entry.digest} but found {etag}"
+                )
+            obj = None
+            object_versions = self._s3.Bucket(bucket).object_versions.filter(Prefix=key)
+            for object_version in object_versions:
+                if manifest_entry.extra.get("etag") == self._etag_from_obj(
+                    object_version
+                ):
+                    obj = object_version.Object()
+                    extra_args["VersionId"] = object_version.version_id
+                    break
+            if obj is None:
+                raise FileNotFoundError(
+                    "Couldn't find object version for {}/{} matching etag {}".format(
+                        bucket, key, manifest_entry.extra.get("etag")
+                    )
+                )
+
+        with cache_open(mode="wb") as f:
+            obj.download_fileobj(f, ExtraArgs=extra_args)
+        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_boto()
+        assert self._s3 is not None  # mypy: unwraps optionality
+
+        # The passed in path might have query string parameters.
+        # We only need to care about a subset, like version, when
+        # parsing. Once we have that, we can store the rest of the
+        # metadata in the artifact entry itself.
+        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:
+            entry_path = name or (key if key != "" else bucket)
+            return [ArtifactManifestEntry(path=entry_path, ref=path, digest=path)]
+
+        # If an explicit version is specified, use that. Otherwise, use the head version.
+        objs = (
+            [self._s3.ObjectVersion(bucket, key, version).Object()]
+            if version
+            else [self._s3.Object(bucket, key)]
+        )
+        start_time = None
+        multi = False
+        if key != "":
+            try:
+                objs[0].load()
+                # S3 doesn't have real folders, however there are cases where the folder key has a valid file which will not
+                # trigger a recursive upload.
+                # we should check the object's metadata says it is a directory and do a multi file upload if it is
+                if "x-directory" in objs[0].content_type:
+                    multi = True
+            except self._botocore.exceptions.ClientError as e:
+                if e.response["Error"]["Code"] == "404":
+                    multi = True
+                else:
+                    raise CommError(
+                        f"Unable to connect to S3 ({e.response['Error']['Code']}): "
+                        f"{e.response['Error']['Message']}. Check that your "
+                        "authentication credentials are valid and that your region is "
+                        "set correctly."
+                    )
+        else:
+            multi = True
+
+        if multi:
+            start_time = time.time()
+            termlog(
+                'Generating checksum for up to %i objects in "%s/%s"... '
+                % (max_objects, bucket, key),
+                newline=False,
+            )
+            if key != "":
+                objs = (
+                    self._s3.Bucket(bucket)
+                    .objects.filter(Prefix=key)
+                    .limit(max_objects)
+                )
+            else:
+                objs = self._s3.Bucket(bucket).objects.limit(max_objects)
+        # Weird iterator scoping makes us assign this to a local function
+        size = self._size_from_obj
+        entries = [
+            self._entry_from_obj(obj, path, name, prefix=key, multi=multi)
+            for obj in objs
+            if size(obj) > 0
+        ]
+        if start_time is not None:
+            termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+        if len(entries) > max_objects:
+            raise ValueError(
+                "Exceeded %i objects tracked, pass max_objects to add_reference"
+                % max_objects
+            )
+        return entries
+
+    def _size_from_obj(self, obj: boto3.s3.Object | boto3.s3.ObjectSummary) -> int:
+        # ObjectSummary has size, Object has content_length
+        size: int
+        if hasattr(obj, "size"):
+            size = obj.size
+        else:
+            size = obj.content_length
+        return size
+
+    def _entry_from_obj(
+        self,
+        obj: boto3.s3.Object | boto3.s3.ObjectSummary,
+        path: str,
+        name: StrPath | None = None,
+        prefix: str = "",
+        multi: bool = False,
+    ) -> ArtifactManifestEntry:
+        """Create an ArtifactManifestEntry from an S3 object.
+
+        Arguments:
+            obj: The S3 object
+            path: The S3-style path (e.g.: "s3://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.key)  # the bucket key
+        posix_path = PurePosixPath(bucket) / 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}://{str(posix_ref)}"),
+            digest=ETag(self._etag_from_obj(obj)),
+            size=self._size_from_obj(obj),
+            extra=self._extra_from_obj(obj),
+        )
+
+    @staticmethod
+    def _etag_from_obj(obj: boto3.s3.Object | boto3.s3.ObjectSummary) -> ETag:
+        etag: ETag
+        etag = obj.e_tag[1:-1]  # escape leading and trailing quote
+        return etag
+
+    def _extra_from_obj(
+        self, obj: boto3.s3.Object | boto3.s3.ObjectSummary
+    ) -> dict[str, str]:
+        extra = {
+            "etag": obj.e_tag[1:-1],  # escape leading and trailing quote
+        }
+        if not hasattr(obj, "version_id"):
+            # Convert ObjectSummary to Object to get the version_id.
+            obj = self._s3.Object(obj.bucket_name, obj.key)  # type: ignore[union-attr]
+        if hasattr(obj, "version_id") and obj.version_id and obj.version_id != "null":
+            extra["versionID"] = obj.version_id
+        return extra

wandb/sdk/artifacts/storage_handlers/tracking_handler.py

@@ -0,0 +1,72 @@
+"""Tracking storage handler."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import urlparse
+
+from wandb.errors.term import termwarn
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.storage_handler import StorageHandler
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    from urllib.parse import ParseResult
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class TrackingHandler(StorageHandler):
+    def __init__(self, scheme: str | None = None) -> None:
+        """Track paths with no modification or special processing.
+
+        Useful when paths being tracked are on file systems mounted at a standardized
+        location.
+
+        For example, if the data to track is located on an NFS share mounted on
+        `/data`, then it is sufficient to just track the paths.
+        """
+        self._scheme = scheme or ""
+
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == self._scheme
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        if local:
+            # Likely a user error. The tracking handler is
+            # oblivious to the underlying paths, so it has
+            # no way of actually loading it.
+            url = urlparse(manifest_entry.ref)
+            raise ValueError(
+                f"Cannot download file at path {str(manifest_entry.ref)}, scheme {str(url.scheme)} not recognized"
+            )
+        # TODO(spencerpearson): should this go through util.to_native_slash_path
+        # instead of just getting typecast?
+        return FilePathStr(manifest_entry.path)
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: StrPath | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        url = urlparse(path)
+        if name is None:
+            raise ValueError(
+                'You must pass name="<entry_name>" when tracking references with unknown schemes. ref: {}'.format(
+                    path
+                )
+            )
+        termwarn(
+            "Artifact references with unsupported schemes cannot be checksummed: {}".format(
+                path
+            )
+        )
+        name = name or url.path[1:]  # strip leading slash
+        return [ArtifactManifestEntry(path=name, ref=path, digest=path)]

wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py

@@ -0,0 +1,135 @@
+"""WB artifact storage handler."""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import urlparse
+
+import wandb
+from wandb import util
+from wandb.apis import PublicApi
+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 StorageHandler
+from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, hex_to_b64_id
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    from urllib.parse import ParseResult
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class WBArtifactHandler(StorageHandler):
+    """Handles loading and storing Artifact reference-type files."""
+
+    _client: PublicApi | None
+
+    def __init__(self) -> None:
+        self._scheme = "wandb-artifact"
+        self._cache = get_artifact_file_cache()
+        self._client = None
+
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == self._scheme
+
+    @property
+    def client(self) -> PublicApi:
+        if self._client is None:
+            self._client = PublicApi()
+        return self._client
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        """Load the file in the specified artifact given its corresponding entry.
+
+        Download the referenced artifact; create and return a new symlink to the caller.
+
+        Arguments:
+            manifest_entry (ArtifactManifestEntry): The index entry to load
+
+        Returns:
+            (os.PathLike): A path to the file represented by `index_entry`
+        """
+        # We don't check for cache hits here. Since we have 0 for size (since this
+        # is a cross-artifact reference which and we've made the choice to store 0
+        # in the size field), we can't confirm if the file is complete. So we just
+        # rely on the dep_artifact entry's download() method to do its own cache
+        # check.
+
+        # Parse the reference path and download the artifact if needed
+        artifact_id = util.host_from_path(manifest_entry.ref)
+        artifact_file_path = util.uri_from_path(manifest_entry.ref)
+
+        dep_artifact = wandb.Artifact._from_id(
+            hex_to_b64_id(artifact_id), self.client.client
+        )
+        assert dep_artifact is not None
+        link_target_path: URIStr | FilePathStr
+        if local:
+            link_target_path = dep_artifact.get_entry(artifact_file_path).download()
+        else:
+            link_target_path = dep_artifact.get_entry(artifact_file_path).ref_target()
+
+        return link_target_path
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: StrPath | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        """Store the file or directory at the given path into the specified artifact.
+
+        Recursively resolves the reference until the result is a concrete asset.
+
+        Arguments:
+            artifact: The artifact doing the storing path (str): The path to store name
+            (str): If specified, the logical name that should map to `path`
+
+        Returns:
+            (list[ArtifactManifestEntry]): A list of manifest entries to store within
+            the artifact
+        """
+        # Recursively resolve the reference until a concrete asset is found
+        # TODO: Consider resolving server-side for performance improvements.
+        iter_path: URIStr | FilePathStr | None = path
+        while iter_path is not None and urlparse(iter_path).scheme == self._scheme:
+            artifact_id = util.host_from_path(iter_path)
+            artifact_file_path = util.uri_from_path(iter_path)
+            target_artifact = wandb.Artifact._from_id(
+                hex_to_b64_id(artifact_id), self.client.client
+            )
+            assert target_artifact is not None
+
+            entry = target_artifact.manifest.get_entry_by_path(artifact_file_path)
+            assert entry is not None
+            iter_path = entry.ref
+
+        # Create the path reference
+        assert target_artifact is not None
+        assert target_artifact.id is not None
+        path = URIStr(
+            "{}://{}/{}".format(
+                self._scheme,
+                b64_to_hex_id(B64MD5(target_artifact.id)),
+                artifact_file_path,
+            )
+        )
+
+        # Return the new entry
+        assert entry is not None
+        return [
+            ArtifactManifestEntry(
+                path=name or os.path.basename(path),
+                ref=path,
+                size=0,
+                digest=entry.digest,
+            )
+        ]

wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py

@@ -0,0 +1,74 @@
+"""WB local artifact storage handler."""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Sequence
+
+import wandb
+from wandb import util
+from wandb.sdk.artifacts.artifact_instance_cache import artifact_instance_cache
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.storage_handler import StorageHandler
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    from urllib.parse import ParseResult
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class WBLocalArtifactHandler(StorageHandler):
+    """Handles loading and storing Artifact reference-type files."""
+
+    def __init__(self) -> None:
+        self._scheme = "wandb-client-artifact"
+
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        return parsed_url.scheme == self._scheme
+
+    def load_path(
+        self,
+        manifest_entry: ArtifactManifestEntry,
+        local: bool = False,
+    ) -> URIStr | FilePathStr:
+        raise NotImplementedError(
+            "Should not be loading a path for an artifact entry with unresolved client id."
+        )
+
+    def store_path(
+        self,
+        artifact: Artifact,
+        path: URIStr | FilePathStr,
+        name: StrPath | None = None,
+        checksum: bool = True,
+        max_objects: int | None = None,
+    ) -> Sequence[ArtifactManifestEntry]:
+        """Store the file or directory at the given path within the specified artifact.
+
+        Arguments:
+            artifact: The artifact doing the storing
+            path (str): The path to store
+            name (str): If specified, the logical name that should map to `path`
+
+        Returns:
+            (list[ArtifactManifestEntry]): A list of manifest entries to store within the artifact
+        """
+        client_id = util.host_from_path(path)
+        target_path = util.uri_from_path(path)
+        target_artifact = artifact_instance_cache.get(client_id)
+        if not isinstance(target_artifact, wandb.Artifact):
+            raise RuntimeError("Local Artifact not found - invalid reference")
+        target_entry = target_artifact._manifest.entries[target_path]  # type: ignore
+        if target_entry is None:
+            raise RuntimeError("Local entry not found - invalid reference")
+
+        # Return the new entry
+        return [
+            ArtifactManifestEntry(
+                path=name or os.path.basename(path),
+                ref=path,
+                size=0,
+                digest=target_entry.digest,
+            )
+        ]

wandb/sdk/artifacts/storage_layout.py

@@ -0,0 +1,6 @@
+"""Storage layout."""
+
+
+class StorageLayout:
+    V1 = "V1"
+    V2 = "V2"

wandb/sdk/artifacts/storage_policies/__init__.py

@@ -0,0 +1,4 @@
+from wandb.sdk.artifacts.storage_policies.register import WANDB_STORAGE_POLICY
+from wandb.sdk.artifacts.storage_policies.wandb_storage_policy import WandbStoragePolicy
+
+__all__ = ["WANDB_STORAGE_POLICY", "WandbStoragePolicy"]

wandb/sdk/artifacts/storage_policies/register.py

@@ -0,0 +1 @@
+WANDB_STORAGE_POLICY = "wandb-storage-policy-v1"