increase timeouts to more than 900s

Author: MalteKurz

doubleml_serverless/double_ml_aws_lambda.py

@@ -1,6 +1,7 @@
 import pandas as pd
 import asyncio
-import aiobotocore
+from aiobotocore.session import get_session
+from aiobotocore.config import AioConfig
 import json
 import time
 
@@ -98,15 +99,16 @@ def invoke_lambdas(self, payloads, smpls, params_names, n_obs, n_rep, n_jobs_cv)
         return preds
 
     async def __invoke_aws_lambdas(self, payloads):
-        session = aiobotocore.get_session()
+        session = get_session()
         tasks = []
         for this_payload in payloads:
             tasks.append(self.__invoke_single_aws_lambda(session, this_payload))
         results = await asyncio.gather(*tasks)
         return results
 
     async def __invoke_single_aws_lambda(self, session, payload):
-        async with session.create_client('lambda', region_name=self.aws_region) as lambda_client:
+        config = AioConfig(connect_timeout=1000, read_timeout=1000)
+        async with session.create_client('lambda', region_name=self.aws_region, config=config) as lambda_client:
             # print(f'Invoking {payload["learner"]} {payload["i_rep"]} {payload["i_fold"]}')
             response = await lambda_client.invoke(
                 FunctionName=self.lambda_function_name,