Cancel batch endpoint

This commit is contained in:
Brandon Rising 2023-07-31 16:05:27 -04:00
parent 55b921818d
commit 02aa93c67c

View File

@ -15,7 +15,7 @@ from ...services.graph import (
GraphExecutionState, GraphExecutionState,
NodeAlreadyExecutedError, NodeAlreadyExecutedError,
) )
from ...services.batch_manager import Batch from ...services.batch_manager import Batch, BatchProcess
from ...services.item_storage import PaginatedResults from ...services.item_storage import PaginatedResults
from ..dependencies import ApiDependencies from ..dependencies import ApiDependencies
@ -42,19 +42,32 @@ async def create_session(
"/batch", "/batch",
operation_id="create_batch", operation_id="create_batch",
responses={ responses={
200: {"model": GraphExecutionState}, 200: {"model": BatchProcess},
400: {"description": "Invalid json"}, 400: {"description": "Invalid json"},
}, },
) )
async def create_batch( async def create_batch(
graph: Optional[Graph] = Body(default=None, description="The graph to initialize the session with"), graph: Optional[Graph] = Body(default=None, description="The graph to initialize the session with"),
batches: list[Batch] = Body(description="Batch config to apply to the given graph") batches: list[Batch] = Body(description="Batch config to apply to the given graph")
) -> GraphExecutionState: ) -> BatchProcess:
"""Creates a new session, optionally initializing it with an invocation graph""" """Creates and starts a new new batch process"""
session = ApiDependencies.invoker.services.batch_manager.run_batch_process(batches, graph) session = ApiDependencies.invoker.services.batch_manager.run_batch_process(batches, graph)
return session return session
@session_router.delete(
"{batch_process_id}/batch",
operation_id="cancel_batch",
responses={202: {"description": "The batch is canceled"}},
)
async def cancel_batch(
batch_process_id: str = Path(description="The id of the batch process to cancel"),
) -> Response:
"""Creates and starts a new new batch process"""
ApiDependencies.invoker.services.batch_manager.cancel_batch_process(batch_process_id)
return Response(status_code=202)
@session_router.get( @session_router.get(
"/", "/",
operation_id="list_sessions", operation_id="list_sessions",