wandb 0.21.1__py3-none-win_amd64.whl → 0.21.2__py3-none-win_amd64.whl

wandb/sdk/artifacts/artifact.py

@@ -17,11 +17,22 @@ import time
 from collections import deque
 from copy import copy
 from dataclasses import dataclass
-from datetime import datetime, timedelta
+from datetime import timedelta
 from functools import partial
 from itertools import filterfalse
-from pathlib import PurePosixPath
-from typing import IO, TYPE_CHECKING, Any, Iterator, Literal, Sequence, Type, final
+from pathlib import Path, PurePosixPath
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Iterator,
+    Literal,
+    Sequence,
+    Type,
+    cast,
+    final,
+)
 from urllib.parse import quote, urljoin, urlparse
 
 import requests
@@ -58,21 +69,45 @@ from wandb.util import (
     vendor_setup,
 )
 
+from ._factories import make_storage_policy
 from ._generated import (
     ADD_ALIASES_GQL,
+    ARTIFACT_BY_ID_GQL,
+    ARTIFACT_BY_NAME_GQL,
+    ARTIFACT_COLLECTION_MEMBERSHIP_FILE_URLS_GQL,
+    ARTIFACT_CREATED_BY_GQL,
+    ARTIFACT_FILE_URLS_GQL,
+    ARTIFACT_TYPE_GQL,
+    ARTIFACT_USED_BY_GQL,
+    ARTIFACT_VIA_MEMBERSHIP_BY_NAME_GQL,
     DELETE_ALIASES_GQL,
+    DELETE_ARTIFACT_GQL,
+    FETCH_ARTIFACT_MANIFEST_GQL,
     FETCH_LINKED_ARTIFACTS_GQL,
     LINK_ARTIFACT_GQL,
+    UNLINK_ARTIFACT_GQL,
     UPDATE_ARTIFACT_GQL,
     ArtifactAliasInput,
+    ArtifactByID,
+    ArtifactByName,
     ArtifactCollectionAliasInput,
+    ArtifactCollectionMembershipFileUrls,
+    ArtifactCreatedBy,
+    ArtifactFileUrls,
+    ArtifactFragment,
+    ArtifactType,
+    ArtifactUsedBy,
+    ArtifactViaMembershipByName,
+    FetchArtifactManifest,
     FetchLinkedArtifacts,
+    FileUrlsFragment,
     LinkArtifact,
     LinkArtifactInput,
+    MembershipWithArtifact,
     TagInput,
     UpdateArtifact,
 )
