Prevent duplicate prediction submissions

Author: chengzeyi

README.md

@@ -93,8 +93,8 @@ from wavespeed import Client
 
 client = Client(
     api_key="your-api-key",
-    max_retries=0,            # Task-level retries (default: 0)
-    max_connection_retries=5, # HTTP connection retries (default: 5)
+    max_retries=0,            # Replacement task attempts (default: 0)
+    max_connection_retries=5, # Result-query GET retries; POST is never retried
     retry_interval=1.0,       # Base delay between retries in seconds (default: 1.0)
 )
 ```

src/wavespeed/api/__init__.py

@@ -47,7 +47,7 @@ def run(
     """Run a model and wait for the output.
 
     Args:
-        model: Model identifier (e.g., "wavespeed-ai/flux-dev").
+        model: Model identifier (e.g., "wavespeed-ai/z-image/turbo").
         input: Input parameters for the model.
         timeout: Maximum time to wait for completion (None = no timeout).
         poll_interval: Interval between status checks in seconds.
@@ -78,11 +78,12 @@ def run(
             enable_sync_mode=True
         )
 
-        # With retry
+        # Keep replacement task attempts disabled unless your workload accepts
+        # a new task after a confirmed terminal failure.
         output = wavespeed.run(
             "wavespeed-ai/z-image/turbo",
             {"prompt": "A cat"},
-            max_retries=3
+            max_retries=0
         )
     """
     return _get_default_client().run(

src/wavespeed/api/client.py

@@ -13,6 +13,10 @@
 _RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}
 
 
+class _SubmissionError(RuntimeError):
+    """A task submission failed and must not be retried automatically."""
+
+
 class Client:
     """WaveSpeed API client.
 
@@ -21,7 +25,7 @@ class Client:
         base_url: Base URL for the API. If not provided, uses wavespeed.config.api.base_url.
         connection_timeout: Timeout for HTTP requests in seconds.
         max_retries: Maximum number of retries for the entire operation.
-        max_connection_retries: Maximum retries for individual HTTP requests.
+        max_connection_retries: Maximum retries for result-query GET requests.
         retry_interval: Base interval between retries in seconds.
 
     Example:
@@ -31,8 +35,9 @@ class Client:
         # With sync mode (best-effort single request, waits for result)
         output = client.run("wavespeed-ai/z-image/turbo", {"prompt": "Cat"}, enable_sync_mode=True)
 
-        # With retry
-        output = client.run("wavespeed-ai/z-image/turbo", {"prompt": "Cat"}, max_retries=3)
+        # Task-level replacement attempts are opt-in; submission POSTs are
+        # always sent at most once.
+        output = client.run("wavespeed-ai/z-image/turbo", {"prompt": "Cat"}, max_retries=1)
     """
 
     def __init__(
@@ -121,67 +126,34 @@ def _submit(
         )
         timeouts = (connect_timeout, request_timeout)
 
-        last_error: Exception | None = None
-
-        for retry in range(self.max_connection_retries + 1):
-            try:
-                response = requests.post(
-                    url, json=body, headers=self._get_headers(), timeout=timeouts
-                )
-
-                if response.status_code != 200:
-                    # Retry on transient server errors (5xx) and rate limiting (429)
-                    if self._is_retryable_status(response.status_code):
-                        last_error = RuntimeError(
-                            f"Failed to submit prediction: HTTP {response.status_code}: "
-                            f"{response.text}"
-                        )
-                        if retry < self.max_connection_retries:
-                            delay = self.retry_interval * (retry + 1)
-                            print(
-                                f"Server error (HTTP {response.status_code}) on attempt "
-                                f"{retry + 1}/{self.max_connection_retries + 1}, "
-                                f"retrying in {delay} seconds..."
-                            )
-                            time.sleep(delay)
-                            continue
-                        raise last_error
-
-                    # Non-retryable HTTP errors (4xx etc.) fail immediately
-                    raise RuntimeError(
-                        f"Failed to submit prediction: HTTP {response.status_code}: "
-                        f"{response.text}"
-                    )
-
-                result = response.json()
+        try:
+            response = requests.post(
+                url, json=body, headers=self._get_headers(), timeout=timeouts
+            )
+        except (
+            requests.exceptions.ConnectionError,
+            requests.exceptions.Timeout,
+        ) as e:
+            raise _SubmissionError(
+                "Prediction submission did not return a response. The task may already "
+                "have been created, so the SDK will not retry the POST automatically."
+            ) from e
 
-                if enable_sync_mode:
-                    return None, result
+        if response.status_code != 200:
+            raise _SubmissionError(
+                f"Failed to submit prediction: HTTP {response.status_code}: {response.text}"
+            )
 
-                request_id = result.get("data", {}).get("id")
-                if not request_id:
-                    raise RuntimeError(f"No request ID in response: {result}")
+        result = response.json()
 
-                return request_id, None
+        if enable_sync_mode:
+            return None, result
 
-            except (
-                requests.exceptions.ConnectionError,
-                requests.exceptions.Timeout,
-            ) as e:
-                last_error = e
-                print(
-                    f"Connection error on attempt {retry + 1}/{self.max_connection_retries + 1}:"
-                )
-                traceback.print_exc()
+        request_id = result.get("data", {}).get("id")
+        if not request_id:
+            raise _SubmissionError(f"No request ID in response: {result}")
 
-                if retry < self.max_connection_retries:
-                    delay = self.retry_interval * (retry + 1)
-                    print(f"Retrying in {delay} seconds...")
-                    time.sleep(delay)
-                else:
-                    raise RuntimeError(
-                        f"Failed to submit prediction after {self.max_connection_retries + 1} attempts"
-                    ) from e
+        return request_id, None
 
     def _get_result(
         self, request_id: str, timeout: float | None = None
@@ -315,7 +287,12 @@ def _is_retryable_error(self, error: Exception) -> bool:
         Returns:
             True if the error is retryable.
         """
-        # Always retry timeout and connection errors
+        # Submission errors are ambiguous: the server may already have created
+        # the task, so never turn them into another POST automatically.
+        if isinstance(error, _SubmissionError):
+            return False
+
+        # Retry timeout and connection errors from result-query GETs.
         if isinstance(
             error,
             (
@@ -366,7 +343,7 @@ def run(
         """Run a model and wait for the output.
 
         Args:
-            model: Model identifier (e.g., "wavespeed-ai/flux-dev").
+            model: Model identifier (e.g., "wavespeed-ai/z-image/turbo").
             input: Input parameters for the model.
             timeout: Maximum time to wait for completion (None = no timeout).
             poll_interval: Interval between status checks in seconds.

src/wavespeed/config.py

@@ -31,7 +31,8 @@ class api:
     # Maximum number of retries for the entire operation (task-level retries)
     max_retries: int = 0
 
-    # Maximum number of retries for individual HTTP requests (connection errors, timeouts)
+    # Maximum retries for idempotent result-query GET requests. Submission
+    # POSTs are intentionally sent at most once.
     max_connection_retries: int = 5
 
     # Base interval between retries in seconds (actual delay = retry_interval * attempt)

tests/test_api.py

@@ -6,6 +6,7 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+import requests
 import wavespeed
 from wavespeed.api import Client
 
@@ -82,6 +83,33 @@ def test_submit_failure(self, mock_post):
         with self.assertRaises(RuntimeError) as ctx:
             client._submit("wavespeed-ai/z-image/turbo", {"prompt": "test"})
         self.assertIn("HTTP 500", str(ctx.exception))
+        mock_post.assert_called_once()
+
+    @patch("wavespeed.api.client.requests.post")
+    def test_submit_connection_error_is_not_retried(self, mock_post):
+        """An ambiguous disconnect must not create a duplicate task."""
+        mock_post.side_effect = requests.exceptions.ConnectionError("disconnected")
+
+        client = Client(api_key="test-key", max_connection_retries=5)
+        with self.assertRaises(RuntimeError) as ctx:
+            client._submit("wavespeed-ai/z-image/turbo", {"prompt": "test"})
+
+        self.assertIn("will not retry the POST", str(ctx.exception))
+        mock_post.assert_called_once()
+
+    @patch("wavespeed.api.client.requests.post")
+    def test_task_retries_do_not_repeat_failed_submission(self, mock_post):
+        """Task retry configuration must not override POST safety."""
+        mock_response = MagicMock()
+        mock_response.status_code = 500
+        mock_response.text = "Internal Server Error"
+        mock_post.return_value = mock_response
+
+        client = Client(api_key="test-key", max_retries=3)
+        with self.assertRaises(RuntimeError):
+            client.run("wavespeed-ai/z-image/turbo", {"prompt": "test"})
+
+        mock_post.assert_called_once()
 
     @patch("wavespeed.api.client.requests.get")
     def test_get_result_success(self, mock_get):