ivcap-lambda 0.7.22__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,17 @@
1
+ #
2
+ # Copyright (c) 2023 Commonwealth Scientific and Industrial Research Organisation (CSIRO). All rights reserved.
3
+ # Use of this source code is governed by a BSD-style license that can be
4
+ # found in the LICENSE file. See the AUTHORS file for names of contributors.
5
+ #
6
+ """A library supporting the development of tools for agents to be deployed on IVCAP"""
7
+
8
+ from .version import __version__
9
+
10
+ from .server import start_tool_server
11
+ from .builder import add_tool_api_route, ToolOptions
12
+ from .executor import ExecutionContext, get_event_reporter, get_job_id
13
+ from .utils import get_public_url_prefix
14
+ from .secret import SecretMgrClient
15
+ from .decorators import ivcap_lambda
16
+ from .decorators import ivcap_lambda as ivcap_ai_tool # backward-compat alias
17
+ from .logger import logging_init
@@ -0,0 +1,260 @@
1
+ #
2
+ # Copyright (c) 2023 Commonwealth Scientific and Industrial Research Organisation (CSIRO). All rights reserved.
3
+ # Use of this source code is governed by a BSD-style license that can be
4
+ # found in the LICENSE file. See the AUTHORS file for names of contributors.
5
+ #
6
+ import asyncio
7
+ from dataclasses import dataclass
8
+ import json
9
+ from fastapi import FastAPI, Response, status, Request
10
+ from pydantic import BaseModel, Field
11
+ from typing import Any, Dict, List, Optional, Callable, Tuple, Type, TypeVar
12
+ from uuid6 import uuid6
13
+
14
+ from ivcap_service import getLogger, get_function_return_type, get_input_type, create_tool_definition
15
+ from ivcap_service import IvcapResult, ToolDefinition, ExecutionError
16
+
17
+ from .executor import ExecutionContext, Executor, ExecutorOpts
18
+ from .utils import get_title_from_path, get_public_url_prefix
19
+
20
+
21
+ class ErrorModel(BaseModel):
22
+ message: str
23
+ code: int
24
+
25
+ class ExecutionErrorModel(BaseModel):
26
+ jschema: str = Field("urn:ivcap:schema.ai-tool.error.1", alias="$schema")
27
+ message: str
28
+ traceback: str
29
+
30
+ JOB_URN_PREFIX = "urn:ivcap:job:"
31
+
32
+ logger = getLogger("wrapper")
33
+
34
+ class ToolOptions(BaseModel):
35
+ name: Optional[str] = Field(None, description="Name to be used for this tool")
36
+ tags: Optional[list[str]] = Field(None, description="OpenAPI tag for this set of functions")
37
+ max_wait_time: Optional[float] = Field(5.0, description="max. time in seconds to wait for result and before returning RetryLater")
38
+ refresh_interval: Optional[int] = Field(3, description="Time in seconds to wait before chacking again for a job result (used in RetryLater)")
39
+ executor_opts: Optional[ExecutorOpts] = Field(None, description="Options for the executor")
40
+ post_route_opts: Optional[Dict[str, Any]] = Field({}, description="Addtitional options given the POST route constructor")
41
+ service_id: Optional[str] = Field(None, description="overriding the default service id")
42
+
43
+ # Define a generic type for Pydantic models
44
+ T = TypeVar("T", bound=BaseModel)
45
+
46
+ WorkerFn = Callable[[BaseModel, Optional[ExecutionContext], Optional[Response]], BaseModel]
47
+
48
+ @dataclass
49
+ class ToolDescription:
50
+ name: str
51
+ path_prefix: str
52
+ worker_fn: WorkerFn
53
+ input: Tuple[Optional[Type[BaseModel]], Dict[str, Any]]
54
+ executor: Executor
55
+
56
+ tools: List[ToolDescription] = []
57
+
58
+ def add_tool_api_route(
59
+ app: FastAPI,
60
+ path_prefix: str,
61
+ worker_fn: WorkerFn,
62
+ *,
63
+ opts: Optional[ToolOptions] = ToolOptions(),
64
+ context: Optional[ExecutionContext] = None
65
+ ):
66
+ """Add a few routes to `app` for use with an AI tool.
67
+
68
+ The tool itself is implemented in `worker_fn` where the first
69
+ argument is a pydantic model describing all the tool's "public" parameters.
70
+ The function is also expected to return it's result as a single pydantic model.
71
+ The tool function can have two optioanl paramters, one with the same type as
72
+ `context` and the second one with `fastapi.Request`. The context paramter will
73
+ be identical to the above `context`, while the `request` will be the incoming
74
+ request.
75
+
76
+ This function then sets up three endpoints:
77
+ - POST {path_prefix}: To request the execution of the tool (the 'job')
78
+ - GET {path_prefix}/{job_id}: To collect the result of a tool execution
79
+ - GET {path_prefix}: To obtain a description of the tool suitable for most agent frameworks
80
+
81
+ The POST request will only wait `opts.max_wait_time` for the tool to finish. If it
82
+ hasn't finished by then, a `204 No Content` code will be returned with additional
83
+ header fields `Location` and `Retry-later` to inform the caller where and approx.
84
+ when the result can be collected later.
85
+
86
+ The `opts` parameter allows for customization of the endpoints. See `ToolOptions`
87
+ for a more detailed description.
88
+
89
+ Args:
90
+ app (FastAPI): The FastAPI context
91
+ path_prefix (str): The path prefix to use for this set of endpoints
92
+ worker_fn (Callable[[BaseModel, Optional[ExecutionContext], Optional[Response]], BaseModel]): _description_
93
+ opts (Optional[ToolOptions], optional): Additional behaviour settings. Defaults to ToolOptions().
94
+ context (Optional[ExecutionContext], optional): An optional context to be provided to every invocation of `worker_fn`. Defaults to None.
95
+ """
96
+ def_name, def_tag = get_title_from_path(path_prefix)
97
+ if opts.tags is None:
98
+ if def_tag == '':
99
+ def_tag = "Tool"
100
+ opts.tags = [def_tag]
101
+ if opts.name is None:
102
+ if def_name == '':
103
+ def_name = "Execute the tool"
104
+ opts.name = def_name
105
+
106
+ output_model = get_function_return_type(worker_fn)
107
+ executor = Executor[output_model](worker_fn, opts=opts.executor_opts, context=context)
108
+
109
+ tools.append(ToolDescription(name=worker_fn.__name__,
110
+ path_prefix=path_prefix,
111
+ worker_fn=worker_fn,
112
+ input=get_input_type(worker_fn),
113
+ executor=executor))
114
+
115
+ _add_do_job_route(app, path_prefix, worker_fn, executor, opts)
116
+ _add_get_job_route(app, path_prefix, worker_fn, executor, opts)
117
+ _add_get_tool_def_route(app, path_prefix, worker_fn, opts)
118
+
119
+ def _add_do_job_route(app: FastAPI, path_prefix: str, worker_fn: Callable, executor: Executor, opts: ToolOptions):
120
+ input_model, _ = get_input_type(worker_fn)
121
+ output_model = get_function_return_type(worker_fn)
122
+ summary, description = (worker_fn.__doc__.lstrip() + "\n").split("\n", 1)
123
+
124
+ async def route(data: input_model, req: Request) -> output_model: # type: ignore
125
+ job_id = req.headers.get("job-id")
126
+ if job_id == None:
127
+ job_id = str(uuid6())
128
+ elif job_id.startswith(JOB_URN_PREFIX):
129
+ job_id = job_id[len(JOB_URN_PREFIX):]
130
+
131
+ if req.headers.get("prefer") == "respond-async":
132
+ timeout = 0
133
+ else:
134
+ toh = req.headers.get("timeout")
135
+ if toh != None:
136
+ timeout = int(toh)
137
+ else:
138
+ timeout = opts.max_wait_time
139
+ logger.info(f"starting job {path_prefix}/jobs/{job_id} - timeout: {timeout} seconds")
140
+
141
+ queue = await executor.execute(data, job_id, req)
142
+ try:
143
+ el = await asyncio.wait_for(queue.get(), timeout=timeout)
144
+ queue.task_done()
145
+ el = _return_job_result(el, job_id)
146
+ return el
147
+ except asyncio.TimeoutError:
148
+ logger.info(f"... defer job result to later - {job_id}")
149
+ return _return_try_later(job_id, path_prefix, opts)
150
+
151
+ responses = {
152
+ 204: {
153
+ "headers": {
154
+ "Location": {
155
+ "description": "The URL where to pick up the result of this request",
156
+ "type": "string",
157
+ },
158
+ "Retry-Later": {
159
+ "description": "The time to wait before checking for a result",
160
+ "type": "integer",
161
+ },
162
+ },
163
+ },
164
+ 400: { "model": ErrorModel, },
165
+ # 400: {"model": Error}, 401: {"model": Error}, 429: {"model": Error}},
166
+ }
167
+ app.add_api_route(
168
+ path_prefix,
169
+ route,
170
+ # name=opts.name,
171
+ summary=summary,
172
+ description=description.strip(),
173
+ methods=["POST"],
174
+ responses=responses,
175
+ tags=opts.tags,
176
+ response_model_exclude_none=True,
177
+ response_model_by_alias=True,
178
+ **opts.post_route_opts,
179
+ )
180
+
181
+ def _add_get_job_route(app: FastAPI, path_prefix: str, worker_fn: Callable, executor: Executor, opts: ToolOptions):
182
+ output_model = get_function_return_type(worker_fn)
183
+ def route(job_id: str) -> output_model: # type: ignore
184
+ if job_id.startswith(JOB_URN_PREFIX):
185
+ job_id = job_id[len(JOB_URN_PREFIX):]
186
+ try:
187
+ result = executor.lookup_job(job_id)
188
+ if result == None:
189
+ return _return_try_later(job_id, path_prefix, opts)
190
+ return _return_job_result(result, job_id)
191
+ except KeyError:
192
+ return Response(status_code=status.HTTP_404_NOT_FOUND,
193
+ content=f"job {job_id} can't be found. It either never existed or its result is no longer cached.")
194
+
195
+ responses = {
196
+ 400: { "model": ErrorModel, },
197
+ }
198
+ path = "/jobs/" + "{job_id}"
199
+ if path_prefix != "/":
200
+ path = path_prefix + path
201
+ app.add_api_route(
202
+ path,
203
+ route,
204
+ summary="Returns the result of a particular job.",
205
+ methods=["GET"],
206
+ responses=responses,
207
+ tags=opts.tags,
208
+ response_model_exclude_none=True,
209
+ response_model_by_alias=True,
210
+ )
211
+
212
+ def _return_job_result(el, job_id):
213
+ h = { "job-id": JOB_URN_PREFIX + job_id }
214
+ if isinstance(el, IvcapResult):
215
+ return Response(status_code=status.HTTP_200_OK, content=el.content, media_type=el.content_type, headers=h)
216
+ elif isinstance(el, ExecutionError):
217
+ if el.type == ValueError:
218
+ m = ErrorModel(message=el.error, code=400)
219
+ status_code=status.HTTP_400_BAD_REQUEST
220
+ else:
221
+ m = ExecutionErrorModel(message=el.error, traceback=el.traceback)
222
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
223
+
224
+ return Response(status_code=status_code, content=m.model_dump_json(indent=2), media_type="application/json", headers=h)
225
+
226
+ msg = json.dumps({"error": f"please report unexpected internal error - unexpected result type {type(el)}"})
227
+ return Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content=msg, media_type="application/json", headers=h)
228
+
229
+
230
+ def _add_get_tool_def_route(app: FastAPI, path_prefix: str, worker_fn: Callable, opts: ToolOptions):
231
+ async def route(req: Request) -> ToolDefinition: # type: ignore
232
+ service_id = opts.service_id
233
+ if service_id != None and service_id.startswith("/"):
234
+ # check if there is a forwarded header and prepand that
235
+ prefix = get_public_url_prefix(req)
236
+ service_id = f"{prefix}{service_id}"
237
+
238
+ return create_tool_definition(worker_fn, service_id=service_id)
239
+
240
+ app.add_api_route(
241
+ path_prefix,
242
+ route,
243
+ summary="Returns the description of this tool. Primarily used by agents.",
244
+ methods=["GET"],
245
+ tags=opts.tags,
246
+ response_model_exclude_none=True,
247
+ response_model_by_alias=True,
248
+ )
249
+
250
+
251
+ def _return_try_later(job_id: str, path_prefix: str, opts: ToolOptions):
252
+ location = f"/jobs/{job_id}"
253
+ if path_prefix != "/":
254
+ location = path_prefix + location
255
+ headers = {
256
+ "Location": location,
257
+ "Retry-Later": f"{opts.refresh_interval}",
258
+ "Ivcap-Self-Report-Result": "true"
259
+ }
260
+ return Response(status_code=status.HTTP_204_NO_CONTENT, headers=headers)
@@ -0,0 +1,60 @@
1
+ #
2
+ # Copyright (c) 2023 Commonwealth Scientific and Industrial Research Organisation (CSIRO). All rights reserved.
3
+ # Use of this source code is governed by a BSD-style license that can be
4
+ # found in the LICENSE file. See the AUTHORS file for names of contributors.
5
+ #
6
+ from typing import Optional
7
+
8
+ from .builder import ToolOptions, add_tool_api_route, WorkerFn
9
+ from .executor import ExecutionContext
10
+ from .server import get_fast_app
11
+
12
+
13
+ def ivcap_lambda(
14
+ path_prefix: str,
15
+ *,
16
+ opts: Optional[ToolOptions] = ToolOptions(),
17
+ context: Optional[ExecutionContext] = None,
18
+ ):
19
+ """Add a few routes to the service for use with an AI tool.
20
+
21
+ The tool itself is implemented in `worker_fn` where the first
22
+ argument is a pydantic model describing all the tool's "public" parameters.
23
+ The function is also expected to return it's result as a single pydantic model.
24
+ The tool function can have two optioanl paramters, one with the same type as
25
+ `context` and the second one with `fastapi.Request`. The context paramter will
26
+ be identical to the above `context`, while the `request` will be the incoming
27
+ request.
28
+
29
+ This function then sets up three endpoints:
30
+ - POST {path_prefix}: To request the execution of the tool (the 'job')
31
+ - GET {path_prefix}/{job_id}: To collect the result of a tool execution
32
+ - GET {path_prefix}: To obtain a description of the tool suitable for most agent frameworks
33
+
34
+ The POST request will only wait `opts.max_wait_time` for the tool to finish. If it
35
+ hasn't finished by then, a `204 No Content` code will be returned with additional
36
+ header fields `Location` and `Retry-later` to inform the caller where and approx.
37
+ when the result can be collected later.
38
+
39
+ The `opts` parameter allows for customization of the endpoints. See `ToolOptions`
40
+ for a more detailed description.
41
+
42
+ Args:
43
+ path_prefix (str): The path prefix to use for this set of endpoints
44
+ opts (Optional[ToolOptions], optional): Additional behaviour settings. Defaults to ToolOptions().
45
+ context (Optional[ExecutionContext], optional): An optional context to be provided to every invocation of `worker_fn`. Defaults to None.
46
+ """
47
+
48
+ def decorator(worker_fn: WorkerFn):
49
+ """
50
+ Args:
51
+ worker_fn (Callable[[BaseModel, Optional[ExecutionContext], Optional[Response]], BaseModel]): _description_
52
+ opts (Optional[ToolOptions], optional): Additional behaviour settings. Defaults to ToolOptions().
53
+ context (Optional[ExecutionContext], optional): An optional context to be provided to every invocation of `worker_fn`. Defaults to None.
54
+ """
55
+ add_tool_api_route(
56
+ get_fast_app(), path_prefix, worker_fn, opts=opts, context=context
57
+ )
58
+ return worker_fn
59
+
60
+ return decorator
@@ -0,0 +1,292 @@
1
+ #
2
+ # Copyright (c) 2023 Commonwealth Scientific and Industrial Research Organisation (CSIRO). All rights reserved.
3
+ # Use of this source code is governed by a BSD-style license that can be
4
+ # found in the LICENSE file. See the AUTHORS file for names of contributors.
5
+ #
6
+ import asyncio
7
+ import contextvars
8
+ import concurrent.futures
9
+ import threading
10
+ from time import sleep
11
+ import traceback
12
+ import contextlib
13
+ from typing import Any, Callable, Generic, List, Optional, TypeVar, Union
14
+ from cachetools import TTLCache
15
+ from fastapi import Request
16
+ from pydantic import BaseModel, Field
17
+ from ivcap_service import getLogger
18
+ from opentelemetry import trace, context
19
+ from opentelemetry.context.context import Context
20
+
21
+ from ivcap_service import get_input_type, push_result, verify_result, EventReporter
22
+ from ivcap_service import ExecutionError, create_event_reporter, JobContext
23
+ from ivcap_client import IVCAP
24
+
25
+ # Number of attempt to deliver job result before giving up
26
+ MAX_DELIVER_RESULT_ATTEMPTS = 4
27
+
28
+ logger = getLogger("executor")
29
+ tracer = trace.get_tracer("executor")
30
+
31
+
32
+ class ExecutionContext:
33
+ pass
34
+
35
+
36
+ class ThreadLocal(threading.local):
37
+ pass
38
+
39
+
40
+ T = TypeVar("T")
41
+
42
+
43
+ class ExecutorOpts(BaseModel):
44
+ job_cache_size: Optional[int] = Field(10000, description="size of job cache")
45
+ job_cache_ttl: Optional[int] = Field(
46
+ 3600, description="TTL of job entries in the job cache"
47
+ )
48
+ max_workers: Optional[int] = Field(
49
+ None,
50
+ description="size of thread pool to use. If None, a new thread pool will be created for each execution",
51
+ )
52
+
53
+
54
+ job_context: contextvars.ContextVar[Optional[JobContext]] = contextvars.ContextVar(
55
+ "ivcap", default=None
56
+ )
57
+
58
+
59
+ def get_job_context() -> Optional[JobContext]:
60
+ return job_context.get()
61
+
62
+
63
+ def get_event_reporter() -> Optional[EventReporter]:
64
+ """Get the current event reporter from the job context."""
65
+ jctxt = job_context.get()
66
+ return jctxt.report if jctxt is not None else None
67
+
68
+
69
+ def get_job_id() -> Optional[str]:
70
+ """Get the current job ID from the job context."""
71
+ jctxt = job_context.get()
72
+ return jctxt.job_id if jctxt is not None else None
73
+
74
+
75
+ class Executor(Generic[T]):
76
+ """
77
+ A generic class that executes a function in a thread pool and returns the result via an asyncio Queue.
78
+ The generic type T represents the return type of the function.
79
+ """
80
+
81
+ # _job_ctxt = JobContext()
82
+ _active_jobs = (
83
+ set()
84
+ ) # keep track of active jobs to block shutdown until they are done
85
+
86
+ @classmethod
87
+ def active_jobs(cls) -> List[str]:
88
+ """Returns a list of IDs of the currently active jobs"""
89
+ return list(cls._active_jobs)
90
+
91
+ @classmethod
92
+ def wait_for_exit_ready(cls):
93
+ """The server is calling this method when a shutdown request arrived. It will
94
+ proceed with the shutdown when this method returns.
95
+
96
+ We may implement functionality to only return when all active jobs have finsihed, as well as not
97
+ accepting any new incoming requests.
98
+ """
99
+ while len(cls._active_jobs) > 0:
100
+ logger.info(
101
+ f"blocking shutdown as {len(cls._active_jobs)} job(s) are still running"
102
+ )
103
+ sleep(5)
104
+ return
105
+
106
+ def __init__(
107
+ self,
108
+ func: Callable[..., T],
109
+ *,
110
+ opts: Optional[ExecutorOpts],
111
+ context: Optional[JobContext] = None,
112
+ ):
113
+ """
114
+ Initialize the Executor with a function and an optional thread pool.
115
+
116
+ Args:
117
+ func: The function to execute, returning type T
118
+ opts:
119
+ - job_cache_size: Optional size of job cache. Defaults to 1000
120
+ - job_cache_ttl: Optional TTL of job entries in the job cache. Defaults to 600 sec
121
+ - max_workers: Optional size of thread pool to use. If None, a new thread pool will be created for each execution.
122
+ """
123
+ self.func = func
124
+ if opts is None:
125
+ opts = ExecutorOpts()
126
+ self.job_cache = TTLCache(maxsize=opts.job_cache_size, ttl=opts.job_cache_ttl)
127
+ self.thread_pool = None
128
+ if opts.max_workers:
129
+ self.thread_pool = concurrent.futures.ThreadPoolExecutor(
130
+ max_workers=opts.max_workers
131
+ )
132
+
133
+ self.context = context
134
+ self.context_param = None
135
+ self.request_param = None
136
+ self.job_ctxt_param = None
137
+ _, extras = get_input_type(func)
138
+ for k, v in extras.items():
139
+ if isinstance(context, v):
140
+ self.context_param = k
141
+ elif v == Request:
142
+ self.request_param = k
143
+ elif v == JobContext:
144
+ self.job_ctxt_param = k
145
+ else:
146
+ raise Exception(f"unexpected function parameter '{k}'")
147
+
148
+ async def execute(
149
+ self, param: Any, job_id: str, req: Request, report_result=True
150
+ ) -> asyncio.Queue[Union[T, ExecutionError]]:
151
+ """
152
+ Execute the function with the given parameter in a thread and return a queue with the result.
153
+
154
+ Args:
155
+ param: Any The parameter to pass to the function
156
+ job_id: str ID of this job
157
+ req: Request FastAPI's request object
158
+
159
+ Returns:
160
+ An asyncio Queue that will contain either the result of type T or an ExecutionError
161
+ """
162
+ result_queue: asyncio.Queue[Union[T, ExecutionError]] = asyncio.Queue()
163
+ event_loop = asyncio.get_running_loop()
164
+ self.job_cache[job_id] = None
165
+
166
+ def _process_result(result):
167
+ """Verify the result, add it to the queue, and report it to IVCAP."""
168
+ try:
169
+ result = verify_result(result, job_id, logger)
170
+ except Exception as e:
171
+ result = ExecutionError(
172
+ error=str(e),
173
+ type=type(e).__name__,
174
+ traceback=traceback.format_exc(),
175
+ )
176
+ logger.warning(f"job {job_id} failed - {result.error}")
177
+ finally:
178
+ self.job_cache[job_id] = result
179
+ logger.info(f"job {job_id} finished, sending result message")
180
+ asyncio.run_coroutine_threadsafe(
181
+ result_queue.put(result),
182
+ event_loop,
183
+ )
184
+ if report_result:
185
+ push_result(result, job_id)
186
+ self.__class__._active_jobs.discard(job_id)
187
+
188
+ def _run(param: Any, ctxt: Context):
189
+ context.attach(ctxt) # OTEL
190
+ authorization = req.headers.get("authorization")
191
+ jctxt = JobContext(
192
+ job_id=job_id,
193
+ job_authorization=authorization,
194
+ report=create_event_reporter(
195
+ job_id=job_id, job_authorization=authorization
196
+ ),
197
+ )
198
+ job_context.set(jctxt)
199
+ kwargs = {}
200
+ if self.context_param is not None:
201
+ kwargs[self.context_param] = self.context
202
+ if self.request_param is not None:
203
+ kwargs[self.request_param] = req
204
+ if self.job_ctxt_param is not None:
205
+ kwargs[self.job_ctxt_param] = jctxt # self._job_ctxt
206
+
207
+ fname = self.func.__name__
208
+ with tracer.start_as_current_span(f"RUN {fname}") as span:
209
+ span.set_attribute("job.id", job_id)
210
+ span.set_attribute("job.name", fname)
211
+ loop = None
212
+ try:
213
+ self.__class__._active_jobs.add(job_id)
214
+ if asyncio.iscoroutinefunction(self.func):
215
+ loop = asyncio.new_event_loop()
216
+ asyncio.set_event_loop(loop)
217
+ try:
218
+ res = loop.run_until_complete(self.func(param, **kwargs))
219
+ except (asyncio.CancelledError, GeneratorExit):
220
+ # Propagate cancellation and generator shutdown properly
221
+ raise
222
+ finally:
223
+ # Gracefully shutdown remaining tasks and async generators to avoid
224
+ # 'Task exception was never retrieved' and similar warnings
225
+ try:
226
+ pending = [
227
+ t for t in asyncio.all_tasks(loop) if not t.done()
228
+ ]
229
+ for t in pending:
230
+ t.cancel()
231
+ if pending:
232
+ with contextlib.suppress(Exception):
233
+ loop.run_until_complete(
234
+ asyncio.gather(
235
+ *pending, return_exceptions=True
236
+ )
237
+ )
238
+ with contextlib.suppress(Exception):
239
+ loop.run_until_complete(loop.shutdown_asyncgens())
240
+ finally:
241
+ loop.close()
242
+ else:
243
+ res = self.func(param, **kwargs)
244
+ except (asyncio.CancelledError, GeneratorExit):
245
+ # Allow cooperative shutdown/cancellation to propagate cleanly
246
+ span.record_exception(Exception("cancelled"))
247
+ raise
248
+ except Exception as ex:
249
+ span.record_exception(ex)
250
+ logger.error(
251
+ f"while executing {job_id} - {type(ex).__name__}: {ex}"
252
+ )
253
+ res = ExecutionError(
254
+ error=str(ex),
255
+ type=type(ex).__name__,
256
+ traceback=traceback.format_exc(),
257
+ )
258
+ finally:
259
+ self.__class__._active_jobs.discard(job_id)
260
+
261
+ try:
262
+ _process_result(res)
263
+ except Exception as ex:
264
+ logger.error(f"while delivering result fo {job_id} - {ex}")
265
+
266
+ job_context.set(None)
267
+
268
+ # Use the provided thread pool or create a new one
269
+ use_pool = self.thread_pool or concurrent.futures.ThreadPoolExecutor(
270
+ max_workers=1
271
+ )
272
+ # Submit the function to the thread pool
273
+ future = use_pool.submit(_run, param, context.get_current())
274
+
275
+ # If we created a new pool, we should clean it up when done
276
+ if self.thread_pool is None:
277
+ future.add_done_callback(lambda _: use_pool.shutdown(wait=False))
278
+ return result_queue
279
+
280
+ def lookup_job(self, job_id: str) -> Union[T, ExecutionError, None]:
281
+ """Return the result of a job
282
+
283
+ Args:
284
+ job_id (str): The id of the job requested
285
+
286
+ Returns:
287
+ Union[T, ExecutionError, None]: Returns the result fo a job, 'None' is still in progress
288
+
289
+ Raises:
290
+ KeyError: Unknown job - may have already expired
291
+ """
292
+ return self.job_cache[job_id]
ivcap_lambda/logger.py ADDED
@@ -0,0 +1,37 @@
1
+ #
2
+ # Copyright (c) 2023 Commonwealth Scientific and Industrial Research Organisation (CSIRO). All rights reserved.
3
+ # Use of this source code is governed by a BSD-style license that can be
4
+ # found in the LICENSE file. See the AUTHORS file for names of contributors.
5
+ #
6
+ import json
7
+ import os
8
+ import logging
9
+ from ivcap_service import set_service_log_config
10
+
11
+ class SuppressPathsFilter(logging.Filter):
12
+ def __init__(self, targets=None):
13
+ super().__init__()
14
+ if targets is None:
15
+ self.targets = []
16
+ else:
17
+ self.targets = targets
18
+
19
+ def filter(self, record):
20
+ # Suppress logs for any request matching a target substring or path
21
+ # For uvicorn.access, HTTP info is in record.args: (client_addr, method, path, http_version, status)
22
+ path = ""
23
+ if hasattr(record, "args") and isinstance(record.args, tuple) and len(record.args) >= 3:
24
+ path = record.args[2]
25
+ for target in self.targets:
26
+ if target == path:
27
+ return False
28
+ return True
29
+
30
+ def logging_init(cfg_path: str=None):
31
+ if not cfg_path:
32
+ script_dir = os.path.dirname(__file__)
33
+ cfg_path = os.path.join(script_dir, "logging.json")
34
+
35
+ with open(cfg_path, 'r') as file:
36
+ config = json.load(file)
37
+ set_service_log_config(config)