Spaces:
Running
Running
File size: 1,347 Bytes
224d30c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """
Shared Bedrock timeout monkey-patch.
Bedrock Converse API can take >60s for long responses (4096 tokens).
The default botocore read_timeout is 60s which causes ReadTimeoutError.
This patch increases it before inspect-ai creates the Bedrock client.
"""
def patch_bedrock_timeout(read_timeout: int = 600) -> None:
"""
Patch the inspect-ai Bedrock provider to use a higher read timeout.
Parameters
----------
read_timeout : int
Read timeout in seconds for Bedrock API calls. Default 600s (10 min).
"""
from inspect_ai.model._providers import bedrock as bedrock_module
_original_generate = bedrock_module.BedrockAPI.generate
async def _patched_generate(self, *args, **kwargs):
from botocore.config import Config
_original_config_init = Config.__init__
def _config_init_with_timeout(config_self, *c_args, **c_kwargs):
c_kwargs.setdefault("read_timeout", read_timeout)
c_kwargs.setdefault("connect_timeout", 60)
return _original_config_init(config_self, *c_args, **c_kwargs)
Config.__init__ = _config_init_with_timeout
try:
return await _original_generate(self, *args, **kwargs)
finally:
Config.__init__ = _original_config_init
bedrock_module.BedrockAPI.generate = _patched_generate
|