wandb 0.21.0__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,
@@ -80,6 +115,7 @@ from ._validators import (
     ensure_logged,
     ensure_not_finalized,
     is_artifact_registry_project,
+    remove_registry_prefix,
     validate_aliases,
     validate_artifact_name,
     validate_artifact_type,
@@ -102,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()
 
@@ -118,6 +151,9 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+_MB: Final[int] = 1024 * 1024
+
+
 @final
 @dataclass
 class _DeferredArtifactManifest:
@@ -131,25 +167,26 @@ class Artifact:
 
     Construct an empty W&B Artifact. Populate an artifacts contents with methods that
     begin with `add`. Once the artifact has all the desired files, you can call
-    `wandb.log_artifact()` to log it.
+    `run.log_artifact()` to log it.
 
     Args:
-        name: A human-readable name for the artifact. Use the name to identify
+        name (str): A human-readable name for the artifact. Use the name to identify
             a specific artifact in the W&B App UI or programmatically. You can
             interactively reference an artifact with the `use_artifact` Public API.
             A name can contain letters, numbers, underscores, hyphens, and dots.
             The name must be unique across a project.
-        type: The artifact's type. Use the type of an artifact to both organize
+        type (str): The artifact's type. Use the type of an artifact to both organize
             and differentiate artifacts. You can use any string that contains letters,
             numbers, underscores, hyphens, and dots. Common types include `dataset` or `model`.
-            Note: Some types are reserved for internal use and cannot be set by users.
-                Such types include `job` and types that start with `wandb-`.
-        description: A description of the artifact. For Model or Dataset Artifacts,
+            Include `model` within your type string if you want to link the artifact
+            to the W&B Model Registry. Note that some types reserved for internal use
+            and cannot be set by users. Such types include `job` and types that start with `wandb-`.
+        description (str | None) = None: A description of the artifact. For Model or Dataset Artifacts,
             add documentation for your standardized team model or dataset card. View
             an artifact's description programmatically with the `Artifact.description`
             attribute or programmatically with the W&B App UI. W&B renders the
             description as markdown in the W&B App.
-        metadata: Additional information about an artifact. Specify metadata as a
+        metadata (dict[str, Any] | None) = None: Additional information about an artifact. Specify metadata as a
             dictionary of key-value pairs. You can specify no more than 100 total keys.
         incremental: Use `Artifact.new_draft()` method instead to modify an
             existing artifact.
@@ -186,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] = {}
@@ -234,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
@@ -255,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(
@@ -291,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(
@@ -299,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(
@@ -373,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()
-        )
-        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()}
-            """
+        query = gql_compat(
+            ARTIFACT_BY_NAME_GQL,
+            omit_variables=omit_vars,
+            omit_fields=omit_artifact_fields(api=api),
         )
-        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(
@@ -433,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:
@@ -443,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()
 
@@ -456,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
@@ -486,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))
@@ -522,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:
@@ -839,7 +822,7 @@ class Artifact:
         )
         return urljoin(
             base_url,
-            f"orgs/{org.display_name}/registry/{self._type}?selectionPath={selection_path}&view=membership&version={self._version}",
+            f"orgs/{org.display_name}/registry/{remove_registry_prefix(self.project)}?selectionPath={selection_path}&view=membership&version={self.version}",
         )
 
     def _construct_model_registry_url(self, base_url: str) -> str:
@@ -924,7 +907,8 @@ class Artifact:
         TTL and there is no custom policy set on an artifact.
 
         Raises:
-            ArtifactNotLoggedError: Unable to fetch inherited TTL if the artifact has not been logged or saved
+            ArtifactNotLoggedError: Unable to fetch inherited TTL if the
+            artifact has not been logged or saved.
         """
         if self._ttl_is_inherited and (self.is_draft() or self._ttl_changed):
             raise ArtifactNotLoggedError(f"{type(self).__name__}.ttl", self)
@@ -976,7 +960,9 @@ class Artifact:
     @property
     @ensure_logged
     def aliases(self) -> list[str]:
-        """List of one or more semantically-friendly references or identifying "nicknames" assigned to an artifact version.
+        """List of one or more semantically-friendly references or
+
+        identifying "nicknames" assigned to an artifact version.
 
         Aliases are mutable references that you can programmatically reference.
         Change an artifact's alias with the W&B App UI or programmatically.
@@ -1012,6 +998,10 @@ class Artifact:
 
     @property
     def distributed_id(self) -> str | None:
+        """The distributed ID of the artifact.
+
+        <!-- lazydoc-ignore: internal -->
+        """
         return self._distributed_id
 
     @distributed_id.setter
@@ -1020,6 +1010,10 @@ class Artifact:
 
     @property
     def incremental(self) -> bool:
+        """Boolean flag indicating if the artifact is an incremental artifact.
+
+        <!-- lazydoc-ignore: internal -->
+        """
         return self._incremental
 
     @property
@@ -1050,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
 
@@ -1099,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
@@ -1139,15 +1116,15 @@ class Artifact:
         """The nearest step at which history metrics were logged for the source run of the artifact.
 
         Examples:
-            ```python
-            run = artifact.logged_by()
-            if run and (artifact.history_step is not None):
-                history = run.sample_history(
-                    min_step=artifact.history_step,
-                    max_step=artifact.history_step + 1,
-                    keys=["my_metric"],
-                )
-            ```
+        ```python
+        run = artifact.logged_by()
+        if run and (artifact.history_step is not None):
+            history = run.sample_history(
+                min_step=artifact.history_step,
+                max_step=artifact.history_step + 1,
+                keys=["my_metric"],
+            )
+        ```
         """
         if self._history_step is None:
             return None
@@ -1168,9 +1145,10 @@ class Artifact:
     def is_draft(self) -> bool:
         """Check if artifact is not saved.
 
-        Returns: Boolean. `False` if artifact is saved. `True` if artifact is not saved.
+        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
@@ -1191,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:
@@ -1252,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:
@@ -1361,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:
@@ -1385,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 = {
@@ -1403,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
 
@@ -1416,7 +1385,7 @@ class Artifact:
             name: The artifact relative name to get.
 
         Returns:
-            W&B object that can be logged with `wandb.log()` and visualized in the W&B UI.
+            W&B object that can be logged with `run.log()` and visualized in the W&B UI.
 
         Raises:
             ArtifactNotLoggedError: If the artifact isn't logged or the run is offline.
@@ -1434,8 +1403,9 @@ class Artifact:
             The added manifest entry
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+            artifact version because it is finalized. Log a new artifact
+            version instead.
         """
         return self.add(item, name)
 
@@ -1452,12 +1422,13 @@ class Artifact:
             encoding: The encoding used to open the new file.
 
         Returns:
-            A new file object that can be written to. Upon closing, the file will be
-            automatically added to the artifact.
+            A new file object that can be written to. Upon closing, the file
+            is automatically added to the artifact.
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+            artifact version because it is finalized. Log a new artifact
+            version instead.
         """
         overwrite: bool = "x" not in mode
 
@@ -1465,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
@@ -1496,22 +1467,26 @@ class Artifact:
 
         Args:
             local_path: The path to the file being added.
-            name: The path within the artifact to use for the file being added. Defaults
-                to the basename of the file.
+            name: The path within the artifact to use for the file being added.
+                Defaults to the basename of the file.
             is_tmp: If true, then the file is renamed deterministically to avoid
                 collisions.
-            skip_cache: If `True`, W&B will not copy files to the cache after uploading.
-            policy: By default, set to "mutable". If set to "mutable", create a temporary copy of the
-                file to prevent corruption during upload. If set to "immutable", disable
-                protection and rely on the user not to delete or change the file.
+            skip_cache: If `True`, do not copy files to the cache
+                after uploading.
+            policy: By default, set to "mutable". If set to "mutable",
+                create a temporary copy of the file to prevent corruption
+                during upload. If set to "immutable", disable
+                protection and rely on the user not to delete or change the
+                file.
             overwrite: If `True`, overwrite the file if it already exists.
 
         Returns:
             The added manifest entry.
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+                artifact version because it is finalized. Log a new artifact
+                version instead.
             ValueError: Policy must be "mutable" or "immutable"
         """
         if not os.path.isfile(local_path):
@@ -1548,48 +1523,47 @@ class Artifact:
 
         Args:
             local_path: The path of the local directory.
-            name: The subdirectory name within an artifact. The name you specify appears
-                in the W&B App UI nested by artifact's `type`.
+            name: The subdirectory name within an artifact. The name you
+                specify appears in the W&B App UI nested by artifact's `type`.
                 Defaults to the root of the artifact.
-            skip_cache: If set to `True`, W&B will not copy/move files to the cache while uploading
-            policy: "mutable" | "immutable". By default, "mutable"
-                "mutable": Create a temporary copy of the file to prevent corruption during upload.
-                "immutable": Disable protection, rely on the user not to delete or change the file.
+            skip_cache: If set to `True`, W&B will not copy/move files to
+                the cache while uploading
+            policy: By default, "mutable".
+            - mutable: Create a temporary copy of the file to prevent corruption during upload.
+            - immutable: Disable protection, rely on the user not to delete or change the file.
             merge: If `False` (default), throws ValueError if a file was already added in a previous add_dir call
                 and its content has changed. If `True`, overwrites existing files with changed content.
                 Always adds new files and never removes files. To replace an entire directory, pass a name when adding the directory
                 using `add_dir(local_path, name=my_prefix)` and call `remove(my_prefix)` to remove the directory, then add it again.
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+                artifact version because it is finalized. Log a new artifact
+                version instead.
             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,
@@ -1597,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(
@@ -1621,13 +1595,14 @@ class Artifact:
 
         - http(s): The size and digest of the file will be inferred by the
           `Content-Length` and the `ETag` response headers returned by the server.
-        - s3: The checksum and size are pulled from the object metadata. If bucket
-          versioning is enabled, then the version ID is also tracked.
+        - s3: The checksum and size are pulled from the object metadata.
+          If bucket versioning is enabled, then the version ID is also tracked.
         - gs: The checksum and size are pulled from the object metadata. If bucket
           versioning is enabled, then the version ID is also tracked.
-        - https, domain matching `*.blob.core.windows.net` (Azure): The checksum and size
-          are be pulled from the blob metadata. If storage account versioning is
-          enabled, then the version ID is also tracked.
+        - https, domain matching `*.blob.core.windows.net`
+        - Azure: The checksum and size are be pulled from the blob metadata.
+          If storage account versioning is enabled, then the version ID is
+          also tracked.
         - file: The checksum and size are pulled from the file system. This scheme
           is useful if you have an NFS share or other externally mounted volume
           containing files you wish to track but not necessarily upload.
@@ -1648,16 +1623,18 @@ class Artifact:
                 setting `checksum=False` when adding reference objects, in which case
                 a new version will only be created if the reference URI changes.
             max_objects: The maximum number of objects to consider when adding a
-                reference that points to directory or bucket store prefix. By default,
-                the maximum number of objects allowed for Amazon S3,
-                GCS, Azure, and local files is 10,000,000. Other URI schemas do not have a maximum.
+                reference that points to directory or bucket store prefix.
+                By default, the maximum number of objects allowed for Amazon S3,
+                GCS, Azure, and local files is 10,000,000. Other URI schemas
+                do not have a maximum.
 
         Returns:
             The added manifest entries.
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+            artifact version because it is finalized. Log a new artifact
+            version instead.
         """
         if name is not None:
             name = LogicalPath(name)
@@ -1674,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,
@@ -1694,17 +1671,19 @@ class Artifact:
 
         Args:
             obj: The object to add. Currently support one of Bokeh, JoinedTable,
-                PartitionedTable, Table, Classes, ImageMask, BoundingBoxes2D, Audio,
-                Image, Video, Html, Object3D
+                PartitionedTable, Table, Classes, ImageMask, BoundingBoxes2D,
+                Audio, Image, Video, Html, Object3D
             name: The path within the artifact to add the object.
-            overwrite: If True, overwrite existing objects with the same file path (if applicable).
+            overwrite: If True, overwrite existing objects with the same file
+                path if applicable.
 
         Returns:
             The added manifest entry
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+            artifact version because it is finalized. Log a new artifact
+            version instead.
         """
         name = LogicalPath(name)
 
@@ -1819,13 +1798,14 @@ class Artifact:
         """Remove an item from the artifact.
 
         Args:
-            item: The item to remove. Can be a specific manifest entry or the name of an
-                artifact-relative path. If the item matches a directory all items in
-                that directory will be removed.
+            item: The item to remove. Can be a specific manifest entry
+                or the name of an artifact-relative path. If the item
+                matches a directory all items in that directory will be removed.
 
         Raises:
-            ArtifactFinalizedError: You cannot make changes to the current artifact
-            version because it is finalized. Log a new artifact version instead.
+            ArtifactFinalizedError: You cannot make changes to the current
+                artifact version because it is finalized. Log a new artifact
+                version instead.
             FileNotFoundError: If the item isn't found in the artifact.
         """
         if isinstance(item, ArtifactManifestEntry):
@@ -1833,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:
@@ -1881,10 +1859,12 @@ class Artifact:
             name: The artifact relative name to retrieve.
 
         Returns:
-            W&B object that can be logged with `wandb.log()` and visualized in the W&B UI.
+            W&B object that can be logged with `run.log()` and
+            visualized in the W&B UI.
 
         Raises:
-            ArtifactNotLoggedError: if the artifact isn't logged or the run is offline
+            ArtifactNotLoggedError: if the artifact isn't logged or the
+                run is offline.
         """
         entry, wb_class = self._get_obj_entry(name)
         if entry is None or wb_class is None:
@@ -1892,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
@@ -1912,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
@@ -1929,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
@@ -1949,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
 
@@ -2080,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(
@@ -2126,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)
@@ -2172,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:
@@ -2320,7 +2279,7 @@ class Artifact:
 
         Args:
             root: The root directory to store the file. Defaults to
-                './artifacts/self.name/'.
+                `./artifacts/self.name/`.
 
         Returns:
             The full path of the downloaded file.
@@ -2386,16 +2345,16 @@ class Artifact:
     def delete(self, delete_aliases: bool = False) -> None:
         """Delete an artifact and its files.
 
-        If called on a linked artifact (i.e. a member of a portfolio collection): only the link is deleted, and the
+        If called on a linked artifact, only the link is deleted, and the
         source artifact is unaffected.
 
         Use `artifact.unlink()` instead of `artifact.delete()` to remove a link between a source artifact and a linked artifact.
 
         Args:
-            delete_aliases: If set to `True`, deletes all aliases associated with the artifact.
-                Otherwise, this raises an exception if the artifact has existing
-                aliases.
-                This parameter is ignored if the artifact is linked (i.e. a member of a portfolio collection).
+            delete_aliases: If set to `True`, deletes all aliases associated
+                with the artifact. Otherwise, this raises an exception if
+                the artifact has existing aliases. This parameter is ignored
+                if the artifact is linked (a member of a portfolio collection).
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
@@ -2410,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:
@@ -2448,14 +2390,14 @@ class Artifact:
                 portfolio inside a project, set `target_path` to the following
                 schema `{"model-registry"}/{Registered Model Name}` or
                 `{entity}/{"model-registry"}/{Registered Model Name}`.
-            aliases: A list of strings that uniquely identifies the artifact inside the
-                specified portfolio.
+            aliases: A list of strings that uniquely identifies the artifact
+                inside the specified portfolio.
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
 
         Returns:
-            The linked artifact if linking was successful, otherwise None.
+            The linked artifact.
         """
         from wandb import Api
 
@@ -2475,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,
@@ -2486,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)
 
@@ -2507,29 +2447,44 @@ 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:
-        """Unlink this artifact if it is currently a member of a portfolio (a promoted collection of artifacts).
+        """Unlink this artifact if it is currently a member of a promoted collection of artifacts.
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
-            ValueError: If the artifact is not linked, i.e. it is not a member of a portfolio collection.
+            ValueError: If the artifact is not linked, in other words,
+            it is not a member of a portfolio collection.
         """
         # Fail early if this isn't a linked artifact to begin with
         if not self.is_link:
@@ -2542,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}"
@@ -2579,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:
@@ -2625,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]:
@@ -2673,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: