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

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -2,14 +2,24 @@
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, Optional, Union
+from urllib.parse import urlparse
 
+import wandb
+from wandb import util
 from wandb.errors.term import termwarn
 from wandb.sdk.lib import filesystem
-from wandb.sdk.lib.hashutil import B64MD5, ETag, b64_to_hex_id, md5_file_b64
+from wandb.sdk.lib.hashutil import (
+    B64MD5,
+    ETag,
+    b64_to_hex_id,
+    hex_to_b64_id,
+    md5_file_b64,
+)
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
 
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+    from wandb.apis.public import RetryingClient
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class ArtifactManifestEntry:
@@ -23,7 +33,8 @@ class ArtifactManifestEntry:
     extra: Dict
     local_path: Optional[str]
 
-    _parent_artifact: Optional["PublicArtifact"] = None
+    _parent_artifact: Optional["Artifact"] = None
+    _download_url: Optional[str] = None
 
     def __init__(
         self,
@@ -51,7 +62,7 @@ class ArtifactManifestEntry:
         termwarn("ArtifactManifestEntry.name is deprecated, use .path instead")
         return self.path
 
-    def parent_artifact(self) -> "PublicArtifact":
+    def parent_artifact(self) -> "Artifact":
         """Get the artifact to which this artifact entry belongs.
 
         Returns:
@@ -109,8 +120,7 @@ class ArtifactManifestEntry:
         if self._parent_artifact is None:
             return self.ref
         return self._parent_artifact.manifest.storage_policy.load_reference(
-            self._parent_artifact.manifest.entries[self.path],
-            local=False,
+            self._parent_artifact.manifest.entries[self.path], local=False
         )
 
     def ref_url(self) -> str:
@@ -137,3 +147,13 @@ class ArtifactManifestEntry:
             + "/"
             + self.path
         )
+
+    def _is_artifact_reference(self) -> bool:
+        return self.ref is not None and urlparse(self.ref).scheme == "wandb-artifact"
+
+    def _get_referenced_artifact(self, client: "RetryingClient") -> "Artifact":
+        artifact: "Artifact" = wandb.Artifact._from_id(
+            hex_to_b64_id(util.host_from_path(self.ref)), client
+        )
+        assert artifact is not None
+        return artifact

wandb/sdk/artifacts/artifact_state.py

@@ -0,0 +1,10 @@
+"""Artifact state."""
+from enum import Enum
+
+
+class ArtifactState(Enum):
+    PENDING = "PENDING"
+    COMMITTED = "COMMITTED"
+    DELETED = "DELETED"
+    GARBAGE_COLLECTED = "GARBAGE_COLLECTED"
+    PENDING_DELETION = "PENDING_DELETION"

wandb/sdk/artifacts/artifacts_cache.py

@@ -15,8 +15,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     import sys
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
-    from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
+    from wandb.sdk.artifacts.artifact import Artifact
 
     if sys.version_info >= (3, 8):
         from typing import Protocol
@@ -36,8 +35,8 @@ class ArtifactsCache:
         mkdir_exists_ok(self._cache_dir)
         self._md5_obj_dir = os.path.join(self._cache_dir, "obj", "md5")
         self._etag_obj_dir = os.path.join(self._cache_dir, "obj", "etag")
-        self._artifacts_by_id: Dict[str, "ArtifactInterface"] = CappedDict()
-        self._artifacts_by_client_id: Dict[str, "LocalArtifact"] = CappedDict()
+        self._artifacts_by_id: Dict[str, "Artifact"] = CappedDict()
+        self._artifacts_by_client_id: Dict[str, "Artifact"] = CappedDict()
 
     def check_md5_obj_path(
         self, b64_md5: B64MD5, size: int
@@ -69,18 +68,18 @@ class ArtifactsCache:
         mkdir_exists_ok(os.path.dirname(path))
         return FilePathStr(path), False, opener
 
-    def get_artifact(self, artifact_id: str) -> Optional["ArtifactInterface"]:
+    def get_artifact(self, artifact_id: str) -> Optional["Artifact"]:
         return self._artifacts_by_id.get(artifact_id)
 
-    def store_artifact(self, artifact: "ArtifactInterface") -> None:
+    def store_artifact(self, artifact: "Artifact") -> None:
         if not artifact.id:
             raise ArtifactNotLoggedError(artifact, "store_artifact")
         self._artifacts_by_id[artifact.id] = artifact
 
-    def get_client_artifact(self, client_id: str) -> Optional["LocalArtifact"]:
+    def get_client_artifact(self, client_id: str) -> Optional["Artifact"]:
         return self._artifacts_by_client_id.get(client_id)
 
-    def store_client_artifact(self, artifact: "LocalArtifact") -> None:
+    def store_client_artifact(self, artifact: "Artifact") -> None:
         self._artifacts_by_client_id[artifact._client_id] = artifact
 
     def cleanup(self, target_size: int, remove_temp: bool = False) -> int:

wandb/sdk/artifacts/exceptions.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Optional
 from wandb import errors
 
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class ArtifactStatusError(AttributeError):
@@ -12,7 +12,7 @@ class ArtifactStatusError(AttributeError):
 
     def __init__(
         self,
-        artifact: Optional["ArtifactInterface"] = None,
+        artifact: Optional["Artifact"] = None,
         attr: Optional[str] = None,
         msg: str = "Artifact is in an invalid state for the requested operation.",
     ):
@@ -28,7 +28,7 @@ class ArtifactNotLoggedError(ArtifactStatusError):
     """Raised for Artifact methods or attributes only available after logging."""
 
     def __init__(
-        self, artifact: Optional["ArtifactInterface"] = None, attr: Optional[str] = None
+        self, artifact: Optional["Artifact"] = None, attr: Optional[str] = None
     ):
         super().__init__(
             artifact,
@@ -42,7 +42,7 @@ class ArtifactFinalizedError(ArtifactStatusError):
     """Raised for Artifact methods or attributes that can't be changed after logging."""
 
     def __init__(
-        self, artifact: Optional["ArtifactInterface"] = None, attr: Optional[str] = None
+        self, artifact: Optional["Artifact"] = None, attr: Optional[str] = None
     ):
         super().__init__(
             artifact,

wandb/sdk/artifacts/storage_handler.py

@@ -6,7 +6,7 @@ from wandb.sdk.lib.paths import FilePathStr, URIStr
 if TYPE_CHECKING:
     from urllib.parse import ParseResult
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 
 DEFAULT_MAX_OBJECTS = 10000
@@ -39,7 +39,7 @@ class StorageHandler:
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[str] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/azure_handler.py

@@ -4,6 +4,7 @@ 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
@@ -12,9 +13,10 @@ 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 as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class AzureHandler(StorageHandler):
@@ -45,8 +47,7 @@ class AzureHandler(StorageHandler):
         )
         version_id = manifest_entry.extra.get("versionID")
         blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient(
-            account_url,
-            credential=self._get_module("azure.identity").DefaultAzureCredential(),
+            account_url, credential=self._get_credential(account_url)
         )
         blob_client = blob_service_client.get_blob_client(
             container=container_name, blob=blob_name
@@ -89,7 +90,7 @@ class AzureHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,
@@ -104,8 +105,7 @@ class AzureHandler(StorageHandler):
             ]
 
         blob_service_client = self._get_module("azure.storage.blob").BlobServiceClient(
-            account_url,
-            credential=self._get_module("azure.identity").DefaultAzureCredential(),
+            account_url, credential=self._get_credential(account_url)
         )
         blob_client = blob_service_client.get_blob_client(
             container=container_name, blob=blob_name
@@ -157,6 +157,16 @@ class AzureHandler(StorageHandler):
         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))

