fix(artifacts): when artifact-commit 409s, retry entire artifact-creation, not just commit

tests/unit_tests/conftest.py

@@ -22,6 +22,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    ContextManager,
     Dict,
     Generator,
     Iterable,
@@ -49,9 +50,9 @@
 from wandb.sdk.lib.git import GitRepo
 
 try:
-    from typing import Literal, TypedDict
+    from typing import Literal, Protocol, TypedDict
 except ImportError:
-    from typing_extensions import Literal, TypedDict
+    from typing_extensions import Literal, Protocol, TypedDict
 
 if TYPE_CHECKING:
 
@@ -74,6 +75,12 @@ class Resolver(TypedDict):
         resolver: Callable[[Any], Optional[Dict[str, Any]]]
 
 
+if TYPE_CHECKING:
+    import wandb_gql
+else:
+    wandb_gql = wandb.util.vendor_import("wandb_gql")
+
+
 class ConsoleFormatter:
     BOLD = "\033[1m"
     CODE = "\033[2m"
@@ -1054,7 +1061,7 @@ def config(self) -> Dict[str, Any]:
         if self._config is not None:
             return deepcopy(self._config)
 
-        self._config = {k: v["config"] for (k, v) in self._entries.items()}
+        self._config = {k: v.get("config") for (k, v) in self._entries.items()}
         return deepcopy(self._config)
 
     # @property
@@ -1142,9 +1149,14 @@ def __init__(self):
                 "name": "upsert_sweep",
                 "resolver": self.resolve_upsert_sweep,
             },
-            # { "name": "create_artifact",
-            #     "resolver": self.resolve_create_artifact,
-            # },
+            {
+                "name": "create_artifact",
+                "resolver": self.resolve_create_artifact,
+            },
+            {
+                "name": "commit_artifact",
+                "resolver": self.resolve_commit_artifact,
+            },
         ]
 
     @staticmethod
