wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -0,0 +1,208 @@
+"""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 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 = get_artifact_file_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:
+                raise ValueError(
+                    f"Exceeded {max_objects} objects tracked, pass max_objects to "
+                    f"add_reference"
+                )
+            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:
+        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,228 @@
+"""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."""
+
+    pass
+
+
+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 if manifest_entry.size is not None else 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(
+                'Generating checksum for up to %i objects with prefix "%s"... '
+                % (max_objects, 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(
+                "Exceeded %i objects tracked, pass max_objects to add_reference"
+                % max_objects
+            )
+        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.
+
+        Arguments:
+            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}://{str(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
+        )

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -0,0 +1,114 @@
+"""HTTP storage handler."""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import ParseResult
+
+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.internal.thread_local_settings import _thread_local_api_settings
+from wandb.sdk.lib.hashutil import ETag
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    import requests
+    from requests.structures import CaseInsensitiveDict
+
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class HTTPHandler(StorageHandler):
+    def __init__(self, session: requests.Session, scheme: str | None = None) -> None:
+        self._scheme = scheme or "http"
+        self._cache = get_artifact_file_cache()
+        self._session = session
+
+    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 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
+
+        response = self._session.get(
+            manifest_entry.ref,
+            stream=True,
+            cookies=_thread_local_api_settings.cookies,
+            headers=_thread_local_api_settings.headers,
+        )
+        response.raise_for_status()
+
+        digest: ETag | FilePathStr | URIStr | None
+        digest, size, extra = self._entry_from_headers(response.headers)
+        digest = digest or manifest_entry.ref
+        if manifest_entry.digest != digest:
+            raise ValueError(
+                f"Digest mismatch for url {manifest_entry.ref}: expected {manifest_entry.digest} but found {digest}"
+            )
+
+        with cache_open(mode="wb") as file:
+            for data in response.iter_content(chunk_size=16 * 1024):
+                file.write(data)
+        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]:
+        name = name or os.path.basename(path)
+        if not checksum:
+            return [ArtifactManifestEntry(path=name, ref=path, digest=path)]
+
+        with self._session.get(
+            path,
+            stream=True,
+            cookies=_thread_local_api_settings.cookies,
+            headers=_thread_local_api_settings.headers,
+        ) as response:
+            response.raise_for_status()
+            digest: ETag | FilePathStr | URIStr | None
+            digest, size, extra = self._entry_from_headers(response.headers)
+            digest = digest or path
+        return [
+            ArtifactManifestEntry(
+                path=name, ref=path, digest=digest, size=size, extra=extra
+            )
+        ]
+
+    def _entry_from_headers(
+        self, headers: CaseInsensitiveDict
+    ) -> tuple[ETag | None, int | None, dict[str, str]]:
+        response_headers = {k.lower(): v for k, v in headers.items()}
+        size = None
+        if response_headers.get("content-length", None):
+            size = int(response_headers["content-length"])
+
+        digest = response_headers.get("etag", None)
+        extra = {}
+        if digest:
+            extra["etag"] = digest
+        if digest and digest[:1] == '"' and digest[-1:] == '"':
+            digest = digest[1:-1]  # trim leading and trailing quotes around etag
+        return digest, size, extra

wandb/sdk/artifacts/storage_handlers/local_file_handler.py

@@ -0,0 +1,141 @@
+"""Local file storage handler."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import time
+from typing import TYPE_CHECKING, Sequence
+from urllib.parse import ParseResult
+
+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 import filesystem
+from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64, md5_string
+from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
+
+if TYPE_CHECKING:
+    from wandb.sdk.artifacts.artifact import Artifact
+
+
+class LocalFileHandler(StorageHandler):
+    """Handles file:// references."""
+
+    def __init__(self, scheme: str | None = None) -> None:
+        """Track files or directories on a local filesystem.
+
+        Expand directories to create an entry for each file contained.
+        """
+        self._scheme = scheme or "file"
+        self._cache = get_artifact_file_cache()
+
+    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 manifest_entry.ref is None:
+            raise ValueError(f"Cannot add path with no ref: {manifest_entry.path}")
+        local_path = util.local_file_uri_to_path(str(manifest_entry.ref))
+        if not os.path.exists(local_path):
+            raise ValueError(
+                "Local file reference: Failed to find file at path {}".format(
+                    local_path
+                )
+            )
+
+        path, hit, cache_open = self._cache.check_md5_obj_path(
+            B64MD5(manifest_entry.digest),  # TODO(spencerpearson): unsafe cast
+            manifest_entry.size if manifest_entry.size is not None else 0,
+        )
+        if hit:
+            return path
+
+        md5 = md5_file_b64(local_path)
+        if md5 != manifest_entry.digest:
+            raise ValueError(
+                f"Local file reference: Digest mismatch for path {local_path}: expected {manifest_entry.digest} but found {md5}"
+            )
+
+        filesystem.mkdir_exists_ok(os.path.dirname(path))
+
+        with cache_open() as f:
+            shutil.copy(local_path, f.name)
+        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]:
+        local_path = util.local_file_uri_to_path(path)
+        max_objects = max_objects or DEFAULT_MAX_OBJECTS
+        # We have a single file or directory
+        # Note, we follow symlinks for files contained within the directory
+        entries = []
+
+        def md5(path: str) -> B64MD5:
+            return (
+                md5_file_b64(path)
+                if checksum
+                else md5_string(str(os.stat(path).st_size))
+            )
+
+        if os.path.isdir(local_path):
+            i = 0
+            start_time = time.time()
+            if checksum:
+                termlog(
+                    'Generating checksum for up to %i files in "%s"...\n'
+                    % (max_objects, local_path),
+                    newline=False,
+                )
+            for root, _, files in os.walk(local_path):
+                for sub_path in files:
+                    i += 1
+                    if i > max_objects:
+                        raise ValueError(
+                            "Exceeded %i objects tracked, pass max_objects to add_reference"
+                            % max_objects
+                        )
+                    physical_path = os.path.join(root, sub_path)
+                    # TODO(spencerpearson): this is not a "logical path" in the sense that
+                    # `LogicalPath` returns a "logical path"; it's a relative path
+                    # **on the local filesystem**.
+                    logical_path = os.path.relpath(physical_path, start=local_path)
+                    if name is not None:
+                        logical_path = os.path.join(name, logical_path)
+
+                    entry = ArtifactManifestEntry(
+                        path=logical_path,
+                        ref=FilePathStr(os.path.join(path, logical_path)),
+                        size=os.path.getsize(physical_path),
+                        digest=md5(physical_path),
+                    )
+                    entries.append(entry)
+            if checksum:
+                termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+        elif os.path.isfile(local_path):
+            name = name or os.path.basename(local_path)
+            entry = ArtifactManifestEntry(
+                path=name,
+                ref=path,
+                size=os.path.getsize(local_path),
+                digest=md5(local_path),
+            )
+            entries.append(entry)
+        else:
+            # TODO: update error message if we don't allow directories.
+            raise ValueError(
+                'Path "{}" must be a valid file or directory path'.format(path)
+            )
+        return entries