wandb/sdk/artifacts/storage_handlers/gcs_handler.py

@@ -16,7 +16,7 @@ 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 as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class GCSHandler(StorageHandler):
@@ -106,7 +106,7 @@ class GCSHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -13,7 +13,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     import requests
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class HTTPHandler(StorageHandler):
@@ -69,7 +69,7 @@ class HTTPHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/local_file_handler.py

@@ -15,7 +15,7 @@ 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 as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class LocalFileHandler(StorageHandler):
@@ -66,7 +66,7 @@ class LocalFileHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/multi_handler.py

@@ -6,7 +6,7 @@ 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 as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 
 
@@ -41,7 +41,7 @@ class MultiHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[str] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -27,7 +27,7 @@ if TYPE_CHECKING:
     import boto3.s3  # type: ignore
     import boto3.session  # type: ignore
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class S3Handler(StorageHandler):
@@ -104,39 +104,42 @@ class S3Handler(StorageHandler):
         version = manifest_entry.extra.get("versionID")
 
         extra_args = {}
-        if version is None:
-            # We don't have version information so just get the latest version
-            # and fallback to listing all versions if we don't have a match.
+        if version:
+            obj = self._s3.ObjectVersion(bucket, key, version).Object()
+            extra_args["VersionId"] = version
+        else:
             obj = self._s3.Object(bucket, key)
