feat(sagemaker): pass through async InvocationTimeoutSeconds

src/tabpfn_client/sagemaker/estimator.py

@@ -79,6 +79,7 @@ def __init__(
         s3_prefix: str = "async-io",
         async_poll_interval_s: float = 2.0,
         async_timeout_s: float = 60 * 60,
+        invocation_timeout_s: Optional[int] = None,
     ):
         self.endpoint_name = endpoint_name
         self.region_name = region_name
@@ -104,6 +105,11 @@ def __init__(
         self.s3_prefix = s3_prefix
         self.async_poll_interval_s = async_poll_interval_s
         self.async_timeout_s = async_timeout_s
+        # Server-side cap (seconds, 1-3600) the endpoint may spend on one async
+        # request before it is failed. Left as the service default when None;
+        # raise it for long-running requests. Distinct from `async_timeout_s`,
+        # which only bounds how long the client polls for the result.
+        self.invocation_timeout_s = invocation_timeout_s
 
     @property
     def _thinking_active(self) -> bool:
@@ -263,13 +269,16 @@ def _invoke_async(self, body_bytes: bytes) -> Dict[str, Any]:
             Body=body_bytes,
             ContentType="application/json",
         )
-        resp = self._runtime_client().invoke_endpoint_async(
-            EndpointName=self.endpoint_name,
-            InputLocation=f"s3://{self.s3_bucket}/{input_key}",
-            ContentType="application/json",
-            Accept="application/json",
-            InferenceId=inference_id,
-        )
+        invoke_kwargs: Dict[str, Any] = {
+            "EndpointName": self.endpoint_name,
+            "InputLocation": f"s3://{self.s3_bucket}/{input_key}",
+            "ContentType": "application/json",
+            "Accept": "application/json",
+            "InferenceId": inference_id,
+        }
+        if self.invocation_timeout_s is not None:
+            invoke_kwargs["InvocationTimeoutSeconds"] = int(self.invocation_timeout_s)
+        resp = self._runtime_client().invoke_endpoint_async(**invoke_kwargs)
         out_bucket, out_key = resp["OutputLocation"].removeprefix("s3://").split("/", 1)
         fail_location = resp.get("FailureLocation")
         fail_bucket, fail_key = (None, None)

tests/unit/conftest.py

@@ -0,0 +1,30 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+#  Licensed under the Apache License, Version 2.0
+"""Isolate process-global client state between unit tests.
+
+Several tests mutate the module-global ``options._opts`` (e.g. via
+``set_access_token`` / ``ServiceClient.authorize``) or the ``TABPFN_TOKEN``
+environment variable. Without isolation that state leaks into later tests and
+makes ordering-dependent assertions (such as ``test_reload_opts``) fail. Snapshot
+both before each test and restore afterwards so tests cannot pollute one another.
+"""
+
+import os
+
+import pytest
+
+import tabpfn_client.options as options
+
+
+@pytest.fixture(autouse=True)
+def _restore_global_client_state():  # pyright: ignore[reportUnusedFunction]
+    saved_opts = options._opts.model_copy(deep=True)
+    saved_token = os.environ.get("TABPFN_TOKEN")
+    try:
+        yield
+    finally:
+        options._opts = saved_opts
+        if saved_token is None:
+            os.environ.pop("TABPFN_TOKEN", None)
+        else:
+            os.environ["TABPFN_TOKEN"] = saved_token

tests/unit/test_sagemaker_invocation_timeout.py

@@ -0,0 +1,58 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+#  Licensed under the Apache License, Version 2.0
+"""Unit tests for the async `InvocationTimeoutSeconds` pass-through on the
+SageMaker estimators. boto3 is mocked, so these run without the extra."""
+
+import json
+import unittest
+from io import BytesIO
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+
+from tabpfn_client.sagemaker import TabPFNClassifier
+
+
+def _fake_s3_client() -> MagicMock:
+    s3 = MagicMock()
+    # `_invoke_async` reads `s3.exceptions.NoSuchKey` as an exception class.
+    s3.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {})
+    s3.get_object.return_value = {
+        "Body": BytesIO(json.dumps({"prediction": [[0, 1]], "metadata": {}}).encode())
+    }
+    return s3
+
+
+class TestSagemakerInvocationTimeout(unittest.TestCase):
+    def _invoke_async_kwargs(self, invocation_timeout_s):
+        clf = TabPFNClassifier(
+            endpoint_name="ep",
+            region_name="us-east-1",
+            use_async=True,
+            s3_bucket="bucket",
+            async_poll_interval_s=0.0,
+            invocation_timeout_s=invocation_timeout_s,
+        )
+        runtime = MagicMock()
+        runtime.invoke_endpoint_async.return_value = {
+            "OutputLocation": "s3://bucket/out.json"
+        }
+        with (
+            patch.object(clf, "_runtime_client", return_value=runtime),
+            patch.object(clf, "_s3_client", return_value=_fake_s3_client()),
+        ):
+            clf.fit(np.zeros((3, 2)), np.array([0, 1, 0]))
+            clf.predict(np.zeros((2, 2)))
+        return runtime.invoke_endpoint_async.call_args.kwargs
+
+    def test_timeout_passed_through(self):
+        kwargs = self._invoke_async_kwargs(1234)
+        self.assertEqual(kwargs["InvocationTimeoutSeconds"], 1234)
+
+    def test_timeout_omitted_by_default(self):
+        kwargs = self._invoke_async_kwargs(None)
+        self.assertNotIn("InvocationTimeoutSeconds", kwargs)
+
+
+if __name__ == "__main__":
+    unittest.main()