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

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()

wandb/sdk/internal/artifacts.py

@@ -1,15 +1,29 @@
+import datetime
 import json
 import os
 import sys
 import tempfile
 import threading
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    cast,
+)
+
+import requests
 
 import wandb
 import wandb.filesync.step_prepare
+import wandb.sdk.lib.retry
 from wandb import util
-
-from ..interface.artifacts import ArtifactEntry, ArtifactManifest
+from wandb.sdk.interface.artifacts import ArtifactEntry, ArtifactManifest
 
 if TYPE_CHECKING:
     from wandb.proto import wandb_internal_pb2
@@ -30,7 +44,13 @@ def __call__(
             pass
 
 
-def _manifest_json_from_proto(manifest: "wandb_internal_pb2.ArtifactManifest") -> Dict:
+class ArtifactCommitError(Exception):
+    def __init__(self, exception: Exception) -> None:
+        super().__init__(f"Artifact commit failed: {exception}")
+        self.exception = exception
+
+
+def manifest_json_from_proto(manifest: "wandb_internal_pb2.ArtifactManifest") -> Dict:
     if manifest.version == 1:
         contents = {
             content.path: {
@@ -61,6 +81,34 @@ def _manifest_json_from_proto(manifest: "wandb_internal_pb2.ArtifactManifest") -
     }
 
 
+def _retry_conflicts(e: Exception) -> Union[bool, datetime.timedelta]:
+    if not isinstance(e, ArtifactCommitError):
+        return False
+    e = e.exception
+    if not isinstance(e, requests.exceptions.HTTPError):
+        return False
+    if e.response.status_code == 409:
+        return datetime.timedelta(minutes=2)
+    return False
+
+
+_T = TypeVar("_T")
+
+
+class _Future(Generic[_T]):
+    def __init__(self) -> None:
+        self._object: Optional[_T] = None
+        self._object_ready = threading.Event()
+
+    def set(self, obj: _T) -> None:
+        self._object = obj
+        self._object_ready.set()
+
+    def get(self) -> _T:
+        self._object_ready.wait()
+        return cast(_T, self._object)
+
+
 class ArtifactSaver:
     _server_artifact: Optional[Dict]  # TODO better define this dict
 
@@ -95,6 +143,55 @@ def save(
         incremental: bool = False,
         history_step: Optional[int] = None,
     ) -> Optional[Dict]:
+
+        # TODO: this retry is messy and should not exist!
+        # It's a hacky fix for https://wandb.atlassian.net/browse/WB-10888 :
+        # as of 2022-09-14, "409 Conflict" errors indicate that some uploaded files
+        # conflict with the artifact's most recent version's files. The right way to
+        # solve this would be for the SDK to ask the server what those files are,
+        # and replace their manifest entries with references to the previous version's
+        # entries; but, as a stopgap solution, we just retry the entire save operation,
+        # even though that might leave a bunch of orphaned files associated with the
+        # aborted artifact.
+        save_retry = wandb.sdk.lib.retry.Retry(
+            self._save,
+            retryable_exceptions=(ArtifactCommitError,),
+            check_retry_fn=_retry_conflicts,
+            num_retries=3,
+        )
+
+        return save_retry(
+            type,
+            name,
+            client_id,
+            sequence_client_id,
+            distributed_id,
+            finalize,
+            metadata,
+            description,
+            aliases,
+            labels,
+            use_after_commit,
+            incremental,
+            history_step,
+        )
+
+    def _save(
+        self,
+        type: str,
+        name: str,
+        client_id: str,
+        sequence_client_id: str,
+        distributed_id: Optional[str] = None,
+        finalize: bool = True,
+        metadata: Optional[Dict] = None,
+        description: Optional[str] = None,
+        aliases: Optional[Sequence[str]] = None,
+        labels: Optional[List[str]] = None,
+        use_after_commit: bool = False,
+        incremental: bool = False,
+        history_step: Optional[int] = None,
+    ) -> Optional[Dict]:
         aliases = aliases or []
         alias_specs = []
         for alias in aliases:
@@ -190,7 +287,7 @@ def save(
             ),
         )
 
-        commit_event = threading.Event()
+        commit_exc: _Future[Optional[Exception]] = _Future()
 
         def before_commit() -> None:
             self._resolve_client_id_manifest_references()
@@ -232,11 +329,12 @@ def before_commit() -> None:
                     extra_headers=extra_headers,
                 )
 
-        def on_commit() -> None:
-            if finalize and use_after_commit:
-                self._api.use_artifact(artifact_id)
+        def on_commit(exc: Optional[Exception]) -> None:
+            if exc is None:
+                if finalize and use_after_commit:
+                    self._api.use_artifact(artifact_id)
             step_prepare.shutdown()
-            commit_event.set()
+            commit_exc.set(exc)
 
         # This will queue the commit. It will only happen after all the file uploads are done
         self._file_pusher.commit_artifact(
@@ -248,8 +346,9 @@ def on_commit() -> None:
 
         # Block until all artifact files are uploaded and the
         # artifact is committed.
-        while not commit_event.is_set():
-            commit_event.wait()
+        exc = commit_exc.get()
+        if exc is not None:
+            raise ArtifactCommitError(exc) from exc
 
         return self._server_artifact
 

wandb/sdk/internal/sender.py

@@ -1256,7 +1256,7 @@ def _send_artifact(
         saver = artifacts.ArtifactSaver(
             api=self._api,
             digest=artifact.digest,
-            manifest_json=artifacts._manifest_json_from_proto(artifact.manifest),
+            manifest_json=artifacts.manifest_json_from_proto(artifact.manifest),
             file_pusher=self._pusher,
             is_user_created=artifact.user_created,
         )