wandb 0.19.12rc1__py3-none-win32.whl → 0.20.1__py3-none-win32.whl

wandb/sdk/artifacts/artifact.py

@@ -19,16 +19,18 @@ from copy import copy
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from functools import partial
+from itertools import filterfalse
 from pathlib import PurePosixPath
-from typing import IO, Any, Dict, Iterator, Literal, Sequence, Type, cast, final
+from typing import IO, TYPE_CHECKING, Any, Iterator, Literal, Sequence, Type, final
 from urllib.parse import quote, urljoin, urlparse
 
 import requests
 
 import wandb
-from wandb import data_types, env, util
+from wandb import data_types, env
+from wandb._iterutils import one
 from wandb.apis.normalize import normalize_exceptions
-from wandb.apis.public import ArtifactCollection, ArtifactFiles, RetryingClient, Run
+from wandb.apis.public import ArtifactCollection, ArtifactFiles, Run
 from wandb.apis.public.utils import gql_compat
 from wandb.data_types import WBValue
 from wandb.errors import CommError
@@ -36,37 +38,6 @@ from wandb.errors.term import termerror, termlog, termwarn
 from wandb.proto import wandb_internal_pb2 as pb
 from wandb.proto.wandb_deprecated import Deprecated
 from wandb.sdk import wandb_setup
-from wandb.sdk.artifacts._generated.fetch_linked_artifacts import FetchLinkedArtifacts
-from wandb.sdk.artifacts._generated.operations import FETCH_LINKED_ARTIFACTS_GQL
-from wandb.sdk.artifacts._graphql_fragments import (
-    _gql_artifact_fragment,
-    omit_artifact_fields,
-)
-from wandb.sdk.artifacts._validators import (
-    LINKED_ARTIFACT_COLLECTION_TYPE,
-    _LinkArtifactFields,
-    ensure_logged,
-    ensure_not_finalized,
-    is_artifact_registry_project,
-    validate_aliases,
-    validate_artifact_name,
-    validate_tags,
-)
-from wandb.sdk.artifacts.artifact_download_logger import ArtifactDownloadLogger
-from wandb.sdk.artifacts.artifact_instance_cache import artifact_instance_cache
-from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest
-from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
-from wandb.sdk.artifacts.artifact_manifests.artifact_manifest_v1 import (
-    ArtifactManifestV1,
-)
-from wandb.sdk.artifacts.artifact_state import ArtifactState
-from wandb.sdk.artifacts.artifact_ttl import ArtifactTTL
-from wandb.sdk.artifacts.exceptions import ArtifactNotLoggedError, WaitTimeoutError
-from wandb.sdk.artifacts.staging import get_staging_dir
-from wandb.sdk.artifacts.storage_handlers.gcs_handler import _GCSIsADirectoryError
-from wandb.sdk.artifacts.storage_layout import StorageLayout
-from wandb.sdk.artifacts.storage_policies import WANDB_STORAGE_POLICY
-from wandb.sdk.artifacts.storage_policy import StoragePolicy
 from wandb.sdk.data_types._dtypes import Type as WBType
 from wandb.sdk.data_types._dtypes import TypeRegistry
 from wandb.sdk.internal.internal_api import Api as InternalApi
@@ -77,23 +48,68 @@ from wandb.sdk.lib.hashutil import B64MD5, b64_to_hex_id, md5_file_b64
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
 from wandb.sdk.lib.runid import generate_id
 from wandb.sdk.mailbox import MailboxHandle
+from wandb.util import (
+    alias_is_version_index,
+    artifact_to_json,
+    fsync_open,
+    json_dumps_safer,
+    uri_from_path,
+    vendor_setup,
+)
 
 from ._generated import (
     ADD_ALIASES_GQL,
     DELETE_ALIASES_GQL,
+    FETCH_LINKED_ARTIFACTS_GQL,
     UPDATE_ARTIFACT_GQL,
     ArtifactAliasInput,
     ArtifactCollectionAliasInput,
+    FetchLinkedArtifacts,
     TagInput,
     UpdateArtifact,
 )
