feat(sagemaker): pass through async InvocationTimeoutSeconds#333
Conversation
Add an optional invocation_timeout_s to the SageMaker estimators. When set it is forwarded to invoke_endpoint_async as InvocationTimeoutSeconds, letting callers raise the server-side per-request cap (1-3600s) for long-running async requests. Defaults to None (service default), so behaviour is unchanged when unset. Distinct from async_timeout_s, which only bounds client-side polling.
test_client's authorize() leaves TABPFN_TOKEN on the module-global options, so test_reload_opts (which asserts it starts as None) fails depending on order. Add an autouse fixture that snapshots and restores options._opts and the TABPFN_TOKEN env var around each test. Pre-existing on main; unrelated to the InvocationTimeoutSeconds change but required for a green suite.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for configuring a server-side invocation timeout (invocation_timeout_s) on SageMaker estimators, passing it to the invoke_endpoint_async call. It also adds a test fixture to isolate global client state between unit tests and includes new unit tests for the timeout parameter. The review feedback suggests validating that invocation_timeout_s is an integer between 1 and 3600, removing a redundant int() cast once validation is in place, and adding unit tests to verify that invalid timeout values correctly raise a ValueError.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
2 participants
The SageMaker async estimators previously invoked
invoke_endpoint_asyncwith noInvocationTimeoutSeconds, so every async request rode the service default and long-running requests could be failed by the platform before completing with no way to extend the window from the client. This adds an optionalinvocation_timeout_sconstructor argument that, when set, is forwarded asInvocationTimeoutSeconds(the server-side 1–3600s per-request cap); it defaults toNoneso behaviour is unchanged when unset, and it is deliberately distinct fromasync_timeout_s, which only bounds how long the client polls S3 for the result. A mocked unit test covers both the pass-through and the default-omitted cases and needs no real boto3.Part of ENG-841.