[Misc] Batch: Extend batch session retry handling with capped backoff to support more retry times.

python/aibrix/aibrix/batch/job_driver/runtime/base.py

@@ -123,6 +123,10 @@ def __deepcopy__(self, memo):
             current_worker_id=self.current_worker_id,
             retry_after_s=self.retry_after_s,
         )
+        # Keep any retry suffix appended after __init__; reconstructing via
+        # original fields would otherwise regenerate the base message.
+        new_copy.message = self.message
+        new_copy.args = self.args
         memo[id(self)] = new_copy
         return new_copy
 
@@ -152,6 +156,10 @@ def __deepcopy__(self, memo):
             runtime_key=self.runtime_key,
             retry_after_s=self.retry_after_s,
         )
+        # Keep any retry suffix appended after __init__; reconstructing via
+        # original fields would otherwise regenerate the base message.
+        new_copy.message = self.message
+        new_copy.args = self.args
         memo[id(self)] = new_copy
         return new_copy
 
@@ -266,8 +274,9 @@ class RuntimeBase:
     """
 
     provisions: bool = False
-    session_retry_attempts: int = 5
+    session_retry_attempts: int = envs.BATCH_SESSION_RETRY_ATTEMPTS
     session_retry_base_delay_s: float = 2.0
+    session_retry_max_delay_s: float = 60.0
     session_liveness_check_interval_s: float = 30.0
     session_liveness_failure_threshold: Optional[int] = None
 
@@ -1003,17 +1012,46 @@ def _should_teardown_failed_wait_ready(self, exc: Exception) -> bool:
         # meaningful left to tear down; retry by provisioning a fresh resource.
         return not self._is_not_found_error(exc)
 
-    async def _sleep_before_session_retry(self, attempt: int) -> None:
-        await asyncio.sleep(self.session_retry_base_delay_s * (2**attempt))
+    def _get_session_retry_delay_s(self, attempt: int) -> float:
+        # set a hard-coded cap for avoiding OverflowError
+        safe_attempt = min(attempt, 10)
+        return min(
+            self.session_retry_base_delay_s * (2**safe_attempt),
+            self.session_retry_max_delay_s,
+        )
+
+    def _annotate_exhausted_session_error(
+        self, exc: Exception, *, retries_completed: int
+    ) -> Exception:
+        if retries_completed <= 0:
+            return exc
+        suffix = (
+            f" [session_retries={retries_completed}, "
+            f"session_attempts={retries_completed + 1}]"
+        )
+        if isinstance(exc, BatchJobError):
+            if suffix not in exc.message:
+                exc.message = f"{exc.message}{suffix}"
+                exc.args = (exc.message,)
+            return exc
+        if exc.args and isinstance(exc.args[0], str):
+            if suffix not in exc.args[0]:
+                exc.args = (f"{exc.args[0]}{suffix}", *exc.args[1:])
+        else:
+            exc.args = (*exc.args, suffix)
+        return exc
 
     async def _sleep_before_session_error_retry(
         self, attempt: int, exc: Exception
     ) -> None:
-        delay = self.session_retry_base_delay_s * (2**attempt)
+        delay = self._get_session_retry_delay_s(attempt)
         if isinstance(
             exc, (RuntimeOwnershipConflictError, RuntimeDeleteInProgressError)
         ):
-            delay = max(delay, exc.retry_after_s)
+            delay = min(
+                max(delay, exc.retry_after_s),
+                self.session_retry_max_delay_s,
+            )
         await asyncio.sleep(delay)
 
     def _should_run_session_liveness_checks(self, handle: Any) -> bool:
@@ -1244,6 +1282,12 @@ async def session(
                     elif not should_retry:
                         handle = None
                     if not should_retry:
+                        self._annotate_exhausted_session_error(
+                            exc,
+                            retries_completed=attempt,
+                        )
+                        # Re-raise the active exception so its original
+                        # traceback is preserved after in-place annotation.
                         raise
                     handle = None
                     runtimeRef = None

python/aibrix/aibrix/envs.py

@@ -119,6 +119,9 @@ def _parse_int_or_none(value: Optional[str]) -> Optional[int]:
 BATCH_ERROR_INJECTION_ENABLED = _is_true(
     os.getenv("AIBRIX_BATCH_ERROR_INJECTION_ENABLED", "0")
 )
+BATCH_SESSION_RETRY_ATTEMPTS = int(
+    os.getenv("AIBRIX_BATCH_SESSION_RETRY_ATTEMPTS", "5")
+)
 BATCH_SESSION_LIVENESS_FAILURE_THRESHOLD = int(
     os.getenv("AIBRIX_BATCH_SESSION_LIVENESS_FAILURE_THRESHOLD", "3")
 )

python/aibrix/tests/batch/job_driver/runtime/test_runtime.py

@@ -14,6 +14,7 @@
 """Unit tests for the Runtime seam + registry (job lifecycle axis A)."""
 
 import asyncio
+import copy
 import inspect
 from datetime import datetime, timezone
 from types import SimpleNamespace
@@ -955,10 +956,18 @@ async def _persist_runtime_ref(
 
 
 @pytest.mark.asyncio
-async def test_runtime_base_session_raises_if_execution_tracking_not_bound():
+async def test_runtime_base_session_raises_if_execution_tracking_not_bound(monkeypatch):
     job = _make_test_job(job_id="job-1")
+    sleeps: list[float] = []
+
+    async def _fake_sleep(delay: float) -> None:
+        sleeps.append(delay)
+
+    monkeypatch.setattr(runtime_base_mod.asyncio, "sleep", _fake_sleep)
 
     class _PersistingRuntime(_R):
+        session_retry_attempts = 1
+
         def _build_runtime_ref(self, job):
             existing = self._load_runtime_ref(job)
             now = datetime.now(timezone.utc)
@@ -974,11 +983,17 @@ def _build_runtime_ref(self, job):
     runtime = _PersistingRuntime(provision=lambda job, job_id: "handle")
     with pytest.raises(
         RuntimeError,
-        match="Execution tracking is not configured for session\\(\\)",
-    ):
+        match=(
+            "Execution tracking is not configured for session\\(\\).*"
+            "session_retries=1.*session_attempts=2"
+        ),
+    ) as exc_info:
         async with runtime.session(job=job, job_id="job-1"):
             pytest.fail("session should fail before entering body")
 
+    assert "session_retries=1" in str(exc_info.value)
+    assert sleeps == [2.0]
+
 
 @pytest.mark.asyncio
 async def test_persist_runtime_ref_updates_execution_only():
@@ -1419,6 +1434,86 @@ def test_runtime_target_enum_values_are_registered():
     assert not missing, f"RuntimeTarget values not registered: {missing}"
 
 
+def test_runtime_base_uses_env_session_retry_attempts():
+    assert (
+        RuntimeBase.session_retry_attempts
+        == runtime_base_mod.envs.BATCH_SESSION_RETRY_ATTEMPTS
+    )
+
+
+@pytest.mark.asyncio
+async def test_session_error_retry_delay_is_capped(monkeypatch):
+    delays: list[float] = []
+
+    async def _capture_sleep(delay: float) -> None:
+        delays.append(delay)
+
+    monkeypatch.setattr(runtime_base_mod.asyncio, "sleep", _capture_sleep)
+    runtime = _R()
+
+    await runtime._sleep_before_session_error_retry(10, Exception("boom"))
+
+    assert delays == [60.0]
+
+
+@pytest.mark.asyncio
+async def test_session_error_retry_after_is_capped(monkeypatch):
+    delays: list[float] = []
+
+    async def _capture_sleep(delay: float) -> None:
+        delays.append(delay)
+
+    monkeypatch.setattr(runtime_base_mod.asyncio, "sleep", _capture_sleep)
+    runtime = _R()
+    exc = runtime_base_mod.RuntimeDeleteInProgressError(
+        job_id="job-1",
+        runtime_key="base",
+        retry_after_s=120.0,
+    )
+
+    await runtime._sleep_before_session_error_retry(1, exc)
+
+    assert delays == [60.0]
+
+
+@pytest.mark.parametrize(
+    ("exc", "expected_message_bits"),
+    [
+        (
+            runtime_base_mod.RuntimeOwnershipConflictError(
+                job_id="job-1",
+                runtime_key="base",
+                owner_worker_id="worker-a",
+                current_worker_id="worker-b",
+                retry_after_s=2.5,
+            ),
+            ["owner_worker_id=worker-a", "current_worker_id=worker-b"],
+        ),
+        (
+            runtime_base_mod.RuntimeDeleteInProgressError(
+                job_id="job-1",
+                runtime_key="base",
+                retry_after_s=2.5,
+            ),
+            ["runtime=base", "retry_after_s=2.50"],
+        ),
+    ],
+)
+def test_runtime_session_error_deepcopy_preserves_annotated_suffix(
+    exc, expected_message_bits
+):
+    runtime = _R()
+
+    annotated = runtime._annotate_exhausted_session_error(exc, retries_completed=2)
+    copied = copy.deepcopy(annotated)
+
+    assert copied.message == annotated.message
+    assert copied.args == annotated.args
+    assert "[session_retries=2, session_attempts=3]" in copied.message
+    for expected_bit in expected_message_bits:
+        assert expected_bit in copied.message
+
+
 def test_registry_create_and_unknown_key():
     assert isinstance(create_runtime("noop"), NoopRuntime)
     src = object()