wandb 0.22.0__py3-none-win_amd64.whl → 0.22.2__py3-none-win_amd64.whl

wandb/sdk/artifacts/artifact.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import atexit
-import concurrent.futures
 import contextlib
 import json
 import logging
@@ -15,10 +14,10 @@ import stat
 import tempfile
 import time
 from collections import deque
+from concurrent.futures import Executor, ThreadPoolExecutor, as_completed
 from copy import copy
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass, replace
 from datetime import timedelta
-from functools import partial
 from itertools import filterfalse
 from pathlib import Path, PurePosixPath
 from typing import (
@@ -30,7 +29,6 @@ from typing import (
     Literal,
     Sequence,
     Type,
-    cast,
     final,
 )
 from urllib.parse import quote, urljoin, urlparse
@@ -51,6 +49,7 @@ 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.storage_policies._multipart import should_multipart_download
 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
@@ -109,10 +108,11 @@ from ._generated import (
     TagInput,
     UpdateArtifact,
 )
-from ._graphql_fragments import omit_artifact_fields
+from ._gqlutils import omit_artifact_fields, supports_enable_tracking_var, type_info
 from ._validators import (
     LINKED_ARTIFACT_COLLECTION_TYPE,
     ArtifactPath,
+    FullArtifactPath,
     _LinkArtifactFields,
     ensure_logged,
     ensure_not_finalized,
@@ -210,6 +210,7 @@ class Artifact:
         metadata: dict[str, Any] | None = None,
         incremental: bool = False,
         use_as: str | None = None,
+        storage_region: str | None = None,
     ) -> None:
         if not re.match(r"^[a-zA-Z0-9_\-.]+$", name):
             raise ValueError(
@@ -268,7 +269,7 @@ class Artifact:
         self._use_as: str | None = None
         self._state: ArtifactState = ArtifactState.PENDING
         self._manifest: ArtifactManifest | _DeferredArtifactManifest | None = (
-            ArtifactManifestV1(storage_policy=make_storage_policy())
+            ArtifactManifestV1(storage_policy=make_storage_policy(storage_region))
         )
         self._commit_hash: str | None = None
         self._file_count: int | None = None
@@ -286,36 +287,33 @@ class Artifact:
 
     @classmethod
     def _from_id(cls, artifact_id: str, client: RetryingClient) -> Artifact | None:
-        if (artifact := artifact_instance_cache.get(artifact_id)) is not None:
-            return artifact
+        if cached_artifact := artifact_instance_cache.get(artifact_id):
+            return cached_artifact
 
-        query = gql_compat(ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields())
+        query = gql_compat(ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields(client))
 
         data = client.execute(query, variable_values={"id": artifact_id})
         result = ArtifactByID.model_validate(data)
 
-        if (art := result.artifact) is None:
+        if (artifact := result.artifact) is None:
             return None
 
-        src_collection = art.artifact_sequence
+        src_collection = artifact.artifact_sequence
         src_project = src_collection.project
 
         entity_name = src_project.entity_name if src_project else ""
         project_name = src_project.name if src_project else ""
 
-        name = f"{src_collection.name}:v{art.version_index}"
-        return cls._from_attrs(entity_name, project_name, name, art, client)
+        name = f"{src_collection.name}:v{artifact.version_index}"
+
+        path = FullArtifactPath(prefix=entity_name, project=project_name, name=name)
+        return cls._from_attrs(path, artifact, client)
 
     @classmethod
     def _membership_from_name(
-        cls,
-        *,
-        entity: str,
-        project: str,
-        name: str,
-        client: RetryingClient,
+        cls, *, path: FullArtifactPath, client: RetryingClient
     ) -> Artifact:
-        if not (api := InternalApi())._server_supports(
+        if not InternalApi()._server_supports(
             pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP
         ):
             raise UnsupportedError(
@@ -325,74 +323,70 @@ class Artifact:
 
         query = gql_compat(
             ARTIFACT_VIA_MEMBERSHIP_BY_NAME_GQL,
-            omit_fields=omit_artifact_fields(api=api),
+            omit_fields=omit_artifact_fields(client),
         )
-
-        gql_vars = {"entityName": entity, "projectName": project, "name": name}
+        gql_vars = {
+            "entityName": path.prefix,
+            "projectName": path.project,
+            "name": path.name,
+        }
         data = client.execute(query, variable_values=gql_vars)
         result = ArtifactViaMembershipByName.model_validate(data)
 
-        if not (project_attrs := result.project):
-            raise ValueError(f"project {project!r} not found under entity {entity!r}")
-
-        if not (acm_attrs := project_attrs.artifact_collection_membership):
-            entity_project = f"{entity}/{project}"
+        if not (project := result.project):
             raise ValueError(
-                f"artifact membership {name!r} not found in {entity_project!r}"
+                f"project {path.project!r} not found under entity {path.prefix!r}"
             )
-
-        target_path = ArtifactPath(prefix=entity, project=project, name=name)
-        return cls._from_membership(acm_attrs, target=target_path, client=client)
+        if not (membership := project.artifact_collection_membership):
+            entity_project = f"{path.prefix}/{path.project}"
+            raise ValueError(
+                f"artifact membership {path.name!r} not found in {entity_project!r}"
+            )
+        return cls._from_membership(membership, target=path, client=client)
 
     @classmethod
     def _from_name(
         cls,
         *,
-        entity: str,
-        project: str,
-        name: str,
+        path: FullArtifactPath,
         client: RetryingClient,
         enable_tracking: bool = False,
     ) -> Artifact:
-        if (api := InternalApi())._server_supports(
+        if InternalApi()._server_supports(
             pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP
         ):
-            return cls._membership_from_name(
-                entity=entity, project=project, name=name, client=client
-            )
-
-        supports_enable_tracking_gql_var = api.server_project_type_introspection()
-        omit_vars = None if supports_enable_tracking_gql_var else {"enableTracking"}
+            return cls._membership_from_name(path=path, client=client)
 
+        omit_vars = None if supports_enable_tracking_var(client) else {"enableTracking"}
         gql_vars = {
-            "entityName": entity,
-            "projectName": project,
-            "name": name,
+            "entityName": path.prefix,
+            "projectName": path.project,
+            "name": path.name,
             "enableTracking": enable_tracking,
         }
         query = gql_compat(
             ARTIFACT_BY_NAME_GQL,
             omit_variables=omit_vars,
-            omit_fields=omit_artifact_fields(api=api),
+            omit_fields=omit_artifact_fields(client),
         )
-
         data = client.execute(query, variable_values=gql_vars)
         result = ArtifactByName.model_validate(data)
 
-        if not (proj_attrs := result.project):
-            raise ValueError(f"project {project!r} not found under entity {entity!r}")
-
-        if not (art_attrs := proj_attrs.artifact):
-            entity_project = f"{entity}/{project}"
-            raise ValueError(f"artifact {name!r} not found in {entity_project!r}")
+        if not (project := result.project):
+            raise ValueError(
+                f"project {path.project!r} not found under entity {path.prefix!r}"
+            )
+        if not (artifact := project.artifact):
+            entity_project = f"{path.prefix}/{path.project}"
+            raise ValueError(f"artifact {path.name!r} not found in {entity_project!r}")
 
-        return cls._from_attrs(entity, project, name, art_attrs, client)
+        return cls._from_attrs(path, artifact, client)
 
     @classmethod
     def _from_membership(
         cls,
         membership: MembershipWithArtifact,
-        target: ArtifactPath,
+        target: FullArtifactPath,
         client: RetryingClient,
     ) -> Artifact:
         if not (
@@ -410,35 +404,31 @@ class Artifact:
                 f"Your request was redirected to the corresponding artifact {name!r} in the new registry. "
                 f"Please update your paths to point to the migrated registry directly, '{proj.name}/{name}'."
             )
-            new_entity, new_project = proj.entity_name, proj.name
+            new_target = replace(target, prefix=proj.entity_name, project=proj.name)
         else:
-            new_entity = cast(str, target.prefix)
-            new_project = cast(str, target.project)
+            new_target = copy(target)
 
         if not (artifact := membership.artifact):
             raise ValueError(f"Artifact {target.to_str()!r} not found in response")
 
-        return cls._from_attrs(new_entity, new_project, target.name, artifact, client)
+        return cls._from_attrs(new_target, artifact, client)
 
     @classmethod
     def _from_attrs(
         cls,
-        entity: str,
-        project: str,
-        name: str,
-        attrs: dict[str, Any] | ArtifactFragment,
+        path: FullArtifactPath,
+        attrs: ArtifactFragment,
         client: RetryingClient,
         aliases: list[str] | None = None,
     ) -> Artifact:
         # Placeholder is required to skip validation.
         artifact = cls("placeholder", type="placeholder")
         artifact._client = client
-        artifact._entity = entity
-        artifact._project = project
-        artifact._name = name
+        artifact._entity = path.prefix
+        artifact._project = path.project
+        artifact._name = path.name
 
-        validated_attrs = ArtifactFragment.model_validate(attrs)
-        artifact._assign_attrs(validated_attrs, aliases)
+        artifact._assign_attrs(attrs, aliases)
 
         artifact.finalize()
 
@@ -743,13 +733,12 @@ class Artifact:
                 raise ValueError("Client is not initialized")
 
             try:
-                artifact = self._from_name(
-                    entity=self.source_entity,
+                path = FullArtifactPath(
+                    prefix=self.source_entity,
                     project=self.source_project,
                     name=self.source_name,
-                    client=self._client,
                 )
-                self._source_artifact = artifact
+                self._source_artifact = self._from_name(path=path, client=self._client)
             except Exception as e:
                 raise ValueError(
                     f"Unable to fetch source artifact for linked artifact {self.name}"
@@ -1234,8 +1223,9 @@ class Artifact:
     def _populate_after_save(self, artifact_id: str) -> None:
         assert self._client is not None
 
-        query = gql_compat(ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields())
-
+        query = gql_compat(
+            ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields(self._client)
+        )
         data = self._client.execute(query, variable_values={"id": artifact_id})
         result = ArtifactByID.model_validate(data)
 
@@ -1258,21 +1248,9 @@ class Artifact:
         collection = self.name.split(":")[0]
 
         aliases = None
-        introspect_query = gql(
-            """
-            query ProbeServerAddAliasesInput {
-               AddAliasesInputInfoType: __type(name: "AddAliasesInput") {
-                   name
-                   inputFields {
-                       name
-                   }
-                }
-            }
-            """
-        )
 
-        data = self._client.execute(introspect_query)
-        if data.get("AddAliasesInputInfoType"):  # wandb backend version >= 0.13.0
+        if type_info(self._client, "AddAliasesInput") is not None:
+            # wandb backend version >= 0.13.0
             alias_props = {
                 "entity_name": entity,
                 "project_name": project,
@@ -1330,7 +1308,7 @@ class Artifact:
                 for alias in self.aliases
             ]
 
-        omit_fields = omit_artifact_fields()
+        omit_fields = omit_artifact_fields(self._client)
         omit_variables = set()
 
         if {"ttlIsInherited", "ttlDurationSeconds"} & omit_fields:
@@ -2068,24 +2046,14 @@ class Artifact:
 
         download_logger = ArtifactDownloadLogger(nfiles=nfiles)
 
-        def _download_entry(
-            entry: ArtifactManifestEntry,
-            executor: concurrent.futures.Executor,
-            api_key: str | None,
-            cookies: dict | None,
-            headers: dict | None,
-        ) -> None:
-            _thread_local_api_settings.api_key = api_key
-            _thread_local_api_settings.cookies = cookies
-            _thread_local_api_settings.headers = headers
-
+        def _download_entry(entry: ArtifactManifestEntry, executor: Executor) -> None:
+            multipart_executor = (
+                executor
+                if should_multipart_download(entry.size, override=multipart)
+                else None
+            )
             try:
-                entry.download(
-                    root,
-                    skip_cache=skip_cache,
-                    executor=executor,
-                    multipart=multipart,
-                )
+                entry.download(root, skip_cache=skip_cache, executor=multipart_executor)
             except FileNotFoundError as e:
                 if allow_missing_references:
                     wandb.termwarn(str(e))
@@ -2096,15 +2064,23 @@ class Artifact:
                 return
             download_logger.notify_downloaded()
 
-        with concurrent.futures.ThreadPoolExecutor(64) as executor:
-            download_entry = partial(
-                _download_entry,
-                executor=executor,
-                api_key=_thread_local_api_settings.api_key,
-                cookies=_thread_local_api_settings.cookies,
-                headers=_thread_local_api_settings.headers,
-            )
+        def _init_thread(
+            api_key: str | None, cookies: dict | None, headers: dict | None
+        ) -> None:
+            """Initialize the thread-local API settings in the CURRENT thread."""
+            _thread_local_api_settings.api_key = api_key
+            _thread_local_api_settings.cookies = cookies
+            _thread_local_api_settings.headers = headers
 
+        with ThreadPoolExecutor(
+            max_workers=64,
+            initializer=_init_thread,
+            initargs=(
+                _thread_local_api_settings.api_key,
+                _thread_local_api_settings.cookies,
+                _thread_local_api_settings.headers,
+            ),
+        ) as executor:
             batch_size = env.get_artifact_fetch_file_url_batch_size()
 
             active_futures = set()
@@ -2125,7 +2101,9 @@ class Artifact:
                     #     continue
                     entry._download_url = node.direct_url
                     if (not path_prefix) or entry.path.startswith(str(path_prefix)):
-                        active_futures.add(executor.submit(download_entry, entry))
+                        active_futures.add(
+                            executor.submit(_download_entry, entry, executor=executor)
+                        )
 
                 # Wait for download threads to catch up.
                 #
@@ -2140,14 +2118,14 @@ class Artifact:
                 #   Consider this for a future change, or (depending on risk and risk tolerance)
                 #   managing this logic via asyncio instead, if viable.
                 if len(active_futures) > batch_size:
-                    for future in concurrent.futures.as_completed(active_futures):
+                    for future in as_completed(active_futures):
                         future.result()  # check for errors
                         active_futures.remove(future)
                         if len(active_futures) <= batch_size:
                             break
 
             # Check for errors.
-            for future in concurrent.futures.as_completed(active_futures):
+            for future in as_completed(active_futures):
                 future.result()
 
         if log:
@@ -2427,7 +2405,7 @@ class Artifact:
 
         # Parse the entity (first part of the path) appropriately,
         # depending on whether we're linking to a registry
-        if target.project and is_artifact_registry_project(target.project):
+        if target.is_registry_path():
             # In a Registry linking, the entity is used to fetch the organization of the artifact
             # therefore the source artifact's entity is passed to the backend
             org = target.prefix or settings.get("organization") or ""
@@ -2435,6 +2413,9 @@ class Artifact:
         else:
             target = target.with_defaults(prefix=self.source_entity)
 
+        # Explicitly convert to FullArtifactPath to ensure all fields are present
+        target = FullArtifactPath(**asdict(target))
+
         # Prepare the validated GQL input, send it
         alias_inputs = [
             ArtifactAliasInput(artifact_collection_name=target.name, alias=a)

wandb/sdk/artifacts/artifact_file_cache.py

@@ -26,6 +26,11 @@ class Opener(Protocol):
         pass
 
 
+def artifacts_cache_dir() -> Path:
+    """Get the artifacts cache directory."""
+    return env.get_cache_dir() / "artifacts"
+
+
 def _get_sys_umask_threadsafe() -> int:
     # Workaround to get the current system umask, since
     # - `os.umask()` isn't thread-safe
@@ -248,4 +253,4 @@ def _build_artifact_file_cache(cache_dir: StrPath) -> ArtifactFileCache:
 
 
 def get_artifact_file_cache() -> ArtifactFileCache:
-    return _build_artifact_file_cache(env.get_cache_dir() / "artifacts")
+    return _build_artifact_file_cache(artifacts_cache_dir())

wandb/sdk/artifacts/artifact_manifest_entry.py

@@ -3,9 +3,11 @@
 from __future__ import annotations
 
 import concurrent.futures
+import hashlib
 import json
 import logging
 import os
+from contextlib import suppress
 from pathlib import Path
 from typing import TYPE_CHECKING
 from urllib.parse import urlparse
@@ -43,6 +45,48 @@ if TYPE_CHECKING:
 _WB_ARTIFACT_SCHEME = "wandb-artifact"
 
 
+def _checksum_cache_path(file_path: str) -> str:
+    """Get path for checksum in central cache directory."""
+    from wandb.sdk.artifacts.artifact_file_cache import artifacts_cache_dir
+
+    # Create a unique cache key based on the file's absolute path
+    abs_path = os.path.abspath(file_path)
+    path_hash = hashlib.sha256(abs_path.encode()).hexdigest()
+
+    # Store in wandb cache directory under checksums subdirectory
+    cache_dir = artifacts_cache_dir() / "checksums"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+
+    return str(cache_dir / f"{path_hash}.checksum")
+
+
+def _read_cached_checksum(file_path: str) -> str | None:
+    """Read checksum from cache if it exists and is valid."""
+    checksum_path = _checksum_cache_path(file_path)
+
+    try:
+        with open(file_path) as f, open(checksum_path) as f_checksum:
+            if os.path.getmtime(f_checksum.name) < os.path.getmtime(f.name):
+                # File was modified after checksum was written
+                return None
+            # Read and return the cached checksum
+            return f_checksum.read().strip()
+    except OSError:
+        # File doesn't exist or couldn't be opened
+        return None
+
+
+def _write_cached_checksum(file_path: str, checksum: str) -> None:
+    """Write checksum to cache directory."""
+    checksum_path = _checksum_cache_path(file_path)
+    try:
+        with open(checksum_path, "w") as f:
+            f.write(checksum)
+    except OSError:
+        # Non-critical failure, just log it
+        logger.debug(f"Failed to write checksum cache for {file_path!r}")
+
+
 class ArtifactManifestEntry:
     """A single entry in an artifact manifest."""
 
@@ -138,7 +182,6 @@ class ArtifactManifestEntry:
         root: str | None = None,
         skip_cache: bool | None = None,
         executor: concurrent.futures.Executor | None = None,
-        multipart: bool | None = None,
     ) -> FilePathStr:
         """Download this artifact entry to the specified root path.
 
@@ -149,11 +192,10 @@ class ArtifactManifestEntry:
         Returns:
             (str): The path of the downloaded artifact entry.
         """
-        if self._parent_artifact is None:
-            raise NotImplementedError
+        artifact = self.parent_artifact()
 
-        root = root or self._parent_artifact._default_root()
-        self._parent_artifact._add_download_root(root)
+        root = root or artifact._default_root()
+        artifact._add_download_root(root)
         path = str(Path(self.path))
         dest_path = os.path.join(root, path)
 
@@ -164,32 +206,43 @@ class ArtifactManifestEntry:
 
         # Skip checking the cache (and possibly downloading) if the file already exists
         # and has the digest we're expecting.
+
+        # Fast integrity check using cached checksum from persistent cache
+        with suppress(OSError):
+            if self.digest == _read_cached_checksum(dest_path):
+                return FilePathStr(dest_path)
+
+        # Fallback to computing/caching the checksum hash
         try:
             md5_hash = md5_file_b64(dest_path)
         except (FileNotFoundError, IsADirectoryError):
-            logger.debug(f"unable to find {dest_path}, skip searching for file")
+            logger.debug(f"unable to find {dest_path!r}, skip searching for file")
         else:
+            _write_cached_checksum(dest_path, md5_hash)
             if self.digest == md5_hash:
                 return FilePathStr(dest_path)
 
         if self.ref is not None:
-            cache_path = self._parent_artifact.manifest.storage_policy.load_reference(
+            cache_path = artifact.manifest.storage_policy.load_reference(
                 self, local=True, dest_path=override_cache_path
             )
         else:
-            cache_path = self._parent_artifact.manifest.storage_policy.load_file(
-                self._parent_artifact,
-                self,
-                dest_path=override_cache_path,
-                executor=executor,
-                multipart=multipart,
+            cache_path = artifact.manifest.storage_policy.load_file(
+                artifact, self, dest_path=override_cache_path, executor=executor
             )
-        return FilePathStr(
+
+        # Determine the final path
+        final_path = (
             dest_path
             if skip_cache
             else copy_or_overwrite_changed(cache_path, dest_path)
         )
 
+        # Cache the checksum for future downloads
+        _write_cached_checksum(str(final_path), self.digest)
+
+        return FilePathStr(final_path)
+
     def ref_target(self) -> FilePathStr | URIStr:
         """Get the reference URL that is targeted by this artifact entry.
 

wandb/sdk/artifacts/storage_handler.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Final
 
 from wandb.sdk.lib.paths import FilePathStr, URIStr
 
@@ -12,18 +13,11 @@ if TYPE_CHECKING:
     from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry
 
-DEFAULT_MAX_OBJECTS = 10**7
+DEFAULT_MAX_OBJECTS: Final[int] = 10_000_000  # 10**7
 
 
-class StorageHandler:
-    def can_handle(self, parsed_url: ParseResult) -> bool:
-        """Checks whether this handler can handle the given url.
-
-        Returns:
-            Whether this handler can handle the given url.
-        """
-        raise NotImplementedError
-
+class _BaseStorageHandler(ABC):
+    @abstractmethod
     def load_path(
         self,
         manifest_entry: ArtifactManifestEntry,
@@ -40,6 +34,7 @@ class StorageHandler:
         """
         raise NotImplementedError
 
+    @abstractmethod
     def store_path(
         self,
         artifact: Artifact,
@@ -47,7 +42,7 @@ class StorageHandler:
         name: str | None = None,
         checksum: bool = True,
         max_objects: int | None = None,
-    ) -> Sequence[ArtifactManifestEntry]:
+    ) -> list[ArtifactManifestEntry]:
         """Store the file or directory at the given path to the specified artifact.
 
         Args:
@@ -60,3 +55,14 @@ class StorageHandler:
             A list of manifest entries to store within the artifact
         """
         raise NotImplementedError
+
+
+class StorageHandler(_BaseStorageHandler, ABC):  # Handles a single storage protocol
+    @abstractmethod
+    def can_handle(self, parsed_url: ParseResult) -> bool:
+        """Checks whether this handler can handle the given url.
+
+        Returns:
+            Whether this handler can handle the given url.
+        """
+        raise NotImplementedError