|
18 | 18 |
|
19 | 19 | import grpc
|
20 | 20 |
|
| 21 | +from . import __version__ |
21 | 22 | from . import bindings
|
22 | 23 | from . import constants
|
23 | 24 | from . import functions
|
24 | 25 | from . import loader
|
25 | 26 | from . import protos
|
26 |
| -from .constants import (CONSOLE_LOG_PREFIX, PYTHON_THREADPOOL_THREAD_COUNT, |
| 27 | +from .constants import (PYTHON_THREADPOOL_THREAD_COUNT, |
27 | 28 | PYTHON_THREADPOOL_THREAD_COUNT_DEFAULT,
|
28 | 29 | PYTHON_THREADPOOL_THREAD_COUNT_MAX,
|
29 | 30 | PYTHON_THREADPOOL_THREAD_COUNT_MIN)
|
30 | 31 | from .logging import disable_console_logging, enable_console_logging
|
31 |
| -from .logging import error_logger, is_system_log_category, logger |
| 32 | +from .logging import (logger, error_logger, is_system_log_category, |
| 33 | + CONSOLE_LOG_PREFIX) |
| 34 | +from .extension import ExtensionManager |
32 | 35 | from .utils.common import get_app_setting
|
33 | 36 | from .utils.tracing import marshall_exception_trace
|
34 | 37 | from .utils.dependency import DependencyManager
|
@@ -255,8 +258,9 @@ async def _dispatch_grpc_request(self, request):
|
255 | 258 | self._grpc_resp_queue.put_nowait(resp)
|
256 | 259 |
|
257 | 260 | async def _handle__worker_init_request(self, req):
|
258 |
| - logger.info('Received WorkerInitRequest, request ID %s', |
259 |
| - self.request_id) |
| 261 | + logger.info('Received WorkerInitRequest, ' |
| 262 | + 'python version %s, worker version %s, request ID %s', |
| 263 | + sys.version, __version__, self.request_id) |
260 | 264 |
|
261 | 265 | capabilities = {
|
262 | 266 | constants.RAW_HTTP_BODY_BYTES: _TRUE,
|
@@ -304,6 +308,11 @@ async def _handle__function_load_request(self, req):
|
304 | 308 | self._functions.add_function(
|
305 | 309 | function_id, func, func_request.metadata)
|
306 | 310 |
|
| 311 | + ExtensionManager.function_load_extension( |
| 312 | + function_name, |
| 313 | + func_request.metadata.directory |
| 314 | + ) |
| 315 | + |
307 | 316 | logger.info('Successfully processed FunctionLoadRequest, '
|
308 | 317 | f'request ID: {self.request_id}, '
|
309 | 318 | f'function ID: {function_id},'
|
@@ -373,20 +382,24 @@ async def _handle__invocation_request(self, req):
|
373 | 382 | pytype=pb_type_info.pytype,
|
374 | 383 | shmem_mgr=self._shmem_mgr)
|
375 | 384 |
|
| 385 | + fi_context = bindings.Context( |
| 386 | + fi.name, fi.directory, invocation_id, trace_context) |
376 | 387 | if fi.requires_context:
|
377 |
| - args['context'] = bindings.Context( |
378 |
| - fi.name, fi.directory, invocation_id, trace_context) |
| 388 | + args['context'] = fi_context |
379 | 389 |
|
380 | 390 | if fi.output_types:
|
381 | 391 | for name in fi.output_types:
|
382 | 392 | args[name] = bindings.Out()
|
383 | 393 |
|
384 | 394 | if fi.is_async:
|
385 |
| - call_result = await fi.func(**args) |
| 395 | + call_result = await self._run_async_func( |
| 396 | + fi_context, fi.func, args |
| 397 | + ) |
386 | 398 | else:
|
387 | 399 | call_result = await self._loop.run_in_executor(
|
388 | 400 | self._sync_call_tp,
|
389 |
| - self.__run_sync_func, invocation_id, fi.func, args) |
| 401 | + self._run_sync_func, |
| 402 | + invocation_id, fi_context, fi.func, args) |
390 | 403 | if call_result is not None and not fi.has_return:
|
391 | 404 | raise RuntimeError(f'function {fi.name!r} without a $return '
|
392 | 405 | 'binding returned a non-None value')
|
@@ -582,15 +595,21 @@ def _create_sync_call_tp(
|
582 | 595 | max_workers=max_worker
|
583 | 596 | )
|
584 | 597 |
|
585 |
| - def __run_sync_func(self, invocation_id, func, params): |
| 598 | + def _run_sync_func(self, invocation_id, context, func, params): |
586 | 599 | # This helper exists because we need to access the current
|
587 | 600 | # invocation_id from ThreadPoolExecutor's threads.
|
588 | 601 | _invocation_id_local.v = invocation_id
|
589 | 602 | try:
|
590 |
| - return func(**params) |
| 603 | + return ExtensionManager.get_sync_invocation_wrapper(context, |
| 604 | + func)(params) |
591 | 605 | finally:
|
592 | 606 | _invocation_id_local.v = None
|
593 | 607 |
|
| 608 | + async def _run_async_func(self, context, func, params): |
| 609 | + return await ExtensionManager.get_async_invocation_wrapper( |
| 610 | + context, func, params |
| 611 | + ) |
| 612 | + |
594 | 613 | def __poll_grpc(self):
|
595 | 614 | options = []
|
596 | 615 | if self._grpc_max_msg_len:
|
|
0 commit comments