Add support for Batch API test (#936)
This commit is contained in:
@@ -308,7 +308,6 @@ class TokenizerManager:
|
||||
event = asyncio.Event()
|
||||
state = ReqState([], False, event)
|
||||
self.rid_to_state[rid] = state
|
||||
|
||||
# Then wait for all responses
|
||||
output_list = []
|
||||
for i in range(batch_size):
|
||||
@@ -341,7 +340,6 @@ class TokenizerManager:
|
||||
)
|
||||
assert state.finished
|
||||
del self.rid_to_state[rid]
|
||||
|
||||
yield output_list
|
||||
|
||||
def _validate_input_length(self, input_ids: List[int]):
|
||||
|
||||
@@ -53,6 +53,7 @@ from sglang.srt.openai_api.protocol import (
|
||||
CompletionStreamResponse,
|
||||
DeltaMessage,
|
||||
ErrorResponse,
|
||||
FileDeleteResponse,
|
||||
FileRequest,
|
||||
FileResponse,
|
||||
LogProbs,
|
||||
@@ -174,6 +175,20 @@ async def v1_files_create(file: UploadFile, purpose: str, file_storage_pth: str
|
||||
return {"error": "Invalid input", "details": e.errors()}
|
||||
|
||||
|
||||
async def v1_delete_file(file_id: str):
|
||||
# Retrieve the file job from the in-memory storage
|
||||
file_response = file_id_response.get(file_id)
|
||||
if file_response is None:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
file_path = file_id_storage.get(file_id)
|
||||
if file_path is None:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
os.remove(file_path)
|
||||
del file_id_response[file_id]
|
||||
del file_id_storage[file_id]
|
||||
return FileDeleteResponse(id=file_id, deleted=True)
|
||||
|
||||
|
||||
async def v1_batches(tokenizer_manager, raw_request: Request):
|
||||
try:
|
||||
body = await raw_request.json()
|
||||
@@ -287,6 +302,13 @@ async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRe
|
||||
retrieve_batch = batch_storage[batch_id]
|
||||
retrieve_batch.output_file_id = output_file_id
|
||||
file_id_storage[output_file_id] = output_file_path
|
||||
file_id_response[output_file_id] = FileResponse(
|
||||
id=output_file_id,
|
||||
bytes=os.path.getsize(output_file_path),
|
||||
created_at=int(time.time()),
|
||||
filename=f"{output_file_id}.jsonl",
|
||||
purpose="batch_result",
|
||||
)
|
||||
# Update batch status to "completed"
|
||||
retrieve_batch.status = "completed"
|
||||
retrieve_batch.completed_at = int(time.time())
|
||||
@@ -380,7 +402,7 @@ def v1_generate_request(all_requests):
|
||||
else:
|
||||
prompt_kwargs = {"input_ids": prompt}
|
||||
else:
|
||||
if isinstance(prompts[0], str) or isinstance(propmt[0][0], str):
|
||||
if isinstance(prompts[0], str):
|
||||
prompt_kwargs = {"text": prompts}
|
||||
else:
|
||||
prompt_kwargs = {"input_ids": prompts}
|
||||
@@ -931,7 +953,6 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
|
||||
).__anext__()
|
||||
except ValueError as e:
|
||||
return create_error_response(str(e))
|
||||
|
||||
if not isinstance(ret, list):
|
||||
ret = [ret]
|
||||
|
||||
|
||||
@@ -95,6 +95,12 @@ class FileResponse(BaseModel):
|
||||
purpose: str
|
||||
|
||||
|
||||
class FileDeleteResponse(BaseModel):
|
||||
id: str
|
||||
object: str = "file"
|
||||
deleted: bool
|
||||
|
||||
|
||||
class BatchRequest(BaseModel):
|
||||
input_file_id: (
|
||||
str # The ID of an uploaded file that contains requests for the new batch
|
||||
|
||||
@@ -59,6 +59,7 @@ from sglang.srt.openai_api.adapter import (
|
||||
v1_batches,
|
||||
v1_chat_completions,
|
||||
v1_completions,
|
||||
v1_delete_file,
|
||||
v1_files_create,
|
||||
v1_retrieve_batch,
|
||||
v1_retrieve_file,
|
||||
@@ -175,6 +176,12 @@ async def openai_v1_files(file: UploadFile = File(...), purpose: str = Form("bat
|
||||
)
|
||||
|
||||
|
||||
@app.delete("/v1/files/{file_id}")
|
||||
async def delete_file(file_id: str):
|
||||
# https://platform.openai.com/docs/api-reference/files/delete
|
||||
return await v1_delete_file(file_id)
|
||||
|
||||
|
||||
@app.post("/v1/batches")
|
||||
async def openai_v1_batches(raw_request: Request):
|
||||
return await v1_batches(tokenizer_manager, raw_request)
|
||||
|
||||
Reference in New Issue
Block a user