+
+        try:
             etag = self._etag_from_obj(obj)
-            if etag != manifest_entry.digest:
-                if self.versioning_enabled(bucket):
-                    # Fallback to listing versions
-                    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")
-                            == object_version.e_tag[1:-1]
-                        ):
-                            obj = object_version.Object()
-                            extra_args["VersionId"] = object_version.version_id
-                            break
-                    if obj is None:
-                        raise ValueError(
-                            "Couldn't find object version for {}/{} matching etag {}".format(
-                                bucket, key, manifest_entry.extra.get("etag")
-                            )
-                        )
-                else:
-                    raise ValueError(
-                        f"Digest mismatch for object {manifest_entry.ref}: expected {manifest_entry.digest} but found {etag}"
+        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 or not self.versioning_enabled(bucket):
+                raise ValueError(
+                    f"Digest mismatch for object {manifest_entry.ref}: 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")
                     )
-        else:
-            obj = self._s3.ObjectVersion(bucket, key, version).Object()
-            extra_args["VersionId"] = version
+                )
 
         with cache_open(mode="wb") as f:
             obj.download_fileobj(f, ExtraArgs=extra_args)
@@ -144,7 +147,7 @@ class S3Handler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/tracking_handler.py

@@ -10,7 +10,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     from urllib.parse import ParseResult
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class TrackingHandler(StorageHandler):
@@ -47,7 +47,7 @@ class TrackingHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_handlers/wb_artifact_handler.py

@@ -3,11 +3,11 @@ import os
 from typing import TYPE_CHECKING, Optional, Sequence, Union
 from urllib.parse import urlparse
 
+import wandb
 from wandb import util
 from wandb.apis import PublicApi
 from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 from wandb.sdk.artifacts.artifacts_cache import get_artifacts_cache
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
 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
@@ -15,7 +15,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     from urllib.parse import ParseResult
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class WBArtifactHandler(StorageHandler):
@@ -62,7 +62,7 @@ class WBArtifactHandler(StorageHandler):
         artifact_id = util.host_from_path(manifest_entry.ref)
         artifact_file_path = util.uri_from_path(manifest_entry.ref)
 