@@ -1242,28 +1254,42 @@ def resolve_upsert_sweep(
             return data
         return None
 
+    @staticmethod
     def resolve_create_artifact(
-        self, request_data: Dict[str, Any], response_data: Dict[str, Any], **kwargs: Any
+        request_data: Dict[str, Any], response_data: Dict[str, Any], **kwargs: Any
     ) -> Optional[Dict[str, Any]]:
-        if not isinstance(request_data, dict):
+        if not isinstance(response_data, dict):
             return None
-        query = (
-            "createArtifact(" in request_data.get("query", "")
-            and request_data.get("variables") is not None
-            and response_data is not None
-        )
+        query = response_data.get("data", {}).get("createArtifact") is not None
         if query:
-            name = request_data["variables"]["runName"]
-            post_processed_data = {
-                "name": name,
+            return {
+                "name": response_data["data"]["createArtifact"]["artifact"]["id"],
                 "create_artifact": [
                     {
                         "variables": request_data["variables"],
-                        "response": response_data["data"]["createArtifact"]["artifact"],
                     }
                 ],
             }
-            return post_processed_data
+        return None
+
+    @staticmethod
+    def resolve_commit_artifact(
+        request_data: Dict[str, Any], response_data: Dict[str, Any], **kwargs: Any
+    ) -> Optional[Dict[str, Any]]:
+        if not isinstance(response_data, dict):
+            return None
+        if "query" not in request_data:
+            return None
+        query = response_data.get("data", {}).get("commitArtifact") is not None
+        if query:
+            return {
+                "name": request_data["variables"]["artifactID"],
+                "commit_artifact": [
+                    {
+                        "variables": request_data["variables"],
+                    }
+                ],
+            }
         return None
 
     def resolve(
@@ -1289,6 +1315,7 @@ class InjectedResponse:
     content_type: str = "text/plain"
     # todo: add more fields for other types of responses?
     counter: int = -1
+    predicate: Optional[Callable[[requests.PreparedRequest], bool]] = None
 
     def __eq__(
         self,
@@ -1307,14 +1334,21 @@ def __eq__(
             return False
         if self.counter == 0:
             return False
-        # todo: add more fields for other types of responses?
-        return self.method == other.method and self.url == other.url
+        if self.method != other.method or self.url != other.url:
+            return False
+        if (
+            isinstance(other, requests.PreparedRequest)
+            and self.predicate is not None
+            and not self.predicate(other)
+        ):
+            return False
+        return True
 
     def to_dict(self):
         return {
             k: self.__getattribute__(k)
             for k in self.__dict__
-            if (not k.startswith("_") and k != "counter")
+            if (not k.startswith("_") and k not in {"counter", "predicate"})
         }
 
 
@@ -1534,6 +1568,13 @@ def storage_file(self, path) -> Mapping[str, str]:
         return relayed_response.json()
 
 
+class RelayServerFixture(Protocol):
+    def __call__(
+        self, inject: Optional[List[InjectedResponse]] = None
+    ) -> ContextManager[RelayServer]:
+        ...
+
+
 @pytest.fixture(scope="function")
 def relay_server(base_url):
     """
@@ -1677,3 +1718,46 @@ def helper(
         )
 
     yield helper
+
+
+class InjectedGraphQLRequestCreator(Protocol):
+    def __call__(
+        self,
+        body: Union[str, Exception] = "{}",
+        status: int = 200,
+        counter: int = -1,
+    ) -> InjectedResponse:
+        ...
+
+
+# Injected responses
+@pytest.fixture(scope="function")
+def inject_graphql_response(base_url: str) -> InjectedGraphQLRequestCreator:
+    def helper(
+        operation_name: str,
+        body: Union[str, Exception] = "{}",
+        status: int = 200,
+        counter: int = -1,
+    ) -> InjectedResponse:
+
+        if status > 299:
+            message = body if isinstance(body, str) else "::".join(body.args)
+            body = DeliberateHTTPError(status_code=status, message=message)
+
+        def predicate(request: requests.PreparedRequest) -> bool:
+            query = wandb_gql.gql(json.loads(request.body)["query"])
+            return query.definitions[0].name.value == operation_name
+
+        return InjectedResponse(
+            method="POST",
+            url=urllib.parse.urljoin(
+                base_url,
+                "/graphql",
+            ),
+            body=body,
+            status=status,
+            counter=counter,
+            predicate=predicate,
+        )
+
+    return helper

tests/unit_tests/test_artifacts.py

@@ -0,0 +1,30 @@
+from typing import TYPE_CHECKING, Callable
+
+import pytest
+import wandb
+
+if TYPE_CHECKING:
+    from .conftest import InjectedGraphQLRequestCreator, RelayServerFixture
+
+
+@pytest.mark.parametrize("status", [409, 500])
+def test_commit_retries_on_right_statuses(
+    relay_server: "RelayServerFixture",
+    wandb_init: Callable[[], wandb.wandb_sdk.wandb_run.Run],
+    inject_graphql_response: "InjectedGraphQLRequestCreator",
+    status: int,
+):
+    art = wandb.Artifact("test", "dataset")
+
+    injected_resp = inject_graphql_response(
+        operation_name="CommitArtifact",
+        status=status,
+        counter=1,
+    )
+    with relay_server(inject=[injected_resp]) as relay:
+        run = wandb_init()
+        logged = run.log_artifact(art).wait()
+        run.finish()
+
+    # even though we made two requests, empirically only the successful one goes into the Context
+    assert relay.context.entries[logged.id]["commit_artifact"]

wandb/filesync/step_upload.py

@@ -36,7 +36,7 @@ class ArtifactStatus(TypedDict):
 
 
 PreCommitFn = Callable[[], None]
-PostCommitFn = Callable[[], None]
+PostCommitFn = Callable[[Optional[Exception]], None]
 OnRequestFinishFn = Callable[[], None]
 SaveFn = Callable[["progress.ProgressFn"], Any]
 
@@ -216,12 +216,16 @@ def _maybe_commit_artifact(self, artifact_id: str) -> None:
             artifact_status["pending_count"] == 0
             and artifact_status["commit_requested"]
         ):
-            for callback in artifact_status["pre_commit_callbacks"]:
-                callback()
+            for pre_callback in artifact_status["pre_commit_callbacks"]:
+                pre_callback()
+            exc = None
             if artifact_status["finalize"]:
-                self._api.commit_artifact(artifact_id)
-            for callback in artifact_status["post_commit_callbacks"]:
-                callback()
+                try:
+                    self._api.commit_artifact(artifact_id)
+                except Exception as e:
+                    exc = e
+            for post_callback in artifact_status["post_commit_callbacks"]:
+                post_callback(exc)
 
     def start(self) -> None:
         self._thread.start()