aiqtoolkit 1.2.0a20250603__py3-none-any.whl → 1.2.0a20250605__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.
- aiq/front_ends/fastapi/fastapi_front_end_config.py +31 -8
- aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +184 -21
- aiq/front_ends/fastapi/job_store.py +47 -25
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/METADATA +3 -1
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/RECORD +10 -10
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/WHEEL +0 -0
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/entry_points.txt +0 -0
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/licenses/LICENSE.md +0 -0
- {aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/top_level.txt +0 -0
|
@@ -77,25 +77,45 @@ class AIQEvaluateRequest(BaseModel):
|
|
|
77
77
|
return config_file
|
|
78
78
|
|
|
79
79
|
|
|
80
|
-
class
|
|
80
|
+
class BaseAsyncResponse(BaseModel):
|
|
81
|
+
"""Base model for async responses."""
|
|
82
|
+
job_id: str = Field(description="Unique identifier for the job")
|
|
83
|
+
status: str = Field(description="Current status of the job")
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
class AIQEvaluateResponse(BaseAsyncResponse):
|
|
81
87
|
"""Response model for the evaluate endpoint."""
|
|
82
|
-
|
|
83
|
-
status: str = Field(description="Current status of the evaluation job")
|
|
88
|
+
pass
|
|
84
89
|
|
|
85
90
|
|
|
86
|
-
class
|
|
87
|
-
"""Response model for the
|
|
91
|
+
class AIQAsyncGenerateResponse(BaseAsyncResponse):
|
|
92
|
+
"""Response model for the async generation endpoint."""
|
|
93
|
+
pass
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class BaseAsyncStatusResponse(BaseModel):
|
|
97
|
+
"""Base model for async status responses."""
|
|
88
98
|
job_id: str = Field(description="Unique identifier for the evaluation job")
|
|
89
99
|
status: str = Field(description="Current status of the evaluation job")
|
|
90
|
-
config_file: str = Field(description="Path to the configuration file used for evaluation")
|
|
91
100
|
error: str | None = Field(default=None, description="Error message if the job failed")
|
|
92
|
-
output_path: str | None = Field(default=None,
|
|
93
|
-
description="Path to the output file if the job completed successfully")
|
|
94
101
|
created_at: datetime = Field(description="Timestamp when the job was created")
|
|
95
102
|
updated_at: datetime = Field(description="Timestamp when the job was last updated")
|
|
96
103
|
expires_at: datetime | None = Field(default=None, description="Timestamp when the job will expire")
|
|
97
104
|
|
|
98
105
|
|
|
106
|
+
class AIQEvaluateStatusResponse(BaseAsyncStatusResponse):
|
|
107
|
+
"""Response model for the evaluate status endpoint."""
|
|
108
|
+
config_file: str = Field(description="Path to the configuration file used for evaluation")
|
|
109
|
+
output_path: str | None = Field(default=None,
|
|
110
|
+
description="Path to the output file if the job completed successfully")
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
class AIQAsyncGenerationStatusResponse(BaseAsyncStatusResponse):
|
|
114
|
+
output: dict | None = Field(
|
|
115
|
+
default=None,
|
|
116
|
+
description="Output of the generate request, this is only available if the job completed successfully.")
|
|
117
|
+
|
|
118
|
+
|
|
99
119
|
class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
|
|
100
120
|
"""
|
|
101
121
|
A FastAPI based front end that allows an AIQ Toolkit workflow to be served as a microservice.
|
|
@@ -153,6 +173,9 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
|
|
|
153
173
|
port: int = Field(default=8000, description="Port to bind the server to", ge=0, le=65535)
|
|
154
174
|
reload: bool = Field(default=False, description="Enable auto-reload for development")
|
|
155
175
|
workers: int = Field(default=1, description="Number of workers to run", ge=1)
|
|
176
|
+
max_running_async_jobs: int = Field(default=10,
|
|
177
|
+
description="Maximum number of async jobs to run concurrently",
|
|
178
|
+
ge=1)
|
|
156
179
|
step_adaptor: StepAdaptorConfig = StepAdaptorConfig()
|
|
157
180
|
|
|
158
181
|
workflow: typing.Annotated[EndpointBase, Field(description="Endpoint for the default workflow.")] = EndpointBase(
|
|
@@ -16,6 +16,7 @@
|
|
|
16
16
|
import asyncio
|
|
17
17
|
import logging
|
|
18
18
|
import os
|
|
19
|
+
import time
|
|
19
20
|
import typing
|
|
20
21
|
from abc import ABC
|
|
21
22
|
from abc import abstractmethod
|
|
@@ -32,6 +33,7 @@ from fastapi.exceptions import HTTPException
|
|
|
32
33
|
from fastapi.middleware.cors import CORSMiddleware
|
|
33
34
|
from fastapi.responses import StreamingResponse
|
|
34
35
|
from pydantic import BaseModel
|
|
36
|
+
from pydantic import Field
|
|
35
37
|
|
|
36
38
|
from aiq.builder.workflow_builder import WorkflowBuilder
|
|
37
39
|
from aiq.data_models.api_server import AIQChatRequest
|
|
@@ -42,6 +44,8 @@ from aiq.data_models.config import AIQConfig
|
|
|
42
44
|
from aiq.eval.config import EvaluationRunOutput
|
|
43
45
|
from aiq.eval.evaluate import EvaluationRun
|
|
44
46
|
from aiq.eval.evaluate import EvaluationRunConfig
|
|
47
|
+
from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerateResponse
|
|
48
|
+
from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerationStatusResponse
|
|
45
49
|
from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateRequest
|
|
46
50
|
from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateResponse
|
|
47
51
|
from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateStatusResponse
|
|
@@ -68,6 +72,9 @@ class FastApiFrontEndPluginWorkerBase(ABC):
|
|
|
68
72
|
|
|
69
73
|
self._front_end_config = config.general.front_end
|
|
70
74
|
|
|
75
|
+
self._cleanup_tasks: list[str] = []
|
|
76
|
+
self._cleanup_tasks_lock = asyncio.Lock()
|
|
77
|
+
|
|
71
78
|
@property
|
|
72
79
|
def config(self) -> AIQConfig:
|
|
73
80
|
return self._config
|
|
@@ -92,10 +99,18 @@ class FastApiFrontEndPluginWorkerBase(ABC):
|
|
|
92
99
|
yield
|
|
93
100
|
|
|
94
101
|
# If a cleanup task is running, cancel it
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
102
|
+
async with self._cleanup_tasks_lock:
|
|
103
|
+
|
|
104
|
+
# Cancel all cleanup tasks
|
|
105
|
+
for task_name in self._cleanup_tasks:
|
|
106
|
+
cleanup_task: asyncio.Task | None = getattr(starting_app.state, task_name, None)
|
|
107
|
+
if cleanup_task is not None:
|
|
108
|
+
logger.info("Cancelling %s cleanup task", task_name)
|
|
109
|
+
cleanup_task.cancel()
|
|
110
|
+
else:
|
|
111
|
+
logger.warning("No cleanup task found for %s", task_name)
|
|
112
|
+
|
|
113
|
+
self._cleanup_tasks.clear()
|
|
99
114
|
|
|
100
115
|
logger.debug("Closing AIQ Toolkit server from process %s", os.getpid())
|
|
101
116
|
|
|
@@ -153,6 +168,32 @@ class RouteInfo(BaseModel):
|
|
|
153
168
|
|
|
154
169
|
class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
155
170
|
|
|
171
|
+
@staticmethod
|
|
172
|
+
async def _periodic_cleanup(name: str, job_store: JobStore, sleep_time_sec: int = 300):
|
|
173
|
+
while True:
|
|
174
|
+
try:
|
|
175
|
+
job_store.cleanup_expired_jobs()
|
|
176
|
+
logger.debug("Expired %s jobs cleaned up", name)
|
|
177
|
+
except Exception as e:
|
|
178
|
+
logger.error("Error during %s job cleanup: %s", name, e)
|
|
179
|
+
await asyncio.sleep(sleep_time_sec)
|
|
180
|
+
|
|
181
|
+
async def create_cleanup_task(self, app: FastAPI, name: str, job_store: JobStore, sleep_time_sec: int = 300):
|
|
182
|
+
# Schedule periodic cleanup of expired jobs on first job creation
|
|
183
|
+
attr_name = f"{name}_cleanup_task"
|
|
184
|
+
|
|
185
|
+
# Cheap check, if it doesn't exist, we will need to re-check after we acquire the lock
|
|
186
|
+
if not hasattr(app.state, attr_name):
|
|
187
|
+
async with self._cleanup_tasks_lock:
|
|
188
|
+
if not hasattr(app.state, attr_name):
|
|
189
|
+
logger.info("Starting %s periodic cleanup task", name)
|
|
190
|
+
setattr(
|
|
191
|
+
app.state,
|
|
192
|
+
attr_name,
|
|
193
|
+
asyncio.create_task(
|
|
194
|
+
self._periodic_cleanup(name=name, job_store=job_store, sleep_time_sec=sleep_time_sec)))
|
|
195
|
+
self._cleanup_tasks.append(attr_name)
|
|
196
|
+
|
|
156
197
|
def get_step_adaptor(self) -> StepAdaptor:
|
|
157
198
|
|
|
158
199
|
return StepAdaptor(self.front_end_config.step_adaptor)
|
|
@@ -198,21 +239,6 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
198
239
|
# Don't run multiple evaluations at the same time
|
|
199
240
|
evaluation_lock = asyncio.Lock()
|
|
200
241
|
|
|
201
|
-
async def periodic_cleanup(job_store: JobStore):
|
|
202
|
-
while True:
|
|
203
|
-
try:
|
|
204
|
-
job_store.cleanup_expired_jobs()
|
|
205
|
-
logger.debug("Expired jobs cleaned up")
|
|
206
|
-
except Exception as e:
|
|
207
|
-
logger.error("Error during job cleanup: %s", str(e))
|
|
208
|
-
await asyncio.sleep(300) # every 5 minutes
|
|
209
|
-
|
|
210
|
-
def create_cleanup_task():
|
|
211
|
-
# Schedule periodic cleanup of expired jobs on first job creation
|
|
212
|
-
if not hasattr(app.state, "cleanup_task"):
|
|
213
|
-
logger.info("Starting periodic cleanup task")
|
|
214
|
-
app.state.cleanup_task = asyncio.create_task(periodic_cleanup(job_store))
|
|
215
|
-
|
|
216
242
|
async def run_evaluation(job_id: str, config_file: str, reps: int, session_manager: AIQSessionManager):
|
|
217
243
|
"""Background task to run the evaluation."""
|
|
218
244
|
async with evaluation_lock:
|
|
@@ -250,7 +276,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
250
276
|
return AIQEvaluateResponse(job_id=job.job_id, status=job.status)
|
|
251
277
|
|
|
252
278
|
job_id = job_store.create_job(request.config_file, request.job_id, request.expiry_seconds)
|
|
253
|
-
create_cleanup_task()
|
|
279
|
+
await self.create_cleanup_task(app=app, name="async_evaluation", job_store=job_store)
|
|
254
280
|
background_tasks.add_task(run_evaluation, job_id, request.config_file, request.reps, session_manager)
|
|
255
281
|
|
|
256
282
|
return AIQEvaluateResponse(job_id=job_id, status="submitted")
|
|
@@ -276,7 +302,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
276
302
|
if not job:
|
|
277
303
|
logger.warning("Job %s not found", job_id)
|
|
278
304
|
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
|
279
|
-
logger.info(
|
|
305
|
+
logger.info("Found job %s with status %s", job_id, job.status)
|
|
280
306
|
return translate_job_to_response(job)
|
|
281
307
|
|
|
282
308
|
async def get_last_job_status(http_request: Request) -> AIQEvaluateStatusResponse:
|
|
@@ -370,9 +396,28 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
370
396
|
GenerateStreamResponseType = workflow.streaming_output_schema # pylint: disable=invalid-name
|
|
371
397
|
GenerateSingleResponseType = workflow.single_output_schema # pylint: disable=invalid-name
|
|
372
398
|
|
|
399
|
+
# Append job_id and expiry_seconds to the input schema, this effectively makes these reserved keywords
|
|
400
|
+
# Consider prefixing these with "aiq_" to avoid conflicts
|
|
401
|
+
class AIQAsyncGenerateRequest(GenerateBodyType):
|
|
402
|
+
job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job")
|
|
403
|
+
sync_timeout: int = Field(
|
|
404
|
+
default=0,
|
|
405
|
+
ge=0,
|
|
406
|
+
le=300,
|
|
407
|
+
description="Attempt to perform the job synchronously up until `sync_timeout` sectonds, "
|
|
408
|
+
"if the job hasn't been completed by then a job_id will be returned with a status code of 202.")
|
|
409
|
+
expiry_seconds: int = Field(default=JobStore.DEFAULT_EXPIRY,
|
|
410
|
+
ge=JobStore.MIN_EXPIRY,
|
|
411
|
+
le=JobStore.MAX_EXPIRY,
|
|
412
|
+
description="Optional time (in seconds) before the job expires. "
|
|
413
|
+
"Clamped between 600 (10 min) and 86400 (24h).")
|
|
414
|
+
|
|
373
415
|
# Ensure that the input is in the body. POD types are treated as query parameters
|
|
374
416
|
if (not issubclass(GenerateBodyType, BaseModel)):
|
|
375
417
|
GenerateBodyType = typing.Annotated[GenerateBodyType, Body()]
|
|
418
|
+
else:
|
|
419
|
+
logger.info("Expecting generate request payloads in the following format: %s",
|
|
420
|
+
GenerateBodyType.model_fields)
|
|
376
421
|
|
|
377
422
|
response_500 = {
|
|
378
423
|
"description": "Internal Server Error",
|
|
@@ -385,6 +430,12 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
385
430
|
},
|
|
386
431
|
}
|
|
387
432
|
|
|
433
|
+
# Create job store for tracking async generation jobs
|
|
434
|
+
job_store = JobStore()
|
|
435
|
+
|
|
436
|
+
# Run up to max_running_async_jobs jobs at the same time
|
|
437
|
+
async_job_concurrency = asyncio.Semaphore(self._front_end_config.max_running_async_jobs)
|
|
438
|
+
|
|
388
439
|
def get_single_endpoint(result_type: type | None):
|
|
389
440
|
|
|
390
441
|
async def get_single(response: Response, request: Request):
|
|
@@ -482,6 +533,96 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
482
533
|
|
|
483
534
|
return post_stream
|
|
484
535
|
|
|
536
|
+
async def run_generation(job_id: str,
|
|
537
|
+
payload: typing.Any,
|
|
538
|
+
session_manager: AIQSessionManager,
|
|
539
|
+
result_type: type):
|
|
540
|
+
"""Background task to run the evaluation."""
|
|
541
|
+
async with async_job_concurrency:
|
|
542
|
+
try:
|
|
543
|
+
result = await generate_single_response(payload=payload,
|
|
544
|
+
session_manager=session_manager,
|
|
545
|
+
result_type=result_type)
|
|
546
|
+
job_store.update_status(job_id, "success", output=result)
|
|
547
|
+
except Exception as e:
|
|
548
|
+
logger.error("Error in evaluation job %s: %s", job_id, e)
|
|
549
|
+
job_store.update_status(job_id, "failure", error=str(e))
|
|
550
|
+
|
|
551
|
+
def _job_status_to_response(job: JobInfo) -> AIQAsyncGenerationStatusResponse:
|
|
552
|
+
job_output = job.output
|
|
553
|
+
if job_output is not None:
|
|
554
|
+
job_output = job_output.model_dump()
|
|
555
|
+
return AIQAsyncGenerationStatusResponse(job_id=job.job_id,
|
|
556
|
+
status=job.status,
|
|
557
|
+
error=job.error,
|
|
558
|
+
output=job_output,
|
|
559
|
+
created_at=job.created_at,
|
|
560
|
+
updated_at=job.updated_at,
|
|
561
|
+
expires_at=job_store.get_expires_at(job))
|
|
562
|
+
|
|
563
|
+
def post_async_generation(request_type: type, final_result_type: type):
|
|
564
|
+
|
|
565
|
+
async def start_async_generation(
|
|
566
|
+
request: request_type, background_tasks: BackgroundTasks, response: Response,
|
|
567
|
+
http_request: Request) -> AIQAsyncGenerateResponse | AIQAsyncGenerationStatusResponse:
|
|
568
|
+
"""Handle async generation requests."""
|
|
569
|
+
|
|
570
|
+
async with session_manager.session(request=http_request):
|
|
571
|
+
|
|
572
|
+
# if job_id is present and already exists return the job info
|
|
573
|
+
if request.job_id:
|
|
574
|
+
job = job_store.get_job(request.job_id)
|
|
575
|
+
if job:
|
|
576
|
+
return AIQAsyncGenerateResponse(job_id=job.job_id, status=job.status)
|
|
577
|
+
|
|
578
|
+
job_id = job_store.create_job(job_id=request.job_id, expiry_seconds=request.expiry_seconds)
|
|
579
|
+
await self.create_cleanup_task(app=app, name="async_generation", job_store=job_store)
|
|
580
|
+
|
|
581
|
+
# The fastapi/starlette background tasks won't begin executing until after the response is sent
|
|
582
|
+
# to the client, so we need to wrap the task in a function, alowing us to start the task now,
|
|
583
|
+
# and allowing the background task function to await the results.
|
|
584
|
+
task = asyncio.create_task(
|
|
585
|
+
run_generation(job_id=job_id,
|
|
586
|
+
payload=request,
|
|
587
|
+
session_manager=session_manager,
|
|
588
|
+
result_type=final_result_type))
|
|
589
|
+
|
|
590
|
+
async def wrapped_task(t: asyncio.Task):
|
|
591
|
+
return await t
|
|
592
|
+
|
|
593
|
+
background_tasks.add_task(wrapped_task, task)
|
|
594
|
+
|
|
595
|
+
now = time.time()
|
|
596
|
+
sync_timeout = now + request.sync_timeout
|
|
597
|
+
while time.time() < sync_timeout:
|
|
598
|
+
job = job_store.get_job(job_id)
|
|
599
|
+
if job is not None and job.status not in job_store.ACTIVE_STATUS:
|
|
600
|
+
# If the job is done, return the result
|
|
601
|
+
response.status_code = 200
|
|
602
|
+
return _job_status_to_response(job)
|
|
603
|
+
|
|
604
|
+
# Sleep for a short time before checking again
|
|
605
|
+
await asyncio.sleep(0.1)
|
|
606
|
+
|
|
607
|
+
response.status_code = 202
|
|
608
|
+
return AIQAsyncGenerateResponse(job_id=job_id, status="submitted")
|
|
609
|
+
|
|
610
|
+
return start_async_generation
|
|
611
|
+
|
|
612
|
+
async def get_async_job_status(job_id: str, http_request: Request) -> AIQAsyncGenerationStatusResponse:
|
|
613
|
+
"""Get the status of an async job."""
|
|
614
|
+
logger.info("Getting status for job %s", job_id)
|
|
615
|
+
|
|
616
|
+
async with session_manager.session(request=http_request):
|
|
617
|
+
|
|
618
|
+
job = job_store.get_job(job_id)
|
|
619
|
+
if not job:
|
|
620
|
+
logger.warning("Job %s not found", job_id)
|
|
621
|
+
raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
|
|
622
|
+
|
|
623
|
+
logger.info("Found job %s with status %s", job_id, job.status)
|
|
624
|
+
return _job_status_to_response(job)
|
|
625
|
+
|
|
485
626
|
if (endpoint.path):
|
|
486
627
|
if (endpoint.method == "GET"):
|
|
487
628
|
|
|
@@ -554,9 +695,31 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
|
|
|
554
695
|
responses={500: response_500},
|
|
555
696
|
)
|
|
556
697
|
|
|
698
|
+
app.add_api_route(
|
|
699
|
+
path=f"{endpoint.path}/async",
|
|
700
|
+
endpoint=post_async_generation(request_type=AIQAsyncGenerateRequest,
|
|
701
|
+
final_result_type=GenerateSingleResponseType),
|
|
702
|
+
methods=[endpoint.method],
|
|
703
|
+
response_model=AIQAsyncGenerateResponse | AIQAsyncGenerationStatusResponse,
|
|
704
|
+
description="Start an async generate job",
|
|
705
|
+
responses={500: response_500},
|
|
706
|
+
)
|
|
557
707
|
else:
|
|
558
708
|
raise ValueError(f"Unsupported method {endpoint.method}")
|
|
559
709
|
|
|
710
|
+
app.add_api_route(
|
|
711
|
+
path=f"{endpoint.path}/async/job/{{job_id}}",
|
|
712
|
+
endpoint=get_async_job_status,
|
|
713
|
+
methods=["GET"],
|
|
714
|
+
response_model=AIQAsyncGenerationStatusResponse,
|
|
715
|
+
description="Get the status of an async job",
|
|
716
|
+
responses={
|
|
717
|
+
404: {
|
|
718
|
+
"description": "Job not found"
|
|
719
|
+
}, 500: response_500
|
|
720
|
+
},
|
|
721
|
+
)
|
|
722
|
+
|
|
560
723
|
if (endpoint.openai_api_path):
|
|
561
724
|
if (endpoint.method == "GET"):
|
|
562
725
|
|
|
@@ -16,6 +16,7 @@
|
|
|
16
16
|
import logging
|
|
17
17
|
import os
|
|
18
18
|
import shutil
|
|
19
|
+
import threading
|
|
19
20
|
from datetime import UTC
|
|
20
21
|
from datetime import datetime
|
|
21
22
|
from datetime import timedelta
|
|
@@ -40,12 +41,13 @@ class JobStatus(str, Enum):
|
|
|
40
41
|
class JobInfo(BaseModel):
|
|
41
42
|
job_id: str
|
|
42
43
|
status: JobStatus
|
|
43
|
-
config_file: str
|
|
44
|
+
config_file: str | None
|
|
44
45
|
error: str | None
|
|
45
46
|
output_path: str | None
|
|
46
47
|
created_at: datetime
|
|
47
48
|
updated_at: datetime
|
|
48
49
|
expiry_seconds: int
|
|
50
|
+
output: BaseModel | None = None
|
|
49
51
|
|
|
50
52
|
|
|
51
53
|
class JobStore:
|
|
@@ -59,8 +61,12 @@ class JobStore:
|
|
|
59
61
|
|
|
60
62
|
def __init__(self):
|
|
61
63
|
self._jobs = {}
|
|
64
|
+
self._lock = threading.Lock() # Ensure thread safety for job operations
|
|
62
65
|
|
|
63
|
-
def create_job(self,
|
|
66
|
+
def create_job(self,
|
|
67
|
+
config_file: str | None = None,
|
|
68
|
+
job_id: str | None = None,
|
|
69
|
+
expiry_seconds: int = DEFAULT_EXPIRY) -> str:
|
|
64
70
|
if job_id is None:
|
|
65
71
|
job_id = str(uuid4())
|
|
66
72
|
|
|
@@ -76,46 +82,62 @@ class JobStore:
|
|
|
76
82
|
error=None,
|
|
77
83
|
output_path=None,
|
|
78
84
|
expiry_seconds=clamped_expiry)
|
|
79
|
-
|
|
85
|
+
|
|
86
|
+
with self._lock:
|
|
87
|
+
self._jobs[job_id] = job
|
|
88
|
+
|
|
80
89
|
logger.info("Created new job %s with config %s", job_id, config_file)
|
|
81
90
|
return job_id
|
|
82
91
|
|
|
83
|
-
def update_status(self,
|
|
92
|
+
def update_status(self,
|
|
93
|
+
job_id: str,
|
|
94
|
+
status: str,
|
|
95
|
+
error: str | None = None,
|
|
96
|
+
output_path: str | None = None,
|
|
97
|
+
output: BaseModel | None = None):
|
|
84
98
|
if job_id not in self._jobs:
|
|
85
99
|
raise ValueError(f"Job {job_id} not found")
|
|
86
100
|
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
101
|
+
with self._lock:
|
|
102
|
+
job = self._jobs[job_id]
|
|
103
|
+
job.status = status
|
|
104
|
+
job.error = error
|
|
105
|
+
job.output_path = output_path
|
|
106
|
+
job.updated_at = datetime.now(UTC)
|
|
107
|
+
job.output = output
|
|
92
108
|
|
|
93
109
|
def get_status(self, job_id: str) -> JobInfo | None:
|
|
94
|
-
|
|
110
|
+
with self._lock:
|
|
111
|
+
return self._jobs.get(job_id)
|
|
95
112
|
|
|
96
113
|
def list_jobs(self):
|
|
97
|
-
|
|
114
|
+
with self._lock:
|
|
115
|
+
return self._jobs
|
|
98
116
|
|
|
99
117
|
def get_job(self, job_id: str) -> JobInfo | None:
|
|
100
118
|
"""Get a job by its ID."""
|
|
101
|
-
|
|
119
|
+
with self._lock:
|
|
120
|
+
return self._jobs.get(job_id)
|
|
102
121
|
|
|
103
122
|
def get_last_job(self) -> JobInfo | None:
|
|
104
123
|
"""Get the last created job."""
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
124
|
+
with self._lock:
|
|
125
|
+
if not self._jobs:
|
|
126
|
+
logger.info("No jobs found in job store")
|
|
127
|
+
return None
|
|
128
|
+
last_job = max(self._jobs.values(), key=lambda job: job.created_at)
|
|
129
|
+
logger.info("Retrieved last job %s created at %s", last_job.job_id, last_job.created_at)
|
|
130
|
+
return last_job
|
|
111
131
|
|
|
112
132
|
def get_jobs_by_status(self, status: str) -> list[JobInfo]:
|
|
113
133
|
"""Get all jobs with the specified status."""
|
|
114
|
-
|
|
134
|
+
with self._lock:
|
|
135
|
+
return [job for job in self._jobs.values() if job.status == status]
|
|
115
136
|
|
|
116
137
|
def get_all_jobs(self) -> list[JobInfo]:
|
|
117
138
|
"""Get all jobs in the store."""
|
|
118
|
-
|
|
139
|
+
with self._lock:
|
|
140
|
+
return list(self._jobs.values())
|
|
119
141
|
|
|
120
142
|
def get_expires_at(self, job: JobInfo) -> datetime | None:
|
|
121
143
|
"""Get the time for a job to expire."""
|
|
@@ -132,7 +154,8 @@ class JobStore:
|
|
|
132
154
|
now = datetime.now(UTC)
|
|
133
155
|
|
|
134
156
|
# Filter out active jobs
|
|
135
|
-
|
|
157
|
+
with self._lock:
|
|
158
|
+
finished_jobs = {job_id: job for job_id, job in self._jobs.items() if job.status not in self.ACTIVE_STATUS}
|
|
136
159
|
|
|
137
160
|
# Sort finished jobs by updated_at descending
|
|
138
161
|
sorted_finished = sorted(finished_jobs.items(), key=lambda item: item[1].updated_at, reverse=True)
|
|
@@ -155,7 +178,6 @@ class JobStore:
|
|
|
155
178
|
elif os.path.isdir(job.output_path):
|
|
156
179
|
shutil.rmtree(job.output_path)
|
|
157
180
|
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
del self._jobs[job_id]
|
|
181
|
+
with self._lock:
|
|
182
|
+
for job_id in expired_ids:
|
|
183
|
+
del self._jobs[job_id]
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: aiqtoolkit
|
|
3
|
-
Version: 1.2.
|
|
3
|
+
Version: 1.2.0a20250605
|
|
4
4
|
Summary: NVIDIA Agent Intelligence toolkit
|
|
5
5
|
Author: NVIDIA Corporation
|
|
6
6
|
Maintainer: NVIDIA Corporation
|
|
@@ -271,6 +271,8 @@ Provides-Extra: profiling
|
|
|
271
271
|
Requires-Dist: matplotlib~=3.9; extra == "profiling"
|
|
272
272
|
Requires-Dist: prefixspan~=0.5.2; extra == "profiling"
|
|
273
273
|
Requires-Dist: scikit-learn~=1.6; extra == "profiling"
|
|
274
|
+
Provides-Extra: gunicorn
|
|
275
|
+
Requires-Dist: gunicorn~=23.0; extra == "gunicorn"
|
|
274
276
|
Dynamic: license-file
|
|
275
277
|
|
|
276
278
|
<!--
|
|
@@ -141,11 +141,11 @@ aiq/front_ends/console/console_front_end_plugin.py,sha256=CzadUoHmzrHC_MWn2Fkgh_
|
|
|
141
141
|
aiq/front_ends/console/register.py,sha256=a84M0jWUFTgOQVyrUiS7UJcxx84i1zhCb1yRkjhapiQ,1159
|
|
142
142
|
aiq/front_ends/cron/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
|
|
143
143
|
aiq/front_ends/fastapi/__init__.py,sha256=Xs1JQ16L9btwreh4pdGKwskffAw1YFO48jKrU4ib_7c,685
|
|
144
|
-
aiq/front_ends/fastapi/fastapi_front_end_config.py,sha256=
|
|
144
|
+
aiq/front_ends/fastapi/fastapi_front_end_config.py,sha256=8Pb0dpICjc_zGbBrA5cZ92QTVjjePK9L0zfnFrrBnWc,9483
|
|
145
145
|
aiq/front_ends/fastapi/fastapi_front_end_plugin.py,sha256=NdE5c4pS5sMYopI3PUaE397K8HF-GuZoOmpAi1ReDRk,4002
|
|
146
|
-
aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py,sha256=
|
|
146
|
+
aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py,sha256=ZbAM3qZFUQJDp2sMyNIyJ6enhB_2gx6aSX7UPkgUtHU,36064
|
|
147
147
|
aiq/front_ends/fastapi/intermediate_steps_subscriber.py,sha256=2Y3ZXfiu8d85qJRWbT3-ex6iFDuXO6TnNtjQ6Cjl-tc,3131
|
|
148
|
-
aiq/front_ends/fastapi/job_store.py,sha256=
|
|
148
|
+
aiq/front_ends/fastapi/job_store.py,sha256=AZC0a8miKqSI9YbW19-AA_GPITVY9bYHXn9H8saNrfQ,6389
|
|
149
149
|
aiq/front_ends/fastapi/main.py,sha256=ftLYy8YEyiY38kxReRwAVKwQsTa7QrKjnBU8V1omXaU,2789
|
|
150
150
|
aiq/front_ends/fastapi/message_handler.py,sha256=3rFDXG633Upom1taW3ab_sC3KKxN8Y_WoM_kXtJ3K6o,12640
|
|
151
151
|
aiq/front_ends/fastapi/message_validator.py,sha256=NqjeIG0InGAS6yooEnTaYwjfy3qtQHNgmdJ4zZlxgSQ,17407
|
|
@@ -308,10 +308,10 @@ aiq/utils/reactive/base/observer_base.py,sha256=UAlyAY_ky4q2t0P81RVFo2Bs_R7z5Nde
|
|
|
308
308
|
aiq/utils/reactive/base/subject_base.py,sha256=Ed-AC6P7cT3qkW1EXjzbd5M9WpVoeN_9KCe3OM3FLU4,2521
|
|
309
309
|
aiq/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
310
310
|
aiq/utils/settings/global_settings.py,sha256=U9TCLdoZsKq5qOVGjREipGVv9e-FlStzqy5zv82_VYk,7454
|
|
311
|
-
aiqtoolkit-1.2.
|
|
312
|
-
aiqtoolkit-1.2.
|
|
313
|
-
aiqtoolkit-1.2.
|
|
314
|
-
aiqtoolkit-1.2.
|
|
315
|
-
aiqtoolkit-1.2.
|
|
316
|
-
aiqtoolkit-1.2.
|
|
317
|
-
aiqtoolkit-1.2.
|
|
311
|
+
aiqtoolkit-1.2.0a20250605.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
|
|
312
|
+
aiqtoolkit-1.2.0a20250605.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
|
|
313
|
+
aiqtoolkit-1.2.0a20250605.dist-info/METADATA,sha256=zSXx7wEwV1TwY_ONAYFykvC5BuygtM79ehAye0K1t5E,20250
|
|
314
|
+
aiqtoolkit-1.2.0a20250605.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
|
|
315
|
+
aiqtoolkit-1.2.0a20250605.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
|
|
316
|
+
aiqtoolkit-1.2.0a20250605.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
|
|
317
|
+
aiqtoolkit-1.2.0a20250605.dist-info/RECORD,,
|
|
File without changes
|
{aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/entry_points.txt
RENAMED
|
File without changes
|
|
File without changes
|
{aiqtoolkit-1.2.0a20250603.dist-info → aiqtoolkit-1.2.0a20250605.dist-info}/licenses/LICENSE.md
RENAMED
|
File without changes
|
|
File without changes
|