-from ._graphql_fragments import _gql_artifact_fragment, omit_artifact_fields
+from ._graphql_fragments import omit_artifact_fields
 from ._validators import (
     LINKED_ARTIFACT_COLLECTION_TYPE,
     ArtifactPath,
@@ -103,9 +138,6 @@ from .exceptions import (
 )
 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 = vendor_setup()
 
@@ -119,6 +151,9 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+_MB: Final[int] = 1024 * 1024
+
+
 @final
 @dataclass
 class _DeferredArtifactManifest:
@@ -188,11 +223,6 @@ class Artifact:
         # Internal.
         self._client: RetryingClient | None = None
 
-        storage_policy_cls = StoragePolicy.lookup_by_name(WANDB_STORAGE_POLICY)
-        layout = StorageLayout.V1 if env.get_use_v1_artifacts() else StorageLayout.V2
-        policy_config = {"storageLayout": layout}
-        self._storage_policy = storage_policy_cls.from_config(config=policy_config)
-
         self._tmp_dir: tempfile.TemporaryDirectory | None = None
         self._added_objs: dict[int, tuple[WBValue, ArtifactManifestEntry]] = {}
         self._added_local_paths: dict[str, ArtifactManifestEntry] = {}
@@ -236,7 +266,7 @@ class Artifact:
         self._use_as: str | None = None
         self._state: ArtifactState = ArtifactState.PENDING
         self._manifest: ArtifactManifest | _DeferredArtifactManifest | None = (
-            ArtifactManifestV1(self._storage_policy)
+            ArtifactManifestV1(storage_policy=make_storage_policy())
         )
         self._commit_hash: str | None = None
         self._file_count: int | None = None
@@ -257,32 +287,22 @@ class Artifact:
         if (artifact := artifact_instance_cache.get(artifact_id)) is not None:
             return artifact
 
-        query = gql(
-            """
-            query ArtifactByID($id: ID!) {
-                artifact(id: $id) {
-                    ...ArtifactFragment
-                }
-            }
-            """
-            + _gql_artifact_fragment()
-        )
-        response = client.execute(
-            query,
-            variable_values={"id": artifact_id},
-        )
-        attrs = response.get("artifact")
-        if attrs is None:
+        query = gql_compat(ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields())
+
+        data = client.execute(query, variable_values={"id": artifact_id})
+        result = ArtifactByID.model_validate(data)
+
+        if (art := result.artifact) is None:
             return None
 
-        src_collection = attrs["artifactSequence"]
-        src_project = src_collection["project"]
+        src_collection = art.artifact_sequence
+        src_project = src_collection.project
 
-        entity_name = src_project["entityName"] if src_project else ""
-        project_name = src_project["name"] if src_project else ""
+        entity_name = src_project.entity_name if src_project else ""
+        project_name = src_project.name if src_project else ""
 
-        name = "{}:v{}".format(src_collection["name"], attrs["versionIndex"])
-        return cls._from_attrs(entity_name, project_name, name, attrs, client)
+        name = f"{src_collection.name}:v{art.version_index}"
+        return cls._from_attrs(entity_name, project_name, name, art, client)
 
     @classmethod
     def _membership_from_name(
@@ -293,7 +313,7 @@ class Artifact:
         name: str,
         client: RetryingClient,
     ) -> Artifact:
-        if not InternalApi()._server_supports(
+        if not (api := InternalApi())._server_supports(
             pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP
         ):
             raise UnsupportedError(
@@ -301,69 +321,26 @@ class Artifact:
                 "by this version of wandb server. Consider updating to the latest version."
             )
 
-        query = gql(
-            f"""
-            query ArtifactByName($entityName: String!, $projectName: String!, $name: String!) {{
-                project(name: $projectName, entityName: $entityName) {{
-                    artifactCollectionMembership(name: $name) {{
-                        id
-                        artifactCollection {{
-                            id
-                            name
-                            project {{
-                                id
-                                entityName
-                                name
-                            }}
-                        }}
-                        artifact {{
-                            ...ArtifactFragment
-                        }}
-                    }}
-                }}
-            }}
-            {_gql_artifact_fragment()}
-            """
+        query = gql_compat(
+            ARTIFACT_VIA_MEMBERSHIP_BY_NAME_GQL,
+            omit_fields=omit_artifact_fields(api=api),
         )
 
-        query_variable_values: dict[str, Any] = {
-            "entityName": entity,
-            "projectName": project,
-            "name": name,
-        }
-        response = client.execute(
-            query,
-            variable_values=query_variable_values,
-        )
-        if not (project_attrs := response.get("project")):
+        gql_vars = {"entityName": entity, "projectName": project, "name": 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.get("artifactCollectionMembership")):
+
+        if not (acm_attrs := project_attrs.artifact_collection_membership):
             entity_project = f"{entity}/{project}"
             raise ValueError(
                 f"artifact membership {name!r} not found in {entity_project!r}"
             )
-        if not (ac_attrs := acm_attrs.get("artifactCollection")):
-            raise ValueError("artifact collection not found")
-        if not (
-            (ac_name := ac_attrs.get("name"))
-            and (ac_project_attrs := ac_attrs.get("project"))
-        ):
-            raise ValueError("artifact collection project not found")
-        ac_project = ac_project_attrs.get("name")
-        ac_entity = ac_project_attrs.get("entityName")
-        if is_artifact_registry_project(ac_project) and project == "model-registry":
-            wandb.termwarn(
-                "This model registry has been migrated and will be discontinued. "
-                f"Your request was redirected to the corresponding artifact `{ac_name}` in the new registry. "
-                f"Please update your paths to point to the migrated registry directly, `{ac_project}/{ac_name}`."
-            )
-            entity = ac_entity
-            project = ac_project
-        if not (attrs := acm_attrs.get("artifact")):
-            entity_project = f"{entity}/{project}"
-            raise ValueError(f"artifact {name!r} not found in {entity_project!r}")
 
-        return cls._from_attrs(entity, project, name, attrs, client)
+        target_path = ArtifactPath(prefix=entity, project=project, name=name)
+        return cls._from_membership(acm_attrs, target=target_path, client=client)
 
     @classmethod
     def _from_name(
@@ -375,59 +352,71 @@ class Artifact:
         client: RetryingClient,
         enable_tracking: bool = False,
     ) -> Artifact:
-        if InternalApi()._server_supports(
+        if (api := InternalApi())._server_supports(
             pb.ServerFeature.PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP
         ):
             return cls._membership_from_name(
-                entity=entity,
-                project=project,
-                name=name,
-                client=client,
+                entity=entity, project=project, name=name, client=client
             )
 
-        query_variable_values: dict[str, Any] = {
+        supports_enable_tracking_gql_var = api.server_project_type_introspection()
+        omit_vars = None if supports_enable_tracking_gql_var else {"enableTracking"}
+
+        gql_vars = {
             "entityName": entity,
             "projectName": project,
             "name": name,
+            "enableTracking": enable_tracking,
         }
-        query_vars = ["$entityName: String!", "$projectName: String!", "$name: String!"]
-        query_args = ["name: $name"]
-
-        server_supports_enabling_artifact_usage_tracking = (
-            InternalApi().server_project_type_introspection()
+        query = gql_compat(
+            ARTIFACT_BY_NAME_GQL,
+            omit_variables=omit_vars,
+            omit_fields=omit_artifact_fields(api=api),
         )
-        if server_supports_enabling_artifact_usage_tracking:
-            query_vars.append("$enableTracking: Boolean")
-            query_args.append("enableTracking: $enableTracking")
-            query_variable_values["enableTracking"] = enable_tracking
-
-        vars_str = ", ".join(query_vars)
-        args_str = ", ".join(query_args)
-
-        query = gql(
-            f"""
-            query ArtifactByName({vars_str}) {{
-                project(name: $projectName, entityName: $entityName) {{
-                    artifact({args_str}) {{
-                        ...ArtifactFragment
-                    }}
-                }}
-            }}
-            {_gql_artifact_fragment()}
-            """
-        )
-        response = client.execute(
-            query,
-            variable_values=query_variable_values,
-        )
-        project_attrs = response.get("project")
-        if not project_attrs:
-            raise ValueError(f"project '{project}' not found under entity '{entity}'")
-        attrs = project_attrs.get("artifact")
-        if not attrs:
-            raise ValueError(f"artifact '{name}' not found in '{entity}/{project}'")
 
-        return cls._from_attrs(entity, project, name, attrs, 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}")
+
+        return cls._from_attrs(entity, project, name, art_attrs, client)
+
+    @classmethod
+    def _from_membership(
+        cls,
+        membership: MembershipWithArtifact,
+        target: ArtifactPath,
+        client: RetryingClient,
+    ) -> Artifact:
+        if not (
+            (collection := membership.artifact_collection)
+            and (name := collection.name)
+            and (proj := collection.project)
+        ):
+            raise ValueError("Missing artifact collection project in GraphQL response")
+
+        if is_artifact_registry_project(proj.name) and (
+            target.project == "model-registry"
+        ):
+            wandb.termwarn(
+                "This model registry has been migrated and will be discontinued. "
+                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
+        else:
+            new_entity = cast(str, target.prefix)
+            new_project = cast(str, target.project)
+
+        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)
 
     @classmethod
     def _from_attrs(
@@ -435,7 +424,7 @@ class Artifact:
         entity: str,
         project: str,
         name: str,
-        attrs: dict[str, Any],
+        attrs: dict[str, Any] | ArtifactFragment,
         client: RetryingClient,
         aliases: list[str] | None = None,
     ) -> Artifact:
@@ -445,7 +434,9 @@ class Artifact:
         artifact._entity = entity
         artifact._project = project
         artifact._name = name
-        artifact._assign_attrs(attrs, aliases)
+
+        validated_attrs = ArtifactFragment.model_validate(attrs)
+        artifact._assign_attrs(validated_attrs, aliases)
 
         artifact.finalize()
 
@@ -458,29 +449,24 @@ class Artifact:
     # doesn't make it clear if the artifact is a link or not and have to manually set it.
     def _assign_attrs(
         self,
-        attrs: dict[str, Any],
+        art: ArtifactFragment,
         aliases: list[str] | None = None,
         is_link: bool | None = None,
     ) -> None:
         """Update this Artifact's attributes using the server response."""
-        self._id = attrs["id"]
-
-        src_version = f"v{attrs['versionIndex']}"
-        src_collection = attrs["artifactSequence"]
-        src_project = src_collection["project"]
+        self._id = art.id
 
-        self._source_entity = src_project["entityName"] if src_project else ""
-        self._source_project = src_project["name"] if src_project else ""
-        self._source_name = f"{src_collection['name']}:{src_version}"
-        self._source_version = src_version
+        src_collection = art.artifact_sequence
+        src_project = src_collection.project
 
-        if self._entity is None:
-            self._entity = self._source_entity
-        if self._project is None:
-            self._project = self._source_project
+        self._source_entity = src_project.entity_name if src_project else ""
+        self._source_project = src_project.name if src_project else ""
+        self._source_name = f"{src_collection.name}:v{art.version_index}"
+        self._source_version = f"v{art.version_index}"
 
-        if self._name is None:
-            self._name = self._source_name
+        self._entity = self._entity or self._source_entity
+        self._project = self._project or self._source_project
+        self._name = self._name or self._source_name
 
         # TODO: Refactor artifact query to fetch artifact via membership instead
         # and get the collection type
@@ -488,33 +474,35 @@ class Artifact:
             self._is_link = (
                 self._entity != self._source_entity
                 or self._project != self._source_project
-                or self._name != self._source_name
+                or self._name.split(":")[0] != self._source_name.split(":")[0]
             )
         else:
             self._is_link = is_link
 
-        self._type = attrs["artifactType"]["name"]
-        self._description = attrs["description"]
-
-        entity = self._entity
-        project = self._project
-        collection, *_ = self._name.split(":")
+        self._type = art.artifact_type.name
+        self._description = art.description
 
-        processed_aliases = []
         # The future of aliases is to move all alias fetches to the membership level
         # so we don't have to do the collection fetches below
         if aliases:
             processed_aliases = aliases
-        else:
+        elif art.aliases:
+            entity = self._entity
+            project = self._project
+            collection = self._name.split(":")[0]
             processed_aliases = [
-                obj["alias"]
-                for obj in attrs["aliases"]
-                if obj["artifactCollection"]
-                and obj["artifactCollection"]["project"]
-                and obj["artifactCollection"]["project"]["entityName"] == entity
-                and obj["artifactCollection"]["project"]["name"] == project
-                and obj["artifactCollection"]["name"] == collection
+                art_alias.alias
+                for art_alias in art.aliases
+                if (
+                    (coll := art_alias.artifact_collection)
+                    and (proj := coll.project)
+                    and proj.entity_name == entity
+                    and proj.name == project
+                    and coll.name == collection
+                )
             ]
+        else:
+            processed_aliases = []
 
         version_aliases = list(filter(alias_is_version_index, processed_aliases))
         other_aliases = list(filterfalse(alias_is_version_index, processed_aliases))
@@ -524,49 +512,42 @@ class Artifact:
                 version_aliases, too_short=TooFewItemsError, too_long=TooManyItemsError
             )
         except TooFewItemsError:
-            version = src_version  # default to the source version
+            version = f"v{art.version_index}"  # 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
-
-        if ":" not in self._name:
-            self._name = f"{self._name}:{version}"
+        self._name = self._name if (":" in self._name) else f"{self._name}:{version}"
 
         self._aliases = other_aliases
-        self._saved_aliases = copy(other_aliases)
+        self._saved_aliases = copy(self._aliases)
 
-        tags = [obj["name"] for obj in (attrs.get("tags") or [])]
-        self._tags = tags
-        self._saved_tags = copy(tags)
+        self._tags = [tag.name for tag in (art.tags or [])]
+        self._saved_tags = copy(self._tags)
 
-        metadata_str = attrs["metadata"]
-        self._metadata = validate_metadata(
-            json.loads(metadata_str) if metadata_str else {}
-        )
+        self._metadata = validate_metadata(art.metadata)
 
         self._ttl_duration_seconds = validate_ttl_duration_seconds(
-            attrs.get("ttlDurationSeconds")
+            art.ttl_duration_seconds
         )
         self._ttl_is_inherited = (
-            True if (attrs.get("ttlIsInherited") is None) else attrs["ttlIsInherited"]
+            True if (art.ttl_is_inherited is None) else art.ttl_is_inherited
         )
 
-        self._state = ArtifactState(attrs["state"])
+        self._state = ArtifactState(art.state)
 
-        try:
-            manifest_url = attrs["currentManifest"]["file"]["directUrl"]
-        except (LookupError, TypeError):
-            self._manifest = None
-        else:
-            self._manifest = _DeferredArtifactManifest(manifest_url)
+        self._manifest = (
+            _DeferredArtifactManifest(manifest.file.direct_url)
+            if (manifest := art.current_manifest)
+            else None
+        )
 
-        self._commit_hash = attrs["commitHash"]
-        self._file_count = attrs["fileCount"]
-        self._created_at = attrs["createdAt"]
-        self._updated_at = attrs["updatedAt"]
-        self._history_step = attrs.get("historyStep", None)
+        self._commit_hash = art.commit_hash
+        self._file_count = art.file_count
+        self._created_at = art.created_at
+        self._updated_at = art.updated_at
+        self._history_step = art.history_step
 
     @ensure_logged
     def new_draft(self) -> Artifact:
@@ -1063,37 +1044,24 @@ class Artifact:
             return self._manifest
 
         if self._manifest is None:
-            query = gql(
-                """
-                query ArtifactManifest(
-                    $entityName: String!,
-                    $projectName: String!,
-                    $name: String!
-                ) {
-                    project(entityName: $entityName, name: $projectName) {
-                        artifact(name: $name) {
-                            currentManifest {
-                                file {
-                                    directUrl
-                                }
-                            }
-                        }
-                    }
-                }
-                """
-            )
-            assert self._client is not None
-            response = self._client.execute(
-                query,
-                variable_values={
-                    "entityName": self._entity,
-                    "projectName": self._project,
-                    "name": self._name,
-                },
-            )
-            attrs = response["project"]["artifact"]
-            manifest_url = attrs["currentManifest"]["file"]["directUrl"]
-            self._manifest = self._load_manifest(manifest_url)
+            if self._client is None:
+                raise RuntimeError("Client not initialized for artifact queries")
+
+            query = gql(FETCH_ARTIFACT_MANIFEST_GQL)
+            gql_vars = {
+                "entityName": self.entity,
+                "projectName": self.project,
+                "name": self.name,
+            }
+            data = self._client.execute(query, variable_values=gql_vars)
+            result = FetchArtifactManifest.model_validate(data)
+            if not (
+                (project := result.project)
+                and (artifact := project.artifact)
+                and (manifest := artifact.current_manifest)
+            ):
+                raise ValueError("Failed to fetch artifact manifest")
+            self._manifest = self._load_manifest(manifest.file.direct_url)
 
         return self._manifest
 
@@ -1112,11 +1080,7 @@ class Artifact:
 
         Includes any references tracked by this artifact.
         """
-        total_size: int = 0
-        for entry in self.manifest.entries.values():
-            if entry.size is not None:
-                total_size += entry.size
-        return total_size
+        return sum(entry.size for entry in self.manifest.entries.values() if entry.size)
 
     @property
     @ensure_logged
@@ -1184,7 +1148,7 @@ class Artifact:
         Returns:
             Boolean. `False` if artifact is saved. `True` if artifact is not saved.
         """
-        return self._state == ArtifactState.PENDING
+        return self._state is ArtifactState.PENDING
 
     def _is_draft_save_started(self) -> bool:
         return self._save_handle is not None
@@ -1205,7 +1169,7 @@ class Artifact:
             settings: A settings object to use when initializing an automatic run. Most
                 commonly used in testing harness.
         """
-        if self._state != ArtifactState.PENDING:
+        if self._state is not ArtifactState.PENDING:
             return self._update()
 
         if self._incremental:
@@ -1266,31 +1230,20 @@ class Artifact:
         return self
 
     def _populate_after_save(self, artifact_id: str) -> None:
-        query_template = """
-            query ArtifactByIDShort($id: ID!) {
-                artifact(id: $id) {
-                    ...ArtifactFragment
-                }
-            }
-        """ + _gql_artifact_fragment()
+        assert self._client is not None
 
-        query = gql(query_template)
+        query = gql_compat(ARTIFACT_BY_ID_GQL, omit_fields=omit_artifact_fields())
 
-        assert self._client is not None
-        response = self._client.execute(
-            query,
-            variable_values={"id": artifact_id},
-        )
+        data = self._client.execute(query, variable_values={"id": artifact_id})
+        result = ArtifactByID.model_validate(data)
 
-        try:
-            attrs = response["artifact"]
-        except LookupError:
+        if not (artifact := result.artifact):
             raise ValueError(f"Unable to fetch artifact with id: {artifact_id!r}")
-        else:
-            # _populate_after_save is only called on source artifacts, not linked artifacts
-            # We have to manually set is_link because we aren't fetching the collection the artifact.
-            # That requires greater refactoring for commitArtifact to return the artifact collection type.
-            self._assign_attrs(attrs, is_link=False)
+
+        # _populate_after_save is only called on source artifacts, not linked artifacts
+        # We have to manually set is_link because we aren't fetching the collection the artifact.
+        # That requires greater refactoring for commitArtifact to return the artifact collection type.
+        self._assign_attrs(artifact, is_link=False)
 
     @normalize_exceptions
     def _update(self) -> None:
@@ -1375,7 +1328,7 @@ class Artifact:
                 for alias in self.aliases
             ]
 
-        omit_fields = omit_artifact_fields(api=InternalApi())
+        omit_fields = omit_artifact_fields()
         omit_variables = set()
 
         if {"ttlIsInherited", "ttlDurationSeconds"} & omit_fields:
@@ -1399,7 +1352,9 @@ class Artifact:
             omit_variables |= {"tagsToAdd", "tagsToDelete"}
 
         mutation = gql_compat(
-            UPDATE_ARTIFACT_GQL, omit_variables=omit_variables, omit_fields=omit_fields
+            UPDATE_ARTIFACT_GQL,
+            omit_variables=omit_variables,
+            omit_fields=omit_fields,
         )
 
         gql_vars = {
@@ -1417,7 +1372,7 @@ class Artifact:
         result = UpdateArtifact.model_validate(data).update_artifact
         if not (result and (artifact := result.artifact)):
             raise ValueError("Unable to parse updateArtifact response")
-        self._assign_attrs(artifact.model_dump())
+        self._assign_attrs(artifact)
 
         self._ttl_changed = False  # Reset after updating artifact
 
@@ -1481,7 +1436,7 @@ class Artifact:
             self._tmp_dir = tempfile.TemporaryDirectory()
         path = os.path.join(self._tmp_dir.name, name.lstrip("/"))
 
-        filesystem.mkdir_exists_ok(os.path.dirname(path))
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
         try:
             with fsync_open(path, mode, encoding) as f:
                 yield f
@@ -1588,30 +1543,27 @@ class Artifact:
             ValueError: Policy must be "mutable" or "immutable"
         """
         if not os.path.isdir(local_path):
-            raise ValueError(f"Path is not a directory: {local_path}")
+            raise ValueError(f"Path is not a directory: {local_path!r}")
 
         termlog(
-            "Adding directory to artifact ({})... ".format(
-                os.path.join(".", os.path.normpath(local_path))
-            ),
+            f"Adding directory to artifact ({Path('.', local_path)})... ",
             newline=False,
         )
-        start_time = time.time()
+        start_time = time.monotonic()
 
-        paths = []
+        paths: deque[tuple[str, str]] = deque()
+        logical_root = name or ""  # shared prefix, if any, for logical paths
         for dirpath, _, filenames in os.walk(local_path, followlinks=True):
             for fname in filenames:
                 physical_path = os.path.join(dirpath, fname)
                 logical_path = os.path.relpath(physical_path, start=local_path)
-                if name is not None:
-                    logical_path = os.path.join(name, logical_path)
+                logical_path = os.path.join(logical_root, logical_path)
                 paths.append((logical_path, physical_path))
 
-        def add_manifest_file(log_phy_path: tuple[str, str]) -> None:
-            logical_path, physical_path = log_phy_path
+        def add_manifest_file(logical_pth: str, physical_pth: str) -> None:
             self._add_local_file(
-                name=logical_path,
-                path=physical_path,
+                name=logical_pth,
+                path=physical_pth,
                 skip_cache=skip_cache,
                 policy=policy,
                 overwrite=merge,
@@ -1619,11 +1571,11 @@ class Artifact:
 
         num_threads = 8
         pool = multiprocessing.dummy.Pool(num_threads)
-        pool.map(add_manifest_file, paths)
+        pool.starmap(add_manifest_file, paths)
         pool.close()
         pool.join()
 
-        termlog("Done. %.1fs" % (time.time() - start_time), prefix=False)
+        termlog("Done. %.1fs" % (time.monotonic() - start_time), prefix=False)
 
     @ensure_not_finalized
     def add_reference(
@@ -1699,7 +1651,7 @@ class Artifact:
                 "References must be URIs. To reference a local file, use file://"
             )
 
-        manifest_entries = self._storage_policy.store_reference(
+        manifest_entries = self.manifest.storage_policy.store_reference(
             self,
             URIStr(uri_str),
             name=name,
@@ -1861,10 +1813,8 @@ class Artifact:
             return
 
         path = str(PurePosixPath(item))
-        entry = self.manifest.get_entry_by_path(path)
-        if entry:
-            self.manifest.remove_entry(entry)
-            return
+        if entry := self.manifest.get_entry_by_path(path):
+            return self.manifest.remove_entry(entry)
 
         entries = self.manifest.get_entries_in_directory(path)
         if not entries:
@@ -1922,8 +1872,7 @@ class Artifact:
 
         # If the entry is a reference from another artifact, then get it directly from
         # that artifact.
-        referenced_id = entry._referenced_artifact_id()
-        if referenced_id:
+        if referenced_id := entry._referenced_artifact_id():
             assert self._client is not None
             artifact = self._from_id(referenced_id, client=self._client)
             assert artifact is not None
@@ -1942,10 +1891,9 @@ class Artifact:
         item_path = item.download()
 
         # Load the object from the JSON blob
-        result = None
-        json_obj = {}
         with open(item_path) as file:
             json_obj = json.load(file)
+
         result = wb_class.from_json(json_obj, self)
         result._set_artifact_source(self, name)
         return result
@@ -1959,10 +1907,9 @@ class Artifact:
         Returns:
             The artifact relative name.
         """
-        entry = self._added_local_paths.get(local_path, None)
-        if entry is None:
-            return None
-        return entry.path
+        if entry := self._added_local_paths.get(local_path):
+            return entry.path
+        return None
 
     def _get_obj_entry(
         self, name: str
@@ -1979,8 +1926,7 @@ class Artifact:
         """
         for wb_class in WBValue.type_mapping().values():
             wandb_file_name = wb_class.with_suffix(name)
-            entry = self.manifest.entries.get(wandb_file_name)
-            if entry is not None:
+            if entry := self.manifest.entries.get(wandb_file_name):
                 return entry, wb_class
         return None, None
 
@@ -2110,14 +2056,14 @@ class Artifact:
         multipart: bool | None = None,
     ) -> FilePathStr:
         nfiles = len(self.manifest.entries)
-        size = sum(e.size or 0 for e in self.manifest.entries.values())
-        log = False
-        if nfiles > 5000 or size > 50 * 1024 * 1024:
-            log = True
+        size_mb = self.size / _MB
+
+        if log := (nfiles > 5000 or size_mb > 50):
             termlog(
-                f"Downloading large artifact {self.name}, {size / (1024 * 1024):.2f}MB. {nfiles} files... ",
+                f"Downloading large artifact {self.name!r}, {size_mb:.2f}MB. {nfiles!r} files...",
             )
-            start_time = datetime.now()
+            start_time = time.monotonic()
+
         download_logger = ArtifactDownloadLogger(nfiles=nfiles)
 
         def _download_entry(
@@ -2156,44 +2102,62 @@ class Artifact:
                 cookies=_thread_local_api_settings.cookies,
                 headers=_thread_local_api_settings.headers,
             )
+
+            batch_size = env.get_artifact_fetch_file_url_batch_size()
+
             active_futures = set()
-            has_next_page = True
-            cursor = None
-            while has_next_page:
-                fetch_url_batch_size = env.get_artifact_fetch_file_url_batch_size()
-                attrs = self._fetch_file_urls(cursor, fetch_url_batch_size)
-                has_next_page = attrs["pageInfo"]["hasNextPage"]
-                cursor = attrs["pageInfo"]["endCursor"]
-                for edge in attrs["edges"]:
-                    entry = self.get_entry(edge["node"]["name"])
+            cursor, has_more = None, True
+            while has_more:
+                files_page = self._fetch_file_urls(cursor=cursor, per_page=batch_size)
+
+                has_more = files_page.page_info.has_next_page
+                cursor = files_page.page_info.end_cursor
+
+                # `File` nodes are formally nullable, so filter them out just in case.
+                file_nodes = (e.node for e in files_page.edges if e.node)
+                for node in file_nodes:
+                    entry = self.get_entry(node.name)
                     # TODO: uncomment once artifact downloads are supported in core
                     # if require_core and entry.ref is None:
                     #     # Handled by core
                     #     continue
-                    entry._download_url = edge["node"]["directUrl"]
+                    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))
+
                 # Wait for download threads to catch up.
-                max_backlog = fetch_url_batch_size
-                if len(active_futures) > max_backlog:
+                #
+                # Extra context and observations (tonyyli):
+                # - Even though the ThreadPoolExecutor limits the number of
+                #   concurrently-executed tasks, its internal task queue is unbounded.
+                #   The code below seems intended to ensure that at most `batch_size`
+                #   "backlogged" futures are held in memory at any given time.  This seems like
+                #   a reasonable safeguard against unbounded memory consumption.
+                #
+                # - We should probably use a builtin (bounded) Queue or Semaphore here instead.
+                #   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):
                         future.result()  # check for errors
                         active_futures.remove(future)
-                        if len(active_futures) <= max_backlog:
+                        if len(active_futures) <= batch_size:
                             break
+
             # Check for errors.
             for future in concurrent.futures.as_completed(active_futures):
                 future.result()
 
         if log:
-            now = datetime.now()
-            delta = abs((now - start_time).total_seconds())
-            hours = int(delta // 3600)
-            minutes = int((delta - hours * 3600) // 60)
-            seconds = delta - hours * 3600 - minutes * 60
-            speed = size / 1024 / 1024 / delta
+            # If you're wondering if we can display a `timedelta`, note that it
+            # doesn't really support custom string format specifiers (compared to
+            # e.g. `datetime` objs).  To truncate the number of decimal places for
+            # the seconds part, we manually convert/format each part below.
+            dt_secs = abs(time.monotonic() - start_time)
+            hrs, mins = divmod(dt_secs, 3600)
+            mins, secs = divmod(mins, 60)
             termlog(
-                f"Done. {hours}:{minutes}:{seconds:.1f} ({speed:.1f}MB/s)",
+                f"Done. {int(hrs):02d}:{int(mins):02d}:{secs:04.1f} ({size_mb / dt_secs:.1f}MB/s)",
                 prefix=False,
             )
         return FilePathStr(root)
@@ -2202,79 +2166,44 @@ class Artifact:
         retry_timedelta=timedelta(minutes=3),
         retryable_exceptions=(requests.RequestException),
     )
-    def _fetch_file_urls(self, cursor: str | None, per_page: int | None = 5000) -> Any:
+    def _fetch_file_urls(
+        self, cursor: str | None, per_page: int = 5000
+    ) -> FileUrlsFragment:
+        if self._client is None:
+            raise RuntimeError("Client not initialized")
+
         if InternalApi()._server_supports(
             pb.ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES
         ):
-            query = gql(
-                """
-                query ArtifactCollectionMembershipFileURLs($entityName: String!, $projectName: String!, \
-                        $artifactName: String!, $artifactVersionIndex: String!, $cursor: String, $perPage: Int) {
-                    project(name: $projectName, entityName: $entityName) {
-                        artifactCollection(name: $artifactName) {
-                            artifactMembership(aliasName: $artifactVersionIndex) {
-                                files(after: $cursor, first: $perPage) {
-                                    pageInfo {
-                                        hasNextPage
-                                        endCursor
-                                    }
-                                    edges {
-                                        node {
-                                            name
-                                            directUrl
-                                        }
-                                    }
-                                }
-                            }
-                        }
-                    }
-                }
-                """
-            )
-            assert self._client is not None
-            response = self._client.execute(
-                query,
-                variable_values={
-                    "entityName": self.entity,
-                    "projectName": self.project,
-                    "artifactName": self.name.split(":")[0],
-                    "artifactVersionIndex": self.version,
-                    "cursor": cursor,
-                    "perPage": per_page,
-                },
-                timeout=60,
-            )
-            return response["project"]["artifactCollection"]["artifactMembership"][
-                "files"
-            ]
+            query = gql(ARTIFACT_COLLECTION_MEMBERSHIP_FILE_URLS_GQL)
+            gql_vars = {
+                "entityName": self.entity,
+                "projectName": self.project,
+                "artifactName": self.name.split(":")[0],
+                "artifactVersionIndex": self.version,
+                "cursor": cursor,
+                "perPage": per_page,
+            }
+            data = self._client.execute(query, variable_values=gql_vars, timeout=60)
+            result = ArtifactCollectionMembershipFileUrls.model_validate(data)
+
+            if not (
+                (project := result.project)
+                and (collection := project.artifact_collection)
+                and (membership := collection.artifact_membership)
+                and (files := membership.files)
+            ):
+                raise ValueError(f"Unable to fetch files for artifact: {self.name!r}")
+            return files
         else:
-            query = gql(
-                """
-                query ArtifactFileURLs($id: ID!, $cursor: String, $perPage: Int) {
-                    artifact(id: $id) {
-                        files(after: $cursor, first: $perPage) {
-                            pageInfo {
-                                hasNextPage
-                                endCursor
-                            }
-                            edges {
-                                node {
-                                    name
-                                    directUrl
-                                }
-                            }
-                        }
-                    }
-                }
-                """
-            )
-            assert self._client is not None
-            response = self._client.execute(
-                query,
-                variable_values={"id": self.id, "cursor": cursor, "perPage": per_page},
-                timeout=60,
-            )
-            return response["artifact"]["files"]
+            query = gql(ARTIFACT_FILE_URLS_GQL)
+            gql_vars = {"id": self.id, "cursor": cursor, "perPage": per_page}
+            data = self._client.execute(query, variable_values=gql_vars, timeout=60)
+            result = ArtifactFileUrls.model_validate(data)
+
+            if not ((artifact := result.artifact) and (files := artifact.files)):
+                raise ValueError(f"Unable to fetch files for artifact: {self.name!r}")
+            return files
 
     @ensure_logged
     def checkout(self, root: str | None = None) -> str:
@@ -2440,33 +2369,16 @@ class Artifact:
 
     @normalize_exceptions
     def _delete(self, delete_aliases: bool = False) -> None:
-        mutation = gql(
-            """
-            mutation DeleteArtifact($artifactID: ID!, $deleteAliases: Boolean) {
-                deleteArtifact(input: {
-                    artifactID: $artifactID
-                    deleteAliases: $deleteAliases
-                }) {
-                    artifact {
-                        id
-                    }
-                }
-            }
-            """
-        )
-        assert self._client is not None
-        self._client.execute(
-            mutation,
-            variable_values={
-                "artifactID": self.id,
-                "deleteAliases": delete_aliases,
-            },
-        )
+        if self._client is None:
+            raise RuntimeError("Client not initialized for artifact mutations")
+
+        mutation = gql(DELETE_ARTIFACT_GQL)
+        gql_vars = {"artifactID": self.id, "deleteAliases": delete_aliases}
+
+        self._client.execute(mutation, variable_values=gql_vars)
 
     @normalize_exceptions
-    def link(
-        self, target_path: str, aliases: list[str] | None = None
-    ) -> Artifact | None:
+    def link(self, target_path: str, aliases: list[str] | None = None) -> Artifact:
         """Link this artifact to a portfolio (a promoted collection of artifacts).
 
         Args:
@@ -2485,7 +2397,7 @@ class Artifact:
             ArtifactNotLoggedError: If the artifact is not logged.
 
         Returns:
-            The linked artifact if linking was successful, otherwise None.
+            The linked artifact.
         """
         from wandb import Api
 
@@ -2505,10 +2417,11 @@ class Artifact:
             # Wait until the artifact is committed before trying to link it.
             self.wait()
 
-        api = Api(overrides={"entity": self.source_entity})
+        api = InternalApi()
+        settings = api.settings()
 
         target = ArtifactPath.from_str(target_path).with_defaults(
-            project=api.settings.get("project") or "uncategorized",
+            project=settings.get("project") or "uncategorized",
         )
 
         # Parse the entity (first part of the path) appropriately,
@@ -2516,11 +2429,8 @@ class Artifact:
         if target.project and is_artifact_registry_project(target.project):
             # 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
-            organization = target.prefix or api.settings.get("organization") or ""
-
-            target.prefix = InternalApi()._resolve_org_entity_name(
-                self.source_entity, organization
-            )
+            org = target.prefix or settings.get("organization") or ""
+            target.prefix = api._resolve_org_entity_name(self.source_entity, org)
         else:
             target = target.with_defaults(prefix=self.source_entity)
 
@@ -2537,21 +2447,35 @@ class Artifact:
             aliases=alias_inputs,
         )
         gql_vars = {"input": gql_input.model_dump(exclude_none=True)}
-        gql_op = gql(LINK_ARTIFACT_GQL)
-        data = self._client.execute(gql_op, variable_values=gql_vars)
 
+        # Newer server versions can return `artifactMembership` directly in the response,
+        # avoiding the need to re-fetch the linked artifact at the end.
+        if api._server_supports(
+            pb.ServerFeature.ARTIFACT_MEMBERSHIP_IN_LINK_ARTIFACT_RESPONSE
+        ):
+            omit_fragments = set()
+        else:
+            # FIXME: Make `gql_compat` omit nested fragment definitions recursively (but safely)
+            omit_fragments = {
+                "MembershipWithArtifact",
+                "ArtifactFragment",
+                "ArtifactFragmentWithoutAliases",
+            }
+
+        gql_op = gql_compat(LINK_ARTIFACT_GQL, omit_fragments=omit_fragments)
+        data = self._client.execute(gql_op, variable_values=gql_vars)
         result = LinkArtifact.model_validate(data).link_artifact
+
+        # Newer server versions can return artifactMembership directly in the response
+        if result and (membership := result.artifact_membership):
+            return self._from_membership(membership, target=target, client=self._client)
+
+        # Fallback to old behavior, which requires re-fetching the linked artifact to return it
         if not (result and (version_idx := result.version_index) is not None):
             raise ValueError("Unable to parse linked artifact version from response")
 
-        # Fetch the linked artifact to return it
-        linked_path = f"{target.to_str()}:v{version_idx}"
-
-        try:
-            return api._artifact(linked_path)
-        except Exception as e:
-            wandb.termerror(f"Error fetching link artifact after linking: {e}")
-            return None
+        link_name = f"{target.to_str()}:v{version_idx}"
+        return Api(overrides={"entity": self.source_entity})._artifact(link_name)
 
     @ensure_logged
     def unlink(self) -> None:
@@ -2573,28 +2497,14 @@ class Artifact:
 
     @normalize_exceptions
     def _unlink(self) -> None:
-        mutation = gql(
-            """
-            mutation UnlinkArtifact($artifactID: ID!, $artifactPortfolioID: ID!) {
-                unlinkArtifact(
-                    input: { artifactID: $artifactID, artifactPortfolioID: $artifactPortfolioID }
-                ) {
-                    artifactID
-                    success
-                    clientMutationId
-                }
-            }
-            """
-        )
-        assert self._client is not None
+        if self._client is None:
+            raise RuntimeError("Client not initialized for artifact mutations")
+
+        mutation = gql(UNLINK_ARTIFACT_GQL)
+        gql_vars = {"artifactID": self.id, "artifactPortfolioID": self.collection.id}
+
         try:
-            self._client.execute(
-                mutation,
-                variable_values={
-                    "artifactID": self.id,
-                    "artifactPortfolioID": self.collection.id,
-                },
-            )
+            self._client.execute(mutation, variable_values=gql_vars)
         except CommError as e:
             raise CommError(
                 f"You do not have permission to unlink the artifact {self.qualified_name}"
@@ -2610,41 +2520,26 @@ class Artifact:
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
         """
-        query = gql(
-            """
-            query ArtifactUsedBy(
-                $id: ID!,
-            ) {
-                artifact(id: $id) {
-                    usedBy {
-                        edges {
-                            node {
-                                name
-                                project {
-                                    name
-                                    entityName
-                                }
-                            }
-                        }
-                    }
-                }
-            }
-            """
-        )
-        assert self._client is not None
-        response = self._client.execute(
-            query,
-            variable_values={"id": self.id},
-        )
-        return [
-            Run(
-                self._client,
-                edge["node"]["project"]["entityName"],
-                edge["node"]["project"]["name"],
-                edge["node"]["name"],
-            )
-            for edge in response.get("artifact", {}).get("usedBy", {}).get("edges", [])
-        ]
+        if self._client is None:
+            raise RuntimeError("Client not initialized for artifact queries")
+
+        query = gql(ARTIFACT_USED_BY_GQL)
+        gql_vars = {"id": self.id}
+        data = self._client.execute(query, variable_values=gql_vars)
+        result = ArtifactUsedBy.model_validate(data)
+
+        if (
+            (artifact := result.artifact)
+            and (used_by := artifact.used_by)
+            and (edges := used_by.edges)
+        ):
+            run_nodes = (e.node for e in edges)
+            return [
+                Run(self._client, proj.entity_name, proj.name, run.name)
+                for run in run_nodes
+                if (proj := run.project)
+            ]
+        return []
 
     @ensure_logged
     def logged_by(self) -> Run | None:
@@ -2656,39 +2551,22 @@ class Artifact:
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
         """
-        query = gql(
-            """
-            query ArtifactCreatedBy(
-                $id: ID!
-            ) {
-                artifact(id: $id) {
-                    createdBy {
-                        ... on Run {
-                            name
-                            project {
-                                name
-                                entityName
-                            }
-                        }
-                    }
-                }
-            }
-        """
-        )
-        assert self._client is not None
-        response = self._client.execute(
-            query,
-            variable_values={"id": self.id},
-        )
-        creator = response.get("artifact", {}).get("createdBy", {})
-        if creator.get("name") is None:
-            return None
-        return Run(
-            self._client,
-            creator["project"]["entityName"],
-            creator["project"]["name"],
-            creator["name"],
-        )
+        if self._client is None:
+            raise RuntimeError("Client not initialized for artifact queries")
+
+        query = gql(ARTIFACT_CREATED_BY_GQL)
+        gql_vars = {"id": self.id}
+        data = self._client.execute(query, variable_values=gql_vars)
+        result = ArtifactCreatedBy.model_validate(data)
+
+        if (
+            (artifact := result.artifact)
+            and (creator := artifact.created_by)
+            and (name := creator.name)
+            and (project := creator.project)
+        ):
+            return Run(self._client, project.entity_name, project.name, name)
+        return None
 
     @ensure_logged
     def json_encode(self) -> dict[str, Any]:
@@ -2704,37 +2582,22 @@ class Artifact:
         entity_name: str, project_name: str, name: str, client: RetryingClient
     ) -> str | None:
         """Returns the expected type for a given artifact name and project."""
-        query = gql(
-            """
-            query ArtifactType(
-                $entityName: String,
-                $projectName: String,
-                $name: String!
-            ) {
-                project(name: $projectName, entityName: $entityName) {
-                    artifact(name: $name) {
-                        artifactType {
-                            name
-                        }
-                    }
-                }
-            }
-            """
-        )
-        if ":" not in name:
-            name += ":latest"
-        response = client.execute(
-            query,
-            variable_values={
-                "entityName": entity_name,
-                "projectName": project_name,
-                "name": name,
-            },
-        )
-        return (
-            ((response.get("project") or {}).get("artifact") or {}).get("artifactType")
-            or {}
-        ).get("name")
+        query = gql(ARTIFACT_TYPE_GQL)
+        gql_vars = {
+            "entityName": entity_name,
+            "projectName": project_name,
+            "name": name if (":" in name) else f"{name}:latest",
+        }
+        data = client.execute(query, variable_values=gql_vars)
+        result = ArtifactType.model_validate(data)
+
+        if (
+            (project := result.project)
+            and (artifact := project.artifact)
+            and (artifact_type := artifact.artifact_type)
+        ):
+            return artifact_type.name
+        return None
 
     def _load_manifest(self, url: str) -> ArtifactManifest:
         with requests.get(url) as response: