diff --git a/src/tabpfn_client/sagemaker/estimator.py b/src/tabpfn_client/sagemaker/estimator.py index 80421f9..4221a56 100644 --- a/src/tabpfn_client/sagemaker/estimator.py +++ b/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) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..b4a10ea --- /dev/null +++ b/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 diff --git a/tests/unit/test_sagemaker_invocation_timeout.py b/tests/unit/test_sagemaker_invocation_timeout.py new file mode 100644 index 0000000..63b27cc --- /dev/null +++ b/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()