+from ._graphql_fragments import _gql_artifact_fragment, omit_artifact_fields
+from ._validators import (
+    LINKED_ARTIFACT_COLLECTION_TYPE,
+    _LinkArtifactFields,
+    ensure_logged,
+    ensure_not_finalized,
+    is_artifact_registry_project,
+    validate_aliases,
+    validate_artifact_name,
+    validate_artifact_type,
+    validate_metadata,
+    validate_tags,
+    validate_ttl_duration_seconds,
+)
+from .artifact_download_logger import ArtifactDownloadLogger
+from .artifact_instance_cache import artifact_instance_cache
+from .artifact_manifest import ArtifactManifest
+from .artifact_manifest_entry import ArtifactManifestEntry
+from .artifact_manifests.artifact_manifest_v1 import ArtifactManifestV1
+from .artifact_state import ArtifactState
+from .artifact_ttl import ArtifactTTL
+from .exceptions import (
+    ArtifactNotLoggedError,
+    TooFewItemsError,
+    TooManyItemsError,
+    WaitTimeoutError,
+)
+from .staging import get_staging_dir
+from .storage_handlers.gcs_handler import _GCSIsADirectoryError
+from .storage_layout import StorageLayout
+from .storage_policies import WANDB_STORAGE_POLICY
+from .storage_policy import StoragePolicy
 
-reset_path = util.vendor_setup()
+reset_path = vendor_setup()
 
 from wandb_gql import gql  # noqa: E402
 
 reset_path()
 
+if TYPE_CHECKING:
+    from wandb.apis.public import RetryingClient
+
 logger = logging.getLogger(__name__)
 
 
@@ -121,8 +137,8 @@ class Artifact:
         type: The artifact's type. Use the type of an artifact to both organize
             and differentiate artifacts. You can use any string that contains letters,
             numbers, underscores, hyphens, and dots. Common types include `dataset` or `model`.
-            Include `model` within your type string if you want to link the artifact
-            to the W&B Model Registry.
+            Note: Some types are reserved for internal use and cannot be set by users.
+                Such types include `job` and types that start with `wandb-`.
         description: A description of the artifact. For Model or Dataset Artifacts,
             add documentation for your standardized team model or dataset card. View
             an artifact's description programmatically with the `Artifact.description`
@@ -132,7 +148,7 @@ class Artifact:
             dictionary of key-value pairs. You can specify no more than 100 total keys.
         incremental: Use `Artifact.new_draft()` method instead to modify an
             existing artifact.
-        use_as: W&B Launch specific parameter. Not recommended for general use.
+        use_as: Deprecated.
         is_link: Boolean indication of if the artifact is a linked artifact(`True`) or source artifact(`False`).
 
     Returns:
@@ -156,12 +172,10 @@ class Artifact:
                 f"Artifact name may only contain alphanumeric characters, dashes, "
                 f"underscores, and dots. Invalid name: {name}"
             )
-        if type == "job" or type.startswith("wandb-"):
-            raise ValueError(
-                "Artifact types 'job' and 'wandb-*' are reserved for internal use. "
-                "Please use a different type."
-            )
-        if incremental:
+
+        from wandb.sdk.artifacts._internal_artifact import InternalArtifact
+
+        if incremental and not isinstance(self, InternalArtifact):
             termwarn("Using experimental arg `incremental`")
 
         # Internal.
@@ -193,9 +207,9 @@ class Artifact:
         self._source_version: str | None = None
         self._source_artifact: Artifact | None = None
         self._is_link: bool = False
-        self._type: str = type
+        self._type: str = validate_artifact_type(type, name)
         self._description: str | None = description
-        self._metadata: dict = self._normalize_metadata(metadata)
+        self._metadata: dict[str, Any] = validate_metadata(metadata)
         self._ttl_duration_seconds: int | None = None
         self._ttl_is_inherited: bool = True
         self._ttl_changed: bool = False
@@ -205,7 +219,14 @@ class Artifact:
         self._saved_tags: list[str] = []
         self._distributed_id: str | None = None
         self._incremental: bool = incremental
