From a51e5eec4df0b80dee13fcd89044be1b42ed5c40 Mon Sep 17 00:00:00 2001 From: Joan Fontanals Date: Tue, 5 Nov 2024 09:58:25 +0100 Subject: [PATCH] fix: fix req handling sagemaker (#6218) --- .../serve/runtimes/worker/request_handling.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/jina/serve/runtimes/worker/request_handling.py b/jina/serve/runtimes/worker/request_handling.py index a813e60bddb95..450690d33bcdc 100644 --- a/jina/serve/runtimes/worker/request_handling.py +++ b/jina/serve/runtimes/worker/request_handling.py @@ -292,15 +292,18 @@ def _init_batchqueue_dict(self): func.fn.__name__: [] for func in self._executor.requests.values() } for endpoint, func in self._executor.requests.items(): - func_endpoints[func.fn.__name__].append(endpoint) + if func.fn.__name__ in func_endpoints: + # For SageMaker, not all endpoints are there + func_endpoints[func.fn.__name__].append(endpoint) for func_name, dbatch_config in dbatch_functions: - for endpoint in func_endpoints[func_name]: - if endpoint not in self._batchqueue_config: - self._batchqueue_config[endpoint] = dbatch_config - else: - # we need to eventually copy the `custom_metric` - if dbatch_config.get('custom_metric', None) is not None: - self._batchqueue_config[endpoint]['custom_metric'] = dbatch_config.get('custom_metric') + if func_name in func_endpoints: # For SageMaker, not all endpoints are there + for endpoint in func_endpoints[func_name]: + if endpoint not in self._batchqueue_config: + self._batchqueue_config[endpoint] = dbatch_config + else: + # we need to eventually copy the `custom_metric` + if dbatch_config.get('custom_metric', None) is not None: + self._batchqueue_config[endpoint]['custom_metric'] = dbatch_config.get('custom_metric') keys_to_remove = [] for k, batch_config in self._batchqueue_config.items():