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.
- ivcap_lambda/__init__.py +17 -0
- ivcap_lambda/builder.py +260 -0
- ivcap_lambda/decorators.py +60 -0
- ivcap_lambda/executor.py +292 -0
- ivcap_lambda/logger.py +37 -0
- ivcap_lambda/logging.json +77 -0
- ivcap_lambda/mcp.py +261 -0
- ivcap_lambda/py.typed +1 -0
- ivcap_lambda/secret.py +44 -0
- ivcap_lambda/server.py +217 -0
- ivcap_lambda/service_definition.py +44 -0
- ivcap_lambda/utils.py +100 -0
- ivcap_lambda/version.py +22 -0
- ivcap_lambda-0.7.22.dist-info/METADATA +134 -0
- ivcap_lambda-0.7.22.dist-info/RECORD +18 -0
- ivcap_lambda-0.7.22.dist-info/WHEEL +4 -0
- ivcap_lambda-0.7.22.dist-info/licenses/AUTHORS.md +5 -0
- ivcap_lambda-0.7.22.dist-info/licenses/LICENSE +29 -0
ivcap_lambda/__init__.py
ADDED
|
@@ -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
|
ivcap_lambda/builder.py
ADDED
|
@@ -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
|
ivcap_lambda/executor.py
ADDED
|
@@ -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)
|