feat(artifacts): add optional timeout parameter to artifacts wait()

tests/unit_tests/test_wandb_artifacts_full.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 import wandb
+from wandb.sdk.wandb_run import WaitTimeoutError
 
 sm = wandb.wandb_sdk.internal.sender.SendManager
 
@@ -193,3 +194,29 @@ def make_table():
     artifact2.add(t1, "t2")
     assert artifact2.manifest.entries["t2.table.json"].ref is not None
     run.finish()
+
+
+def test_artifact_wait_success(wandb_init):
+    # Test artifact wait() timeout parameter
+    timeout = 2
+    leeway = 0.50
+    run = wandb_init()
+    artifact = wandb.Artifact("art", type="dataset")
+    start_timestamp = time.time()
+    run.log_artifact(artifact).wait(wait_timeout_secs=timeout)
+    elapsed_time = time.time() - start_timestamp
+    assert elapsed_time < timeout * (1 + leeway)
+    run.finish()
+
+
+@pytest.mark.parametrize("timeout", [0, 1e-1])
+def test_artifact_wait_failure(wandb_init, timeout):
+    # Test to expect WaitTimeoutError when wait timeout is reached and large image
+    # wasn't uploaded yet
+    large_image = wandb.Image(np.zeros((1000, 1000)))
+    run = wandb_init()
+    with pytest.raises(WaitTimeoutError):
+        artifact = wandb.Artifact("art", type="image")
+        artifact.add(large_image, "image")
+        run.log_artifact(artifact).wait(wait_timeout_secs=timeout)
+    run.finish()

tests/unit_tests/test_wandb_run.py

@@ -2,6 +2,7 @@
 import pickle
 import platform
 import sys
+import time
 from unittest import mock
 
 import numpy as np

wandb/sdk/wandb_artifacts.py

@@ -647,9 +647,9 @@ def delete(self) -> None:
             "Cannot call delete on an artifact before it has been logged or in offline mode"
         )
 
-    def wait(self) -> ArtifactInterface:
+    def wait(self, wait_timeout_secs: int = None) -> ArtifactInterface:
         if self._logged_artifact:
-            return self._logged_artifact.wait()
+            return self._logged_artifact.wait(wait_timeout_secs)
 
         raise ValueError(
             "Cannot call wait on an artifact before it has been logged or in offline mode"

wandb/sdk/wandb_run.py

@@ -135,6 +135,12 @@ class TeardownHook(NamedTuple):
     stage: TeardownStage
 
 
+class WaitTimeoutError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message
+
+
 class RunStatusChecker:
     """Periodically polls the background process for relevant updates.
 
@@ -3575,9 +3581,14 @@ def __getattr__(self, item: str) -> Any:
         self._assert_instance()
         return getattr(self._instance, item)
 
-    def wait(self) -> ArtifactInterface:
+    def wait(self, wait_timeout_secs: int = None) -> ArtifactInterface:
         if not self._instance:
-            resp = self._future.get().response.log_artifact_response
+            future_get = self._future.get(wait_timeout_secs)
+            if not future_get:
+                raise WaitTimeoutError(
+                    "Artifact wait timed out, failed to fetch Artifact response"
+                )
+            resp = future_get.response.log_artifact_response
             if resp.error_message:
                 raise ValueError(resp.error_message)
             self._instance = public.Artifact.from_id(resp.artifact_id, self._api.client)