wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -0,0 +1,192 @@
+"""Azure storage handler."""
+from pathlib import PurePosixPath
+from types import ModuleType
+from typing import TYPE_CHECKING, Dict, Optional, Sequence, Tuple, Union
+from urllib.parse import ParseResult, parse_qsl, urlparse
+
+import wandb
+from wandb import util
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
+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,
+    ) -> Union[URIStr, FilePathStr]:
+        assert manifest_entry.ref is not None
+        if not local:
+            return manifest_entry.ref
+
+        path, hit, cache_open = get_artifacts_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: Union[URIStr, FilePathStr],
+        name: Optional[StrPath] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = 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")
+            )
+            return [
+                self._create_entry(
+                    blob_properties,
+                    path=name or PurePosixPath(blob_name).name,
+                    ref=URIStr(
+                        f"{account_url}/{container_name}/{blob_properties.name}"
+                    ),
+                )
+            ]
+
+        entries = []
+        container_client = blob_service_client.get_container_client(container_name)
+        max_objects = max_objects or DEFAULT_MAX_OBJECTS
+        for i, blob_properties in enumerate(
+            container_client.list_blobs(name_starts_with=f"{blob_name}/")
+        ):
+            if i >= max_objects:
+                raise ValueError(
+                    f"Exceeded {max_objects} objects tracked, pass max_objects to "
+                    f"add_reference"
+                )
+            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
+    ) -> Union["azure.identity.DefaultAzureCredential", str]:
+        if (
+            wandb.run
+            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,
+        )

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -0,0 +1,224 @@
+"""GCS storage handler."""
+import base64
+import time
+from pathlib import PurePosixPath
+from typing import TYPE_CHECKING, Dict, Optional, Sequence, Tuple, Union
+from urllib.parse import ParseResult, urlparse
+
+from wandb import util
+from wandb.errors.term import termlog
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
+from wandb.sdk.artifacts.storage_handler import DEFAULT_MAX_OBJECTS, StorageHandler
+from wandb.sdk.lib.hashutil import B64MD5
+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 GCSHandler(StorageHandler):
+    _client: Optional["gcs_module.client.Client"]
+    _versioning_enabled: Optional[bool]
+
+    def __init__(self, scheme: Optional[str] = None) -> None:
+        self._scheme = scheme or "gs"
+        self._client = None
+        self._versioning_enabled = None
+        self._cache = get_artifacts_cache()
+
+    def versioning_enabled(self, bucket_path: str) -> bool:
+        if self._versioning_enabled is not None:
+            return self._versioning_enabled
+        self.init_gcs()
+        assert self._client is not None  # mypy: unwraps optionality
+        bucket = self._client.bucket(bucket_path)
+        bucket.reload()
+        self._versioning_enabled = bucket.versioning_enabled
+        return self._versioning_enabled
+
+    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, Optional[str]]:
+        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,
+    ) -> Union[URIStr, FilePathStr]:
+        if not local:
+            assert manifest_entry.ref is not None
+            return manifest_entry.ref
+
+        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
+
+        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")
+
+        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}"
+                )
+            md5 = obj.md5_hash
+            if md5 != manifest_entry.digest:
+                raise ValueError(
+                    f"Digest mismatch for object {manifest_entry.ref}: expected {manifest_entry.digest} but found {md5}"
+                )
+
+        with cache_open(mode="wb") as f:
+            obj.download_to_file(f)
+        return path
+
+    def store_path(
+        self,
+        artifact: "Artifact",
+        path: Union[URIStr, FilePathStr],
+        name: Optional[StrPath] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = 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 self.versioning_enabled(bucket) and version:
+            raise ValueError(
+                f"Specifying a versionId is not valid for s3://{bucket} as it does not have versioning enabled."
+            )
+
+        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)
+        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 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: Optional[StrPath] = 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.md5_hash,
+            size=obj.size,
+            extra=self._extra_from_obj(obj),
+        )
+
+    @staticmethod
+    def _extra_from_obj(obj: "gcs_module.blob.Blob") -> Dict[str, str]:
+        return {
+            "etag": obj.etag,
+            "versionID": obj.generation,
+        }
+
+    @staticmethod
+    def _content_addressed_path(md5: str) -> FilePathStr:
+        # TODO: is this the structure we want? not at all human
+        # readable, but that's probably OK. don't want people
+        # poking around in the bucket
+        return FilePathStr(
+            "wandb/%s" % base64.b64encode(md5.encode("ascii")).decode("ascii")
+        )

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -0,0 +1,112 @@
+"""HTTP storage handler."""
+import os
+from typing import TYPE_CHECKING, Dict, Optional, Sequence, Tuple, Union
+from urllib.parse import ParseResult
+
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
+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 wandb.sdk.artifacts.artifact import Artifact
+
+
+class HTTPHandler(StorageHandler):
+    def __init__(
+        self, session: "requests.Session", scheme: Optional[str] = None
+    ) -> None:
+        self._scheme = scheme or "http"
+        self._cache = get_artifacts_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,
+    ) -> Union[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: Optional[Union[ETag, FilePathStr, URIStr]]
+        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: Union[URIStr, FilePathStr],
+        name: Optional[StrPath] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = 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: Optional[Union[ETag, FilePathStr, URIStr]]
+            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: "requests.structures.CaseInsensitiveDict"
+    ) -> Tuple[Optional[ETag], Optional[int], 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,134 @@
+"""Local file storage handler."""
+import os
+import shutil
+import time
+from typing import TYPE_CHECKING, Optional, Sequence, Union
+from urllib.parse import ParseResult
+
+from wandb import util
+from wandb.errors.term import termlog
+from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
+from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
+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: Optional[str] = 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_artifacts_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,
+    ) -> Union[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 %s" % 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: Union[URIStr, FilePathStr],
+        name: Optional[StrPath] = None,
+        checksum: bool = True,
+        max_objects: Optional[int] = 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 "%s" must be a valid file or directory path' % path)
+        return entries