-        self._use_as: str | None = use_as
+        if use_as is not None:
+            deprecate(
+                field_name=Deprecated.artifact__init_use_as,
+                warning_message=(
+                    "`use_as` argument is deprecated and does not affect the behaviour of `wandb.Artifact()`"
+                ),
+            )
+        self._use_as: str | None = None
         self._state: ArtifactState = ArtifactState.PENDING
         self._manifest: ArtifactManifest | _DeferredArtifactManifest | None = (
             ArtifactManifestV1(self._storage_policy)
@@ -226,8 +247,7 @@ class Artifact:
 
     @classmethod
     def _from_id(cls, artifact_id: str, client: RetryingClient) -> Artifact | None:
-        artifact = artifact_instance_cache.get(artifact_id)
-        if artifact is not None:
+        if (artifact := artifact_instance_cache.get(artifact_id)) is not None:
             return artifact
 
         query = gql(
@@ -398,23 +418,18 @@ class Artifact:
                 and obj["artifactCollection"]["name"] == collection
             ]
 
-        version_aliases = [
-            alias for alias in processed_aliases if util.alias_is_version_index(alias)
-        ]
-        other_aliases = [
-            alias
-            for alias in processed_aliases
-            if not util.alias_is_version_index(alias)
-        ]
-        if version_aliases:
-            try:
-                [version] = version_aliases
-            except ValueError:
-                raise ValueError(
-                    f"Expected at most one version alias, got {len(version_aliases)}: {version_aliases!r}"
-                )
-        else:
-            version = src_version
+        version_aliases = list(filter(alias_is_version_index, processed_aliases))
+        other_aliases = list(filterfalse(alias_is_version_index, processed_aliases))
+
+        try:
+            version = one(
+                version_aliases, too_short=TooFewItemsError, too_long=TooManyItemsError
+            )
+        except TooFewItemsError:
+            version = src_version  # default to the source version
+        except TooManyItemsError:
+            msg = f"Expected at most one version alias, got {len(version_aliases)}: {version_aliases!r}"
+            raise ValueError(msg) from None
 
         self._version = version
 
@@ -429,11 +444,11 @@ class Artifact:
         self._saved_tags = copy(tags)
 
         metadata_str = attrs["metadata"]
-        self._metadata = self._normalize_metadata(
+        self._metadata = validate_metadata(
             json.loads(metadata_str) if metadata_str else {}
         )
 
-        self._ttl_duration_seconds = _ttl_duration_seconds_from_gql(
+        self._ttl_duration_seconds = validate_ttl_duration_seconds(
             attrs.get("ttlDurationSeconds")
         )
         self._ttl_is_inherited = (
@@ -643,9 +658,10 @@ class Artifact:
         if not self.is_link:
             return self
         if self._source_artifact is None:
+            if self._client is None:
+                raise ValueError("Client is not initialized")
+
             try:
-                if self._client is None:
-                    raise ValueError("Client is not initialized")
                 artifact = self._from_name(
                     entity=self.source_entity,
                     project=self.source_project,
@@ -798,7 +814,7 @@ class Artifact:
             wandb.termwarn(
                 "Editing the metadata of this linked artifact will edit the metadata for the source artifact and it's linked artifacts as well."
             )
-        self._metadata = self._normalize_metadata(metadata)
+        self._metadata = validate_metadata(metadata)
 
     @property
     def ttl(self) -> timedelta | None:
@@ -912,6 +928,11 @@ class Artifact:
 
     @property
     def use_as(self) -> str | None:
+        """Deprecated."""
+        deprecate(
+            field_name=Deprecated.artifact__use_as,
+            warning_message=("The use_as property of Artifact is deprecated."),
+        )
         return self._use_as
 
     @property
@@ -1081,9 +1102,7 @@ class Artifact:
             with telemetry.context() as tel:
                 tel.feature.artifact_incremental = True
 
-        singleton = wandb_setup._setup(start_service=False)
-
-        if run := singleton.most_recent_active_run:
+        if run := wandb_setup.singleton().most_recent_active_run:
             # TODO: Deprecate and encourage explicit log_artifact().
             run.log_artifact(self)
         else:
@@ -1276,7 +1295,7 @@ class Artifact:
         gql_vars = {
             "artifactID": self.id,
             "description": self.description,
-            "metadata": util.json_dumps_safer(self.metadata),
+            "metadata": json_dumps_safer(self.metadata),
             "ttlDurationSeconds": self._ttl_duration_seconds_to_gql(),
             "aliases": aliases,
             "tagsToAdd": [TagInput(tag_name=t).model_dump() for t in tags_to_add],
@@ -1352,7 +1371,7 @@ class Artifact:
 
         filesystem.mkdir_exists_ok(os.path.dirname(path))
         try:
-            with util.fsync_open(path, mode, encoding) as f:
+            with fsync_open(path, mode, encoding) as f:
                 yield f
         except FileExistsError:
             raise ValueError(f"File with name {name!r} already exists at {path!r}")
@@ -1361,7 +1380,7 @@ class Artifact:
                 f"Failed to open the provided file ({type(e).__name__}: {e}). Please "
                 f"provide the proper encoding."
             )
-            raise e
+            raise
 
         self.add_file(
             path, name=name, policy="immutable", skip_cache=True, overwrite=overwrite
@@ -1427,6 +1446,7 @@ class Artifact:
         name: str | None = None,
         skip_cache: bool | None = False,
         policy: Literal["mutable", "immutable"] | None = "mutable",
+        merge: bool = False,
     ) -> None:
         """Add a local directory to the artifact.
 
@@ -1439,6 +1459,10 @@ class Artifact:
             policy: "mutable" | "immutable". By default, "mutable"
                 "mutable": Create a temporary copy of the file to prevent corruption during upload.
                 "immutable": Disable protection, rely on the user not to delete or change the file.
+            merge: If `False` (default), throws ValueError if a file was already added in a previous add_dir call
+                and its content has changed. If `True`, overwrites existing files with changed content.
+                Always adds new files and never removes files. To replace an entire directory, pass a name when adding the directory
+                using `add_dir(local_path, name=my_prefix)` and call `remove(my_prefix)` to remove the directory, then add it again.
 
         Raises:
             ArtifactFinalizedError: You cannot make changes to the current artifact
@@ -1446,7 +1470,7 @@ class Artifact:
             ValueError: Policy must be "mutable" or "immutable"
         """
         if not os.path.isdir(local_path):
-            raise ValueError("Path is not a directory: {}".format(local_path))
+            raise ValueError(f"Path is not a directory: {local_path}")
 
         termlog(
             "Adding directory to artifact ({})... ".format(
@@ -1472,6 +1496,7 @@ class Artifact:
                 path=physical_path,
                 skip_cache=skip_cache,
                 policy=policy,
+                overwrite=merge,
             )
 
         num_threads = 8
@@ -1612,8 +1637,9 @@ class Artifact:
             data_types._SavedModel,
         )
         if not isinstance(obj, allowed_types):
-            raise ValueError(
-                f"Found object of type {obj.__class__}, expected one of: {allowed_types}"
+            raise TypeError(
+                f"Found object of type {obj.__class__}, expected one of:"
+                f" {allowed_types}"
             )
 
         obj_id = id(obj)
@@ -1747,7 +1773,7 @@ class Artifact:
         name = LogicalPath(name)
         entry = self.manifest.entries.get(name) or self._get_obj_entry(name)[0]
         if entry is None:
-            raise KeyError("Path not contained in artifact: {}".format(name))
+            raise KeyError(f"Path not contained in artifact: {name}")
         entry._parent_artifact = self
         return entry
 
@@ -1775,7 +1801,7 @@ class Artifact:
             assert self._client is not None
             artifact = self._from_id(referenced_id, client=self._client)
             assert artifact is not None
-            return artifact.get(util.uri_from_path(entry.ref))
+            return artifact.get(uri_from_path(entry.ref))
 
         # Special case for wandb.Table. This is intended to be a short term
         # optimization. Since tables are likely to download many other assets in
@@ -1901,7 +1927,7 @@ class Artifact:
 
         # TODO: Create a special stream instead of relying on an existing run.
         if wandb.run is None:
-            wl = wandb.setup()
+            wl = wandb_setup.singleton()
 
             stream_id = generate_id()
 
@@ -1963,9 +1989,7 @@ class Artifact:
         if nfiles > 5000 or size > 50 * 1024 * 1024:
             log = True
             termlog(
-                "Downloading large artifact {}, {:.2f}MB. {} files... ".format(
-                    self.name, size / (1024 * 1024), nfiles
-                ),
+                f"Downloading large artifact {self.name}, {size / (1024 * 1024):.2f}MB. {nfiles} files... ",
             )
             start_time = datetime.now()
         download_logger = ArtifactDownloadLogger(nfiles=nfiles)
@@ -2053,8 +2077,8 @@ class Artifact:
         retryable_exceptions=(requests.RequestException),
     )
     def _fetch_file_urls(self, cursor: str | None, per_page: int | None = 5000) -> Any:
-        if InternalApi()._check_server_feature_with_fallback(
-            pb.ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES  # type: ignore
+        if InternalApi()._server_supports(
+            pb.ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES
         ):
             query = gql(
                 """
@@ -2165,7 +2189,7 @@ class Artifact:
 
         Args:
             root: The directory to verify. If None artifact will be downloaded to
-                './artifacts/self.name/'
+                './artifacts/self.name/'.
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
@@ -2181,16 +2205,14 @@ class Artifact:
                     self.get_entry(artifact_path)
                 except KeyError:
                     raise ValueError(
-                        "Found file {} which is not a member of artifact {}".format(
-                            full_path, self.name
-                        )
+                        f"Found file {full_path} which is not a member of artifact {self.name}"
                     )
 
         ref_count = 0
         for entry in self.manifest.entries.values():
             if entry.ref is None:
                 if md5_file_b64(os.path.join(root, entry.path)) != entry.digest:
-                    raise ValueError("Digest mismatch for file: {}".format(entry.path))
+                    raise ValueError(f"Digest mismatch for file: {entry.path}")
             else:
                 ref_count += 1
         if ref_count > 0:
@@ -2344,9 +2366,7 @@ class Artifact:
                 "Linking to a link artifact will result in directly linking to the source artifact of that link artifact."
             )
 
-        singleton = wandb_setup._setup(start_service=False)
-
-        if run := singleton.most_recent_active_run:
+        if run := wandb_setup.singleton().most_recent_active_run:
             # TODO: Deprecate and encourage explicit link_artifact().
             return run.link_artifact(self, target_path, aliases)
 
@@ -2504,7 +2524,7 @@ class Artifact:
         Returns:
             A `dict` with `string` keys representing attributes of the artifact.
         """
-        return util.artifact_to_json(self)
+        return artifact_to_json(self)
 
     @staticmethod
     def _expected_type(
@@ -2543,16 +2563,6 @@ class Artifact:
             or {}
         ).get("name")
 
-    @staticmethod
-    def _normalize_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]:
-        if metadata is None:
-            return {}
-        if not isinstance(metadata, dict):
-            raise TypeError(f"metadata must be dict, not {type(metadata)}")
-        return cast(
-            Dict[str, Any], json.loads(json.dumps(util.json_friendly_val(metadata)))
-        )
-
     def _load_manifest(self, url: str) -> ArtifactManifest:
         with requests.get(url) as response:
             response.raise_for_status()
@@ -2648,14 +2658,6 @@ class Artifact:
         return linked_artifact
 
 
-def _ttl_duration_seconds_from_gql(gql_ttl_duration_seconds: int | None) -> int | None:
-    # If gql_ttl_duration_seconds is not positive, its indicating that TTL is DISABLED(-2)
-    # gql_ttl_duration_seconds only returns None if the server is not compatible with setting Artifact TTLs
-    if gql_ttl_duration_seconds and gql_ttl_duration_seconds > 0:
-        return gql_ttl_duration_seconds
-    return None
-
-
 class _ArtifactVersionType(WBType):
     name = "artifactVersion"
     types = [Artifact]

wandb/sdk/artifacts/artifact_download_logger.py

@@ -1,5 +1,7 @@
 """Artifact download logger."""
 
+from __future__ import annotations
+
 import multiprocessing.dummy
 import time
 from typing import Callable

wandb/sdk/artifacts/artifact_saver.py

@@ -273,7 +273,5 @@ class ArtifactSaver:
                     if artifact_id is None:
                         raise RuntimeError(f"Could not resolve client id {client_id}")
                     entry.ref = URIStr(
-                        "wandb-artifact://{}/{}".format(
-                            b64_to_hex_id(B64MD5(artifact_id)), artifact_file_path
-                        )
+                        f"wandb-artifact://{b64_to_hex_id(B64MD5(artifact_id))}/{artifact_file_path}"
                     )

wandb/sdk/artifacts/artifact_state.py

@@ -1,5 +1,7 @@
 """Artifact state."""
 
+from __future__ import annotations
+
 from enum import Enum
 
 

wandb/sdk/artifacts/artifact_ttl.py

@@ -1,5 +1,7 @@
 """Artifact TTL."""
 
+from __future__ import annotations
+
 from enum import Enum
 
 

wandb/sdk/artifacts/exceptions.py

@@ -55,3 +55,17 @@ class ArtifactFinalizedError(ArtifactStatusError):
 
 class WaitTimeoutError(errors.Error):
     """Raised when wait() timeout occurs before process is finished."""
+
+
+class TooFewItemsError(ValueError):
+    """Raised when there are fewer items than expected in a collection.
+
+    Intended for internal use only.
+    """
+
+
+class TooManyItemsError(ValueError):
+    """Raised when there are more items than expected in a collection.
+
+    Intended for internal use only.
+    """

wandb/sdk/artifacts/staging.py

@@ -5,6 +5,8 @@ in order to avoid file changes corrupting the artifact. Once the upload is compl
 file should be moved to the artifact cache.
 """
 
+from __future__ import annotations
+
 import os
 
 from wandb import env

wandb/sdk/artifacts/storage_handlers/local_file_handler.py

@@ -46,9 +46,7 @@ class LocalFileHandler(StorageHandler):
         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
-                )
+                f"Local file reference: Failed to find file at path {local_path}"
             )
 
         path, hit, cache_open = self._cache.check_md5_obj_path(
@@ -140,7 +138,5 @@ class LocalFileHandler(StorageHandler):
             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)
-            )
+            raise ValueError(f'Path "{path}" must be a valid file or directory path')
         return entries

wandb/sdk/artifacts/storage_handlers/multi_handler.py

@@ -31,7 +31,7 @@ class MultiHandler(StorageHandler):
                 return handler
         if self._default_handler is not None:
             return self._default_handler
-        raise ValueError('No storage handler registered for url "{}"'.format(str(url)))
+        raise ValueError(f'No storage handler registered for url "{url!s}"')
 
     def load_path(
         self,

wandb/sdk/artifacts/storage_handlers/tracking_handler.py

@@ -59,14 +59,10 @@ class TrackingHandler(StorageHandler):
         url = urlparse(path)
         if name is None:
             raise ValueError(
-                'You must pass name="<entry_name>" when tracking references with unknown schemes. ref: {}'.format(
-                    path
-                )
+                f'You must pass name="<entry_name>" when tracking references with unknown schemes. ref: {path}'
             )
         termwarn(
-            "Artifact references with unsupported schemes cannot be checksummed: {}".format(
-                path
-            )
+            f"Artifact references with unsupported schemes cannot be checksummed: {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

@@ -116,11 +116,7 @@ class WBArtifactHandler(StorageHandler):
         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,
-            )
+            f"{self._scheme}://{b64_to_hex_id(B64MD5(target_artifact.id))}/{artifact_file_path}"
         )
 
         # Return the new entry

wandb/sdk/artifacts/storage_handlers/wb_local_artifact_handler.py

@@ -58,7 +58,7 @@ class WBLocalArtifactHandler(StorageHandler):
         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")
+            raise TypeError("Artifact passed to store_path() must be a wandb.Artifact.")
         target_entry = target_artifact.manifest.entries[target_path]  # type: ignore
         if target_entry is None:
             raise RuntimeError("Local entry not found - invalid reference")

wandb/sdk/artifacts/storage_layout.py

@@ -1,5 +1,7 @@
 """Storage layout."""
 
+from __future__ import annotations
+
 
 class StorageLayout:
     V1 = "V1"

wandb/sdk/artifacts/storage_policies/wandb_storage_policy.py

@@ -254,7 +254,7 @@ class WandbStoragePolicy(StoragePolicy):
                     if env.is_debug():
                         logger.debug(f"Error writing chunk to file: {e}")
                     download_has_error.set()
-                    raise e
+                    raise
             else:
                 raise ValueError(f"Unknown queue item type: {type(item)}")
 
@@ -340,7 +340,7 @@ class WandbStoragePolicy(StoragePolicy):
                 if env.is_debug():
                     logger.debug(f"Error downloading file: {e}")
                 download_has_error.set()
-                raise e
+                raise
             finally:
                 # Always signal the writer to stop
                 q.put(_CHUNK_QUEUE_SENTINEL)
@@ -387,7 +387,7 @@ class WandbStoragePolicy(StoragePolicy):
                 api.settings("base_url"), entity_name, md5_hex
             )
         elif storage_layout == StorageLayout.V2:
-            if api._check_server_feature_with_fallback(
+            if api._server_supports(
                 ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER  # type: ignore
             ):
                 return "{}/artifactsV2/{}/{}/{}/{}/{}/{}/{}".format(