wandb 0.21.4__py3-none-musllinux_1_2_aarch64.whl → 0.22.1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/artifacts/artifact.py

@@ -16,7 +16,7 @@ import tempfile
 import time
 from collections import deque
 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
@@ -30,7 +30,6 @@ from typing import (
     Literal,
     Sequence,
     Type,
-    cast,
     final,
 )
 from urllib.parse import quote, urljoin, urlparse
@@ -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:
@@ -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."""
 
@@ -164,11 +208,19 @@ 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)
 
@@ -184,12 +236,19 @@ class ArtifactManifestEntry:
                 executor=executor,
                 multipart=multipart,
             )
-        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_handlers/gcs_handler.py

@@ -65,7 +65,7 @@ class GCSHandler(StorageHandler):
         path, hit, cache_open = self._cache.check_etag_obj_path(
             url=URIStr(manifest_entry.ref),
             etag=ETag(manifest_entry.digest),
-            size=manifest_entry.size if manifest_entry.size is not None else 0,
+            size=manifest_entry.size or 0,
         )
         if hit:
             return path

wandb/sdk/artifacts/storage_handlers/http_handler.py

@@ -43,7 +43,7 @@ class HTTPHandler(StorageHandler):
         path, hit, cache_open = self._cache.check_etag_obj_path(
             URIStr(manifest_entry.ref),
             ETag(manifest_entry.digest),
-            manifest_entry.size if manifest_entry.size is not None else 0,
+            manifest_entry.size or 0,
         )
         if hit:
             return path
@@ -54,7 +54,6 @@ class HTTPHandler(StorageHandler):
             cookies=_thread_local_api_settings.cookies,
             headers=_thread_local_api_settings.headers,
         )
-        response.raise_for_status()
 
         digest: ETag | FilePathStr | URIStr | None
         digest, size, extra = self._entry_from_headers(response.headers)
@@ -87,7 +86,6 @@ class HTTPHandler(StorageHandler):
             cookies=_thread_local_api_settings.cookies,
             headers=_thread_local_api_settings.headers,
         ) as response:
-            response.raise_for_status()
             digest: ETag | FilePathStr | URIStr | None
             digest, size, extra = self._entry_from_headers(response.headers)
             digest = digest or path

wandb/sdk/artifacts/storage_handlers/local_file_handler.py

@@ -51,7 +51,7 @@ class LocalFileHandler(StorageHandler):
 
         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,
+            manifest_entry.size or 0,
         )
         if hit:
             return path

wandb/sdk/artifacts/storage_handlers/s3_handler.py

@@ -96,7 +96,7 @@ class S3Handler(StorageHandler):
         path, hit, cache_open = self._cache.check_etag_obj_path(
             URIStr(manifest_entry.ref),
             ETag(manifest_entry.digest),
-            manifest_entry.size if manifest_entry.size is not None else 0,
+            manifest_entry.size or 0,
         )
         if hit:
             return path

wandb/sdk/artifacts/storage_policies/_factories.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from typing import Final
+
+from requests import Response, Session
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from ..storage_handler import StorageHandler
+from ..storage_handlers.azure_handler import AzureHandler
+from ..storage_handlers.gcs_handler import GCSHandler
+from ..storage_handlers.http_handler import HTTPHandler
+from ..storage_handlers.local_file_handler import LocalFileHandler
+from ..storage_handlers.s3_handler import S3Handler
+from ..storage_handlers.wb_artifact_handler import WBArtifactHandler
+from ..storage_handlers.wb_local_artifact_handler import WBLocalArtifactHandler
+
+# 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.
+HTTP_RETRY_STRATEGY: Final[Retry] = Retry(
+    backoff_factor=1,
+    total=16,
+    status_forcelist=(308, 408, 409, 429, 500, 502, 503, 504),
+)
+HTTP_POOL_CONNECTIONS: Final[int] = 64
+HTTP_POOL_MAXSIZE: Final[int] = 64
+
+
+def raise_for_status(response: Response, *_, **__) -> None:
+    """A `requests.Session` hook to raise for status on all requests."""
+    response.raise_for_status()
+
+
+def make_http_session() -> Session:
+    """A factory that returns a `requests.Session` for use with artifact storage handlers."""
+    session = Session()
+
+    # Explicitly configure the retry strategy for http/https adapters.
+    adapter = HTTPAdapter(
+        max_retries=HTTP_RETRY_STRATEGY,
+        pool_connections=HTTP_POOL_CONNECTIONS,
+        pool_maxsize=HTTP_POOL_MAXSIZE,
+    )
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    # Always raise on HTTP status errors.
+    session.hooks["response"].append(raise_for_status)
+    return session
+
+
+def make_storage_handlers(session: Session) -> list[StorageHandler]:
+    """A factory that returns the default artifact storage handlers."""
+    return [
+        S3Handler(),  # s3
+        GCSHandler(),  # gcs
+        AzureHandler(),  # azure
+        HTTPHandler(session, scheme="http"),  # http
+        HTTPHandler(session, scheme="https"),  # https
+        WBArtifactHandler(),  # artifact
+        WBLocalArtifactHandler(),  # local_artifact
+        LocalFileHandler(),  # file_handler
+    ]