-        dep_artifact = PublicArtifact.from_id(
+        dep_artifact = wandb.Artifact._from_id(
             hex_to_b64_id(artifact_id), self.client.client
         )
         assert dep_artifact is not None
@@ -76,7 +76,7 @@ class WBArtifactHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,
@@ -100,15 +100,11 @@ class WBArtifactHandler(StorageHandler):
         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 = PublicArtifact.from_id(
+            target_artifact = wandb.Artifact._from_id(
                 hex_to_b64_id(artifact_id), self.client.client
             )
             assert target_artifact is not None
 
-            # this should only have an effect if the user added the reference by url
-            # string directly (in other words they did not already load the artifact into ram.)
-            target_artifact._load_manifest()
-
             entry = target_artifact.manifest.get_entry_by_path(artifact_file_path)
             assert entry is not None
             iter_path = entry.ref

wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py

@@ -12,7 +12,7 @@ from wandb.sdk.lib.paths import FilePathStr, StrPath, URIStr
 if TYPE_CHECKING:
     from urllib.parse import ParseResult
 
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
 
 
 class WBLocalArtifactHandler(StorageHandler):
@@ -36,7 +36,7 @@ class WBLocalArtifactHandler(StorageHandler):
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[StrPath] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_policies/s3_bucket_policy.py

@@ -9,7 +9,7 @@ from wandb.sdk.artifacts.storage_handlers.tracking_handler import TrackingHandle
 from wandb.sdk.artifacts.storage_policy import StoragePolicy
 
 if TYPE_CHECKING:
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 
@@ -50,7 +50,7 @@ class __S3BucketPolicy(StoragePolicy):  # noqa: N801
 
     def store_path(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[str] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -30,12 +30,12 @@ from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 if TYPE_CHECKING:
     from wandb.filesync.step_prepare import StepPrepare
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal import progress
 
-# This makes the first sleep 1s, and then doubles it up to total times,
-# which makes for ~18 hours.
+# Sleep length: 0, 2, 4, 8, 16, 32, 64, 120, 120, 120, 120, 120, 120, 120, 120, 120
+# seconds, i.e. a total of 20min 6s.
 _REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry(
     backoff_factor=1,
     total=16,
@@ -105,7 +105,7 @@ class WandbStoragePolicy(StoragePolicy):
 
     def load_file(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         manifest_entry: "ArtifactManifestEntry",
     ) -> FilePathStr:
         path, hit, cache_open = self._cache.check_md5_obj_path(
@@ -115,17 +115,25 @@ class WandbStoragePolicy(StoragePolicy):
         if hit:
             return path
 
-        auth = None
-        if not _thread_local_api_settings.cookies:
-            auth = ("api", self._api.api_key)
-        response = self._session.get(
-            self._file_url(self._api, artifact.entity, manifest_entry),
-            auth=auth,
-            cookies=_thread_local_api_settings.cookies,
-            headers=_thread_local_api_settings.headers,
-            stream=True,
-        )
-        response.raise_for_status()
+        if manifest_entry._download_url is not None:
+            response = self._session.get(manifest_entry._download_url, stream=True)
+            try:
+                response.raise_for_status()
+            except Exception:
+                # Signed URL might have expired, fall back to fetching it one by one.
+                manifest_entry._download_url = None
+        if manifest_entry._download_url is None:
+            auth = None
+            if not _thread_local_api_settings.cookies:
+                auth = ("api", self._api.api_key)
+            response = self._session.get(
+                self._file_url(self._api, artifact.entity, manifest_entry),
+                auth=auth,
+                cookies=_thread_local_api_settings.cookies,
+                headers=_thread_local_api_settings.headers,
+                stream=True,
+            )
+            response.raise_for_status()
 
         with cache_open(mode="wb") as file:
             for data in response.iter_content(chunk_size=16 * 1024):
@@ -134,7 +142,7 @@ class WandbStoragePolicy(StoragePolicy):
 
     def store_reference(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[str] = None,
         checksum: bool = True,

wandb/sdk/artifacts/storage_policy.py

@@ -5,7 +5,7 @@ from wandb.sdk.lib.paths import FilePathStr, URIStr
 
 if TYPE_CHECKING:
     from wandb.filesync.step_prepare import StepPrepare
-    from wandb.sdk.artifacts.artifact import Artifact as ArtifactInterface
+    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
     from wandb.sdk.internal.progress import ProgressFn
 
@@ -30,7 +30,7 @@ class StoragePolicy:
         raise NotImplementedError
 
     def load_file(
-        self, artifact: "ArtifactInterface", manifest_entry: "ArtifactManifestEntry"
+        self, artifact: "Artifact", manifest_entry: "ArtifactManifestEntry"
     ) -> FilePathStr:
         raise NotImplementedError
 
@@ -57,7 +57,7 @@ class StoragePolicy:
 
     def store_reference(
         self,
-        artifact: "ArtifactInterface",
+        artifact: "Artifact",
         path: Union[URIStr, FilePathStr],
         name: Optional[str] = None,
         checksum: bool = True,