From 8fb1f9696efcf7b51ec3851f2d295aa361ea501d Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Mon, 10 Mar 2025 23:29:04 +0000 Subject: [PATCH] feat: Switch synthetic data generation API to jobs pattern The API should behave in a way similar to existing training and eval flows where a long running task is sent to background; the client receives a job ID to follow status and extract artifacts. Note: there are no providers for this API implemented yet, so no implementation changes seem to be needed. Signed-off-by: Ihar Hrachyshka --- docs/_static/llama-stack-spec.html | 446 +++++++++++++++--- docs/_static/llama-stack-spec.yaml | 266 +++++++++-- .../synthetic_data_generation.py | 52 +- 3 files changed, 643 insertions(+), 121 deletions(-) diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 1a8169090..f63b8ec97 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -230,6 +230,42 @@ } } }, + "/v1/synthetic-data-generation/job/cancel": { + "post": { + "responses": { + "200": { + "description": "OK" + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "SyntheticDataGeneration (Coming Soon)" + ], + "description": "", + "parameters": [], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CancelSyntheticDataGenerationJobRequest" + } + } + }, + "required": true + } + } + }, "/v1/post-training/job/cancel": { "post": { "responses": { @@ -1493,6 +1529,137 @@ } } }, + "/v1/synthetic-data-generation/job/artifacts": { + "get": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "$ref": "#/components/schemas/SyntheticDataGenerationJobArtifactsResponse" + }, + { + "type": "null" + } + ] + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "SyntheticDataGeneration (Coming Soon)" + ], + "description": "", + "parameters": [ + { + "name": "job_uuid", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, + "/v1/synthetic-data-generation/job/status": { + "get": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "oneOf": [ + { + "$ref": "#/components/schemas/SyntheticDataGenerationJobStatusResponse" + }, + { + "type": "null" + } + ] + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "SyntheticDataGeneration (Coming Soon)" + ], + "description": "", + "parameters": [ + { + "name": "job_uuid", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, + "/v1/synthetic-data-generation/jobs": { + "get": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ListSyntheticDataGenerationJobsResponse" + } + } + } + }, + "400": { + "$ref": "#/components/responses/BadRequest400" + }, + "429": { + "$ref": "#/components/responses/TooManyRequests429" + }, + "500": { + "$ref": "#/components/responses/InternalServerError500" + }, + "default": { + "$ref": "#/components/responses/DefaultError" + } + }, + "tags": [ + "SyntheticDataGeneration (Coming Soon)" + ], + "description": "", + "parameters": [] + } + }, "/v1/tools/{tool_name}": { "get": { "responses": { @@ -3563,7 +3730,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/SyntheticDataGenerationResponse" + "$ref": "#/components/schemas/SyntheticDataGenerationJob" } } } @@ -4629,6 +4796,19 @@ "title": "CompletionResponse", "description": "Response from a completion request." }, + "CancelSyntheticDataGenerationJobRequest": { + "type": "object", + "properties": { + "job_uuid": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "job_uuid" + ], + "title": "CancelSyntheticDataGenerationJobRequest" + }, "CancelTrainingJobRequest": { "type": "object", "properties": { @@ -7347,6 +7527,196 @@ ], "title": "QuerySpanTreeResponse" }, + "SyntheticDataGenerationJobArtifactsResponse": { + "type": "object", + "properties": { + "job_uuid": { + "type": "string" + }, + "synthetic_data": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } + } + }, + "statistics": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } + } + }, + "additionalProperties": false, + "required": [ + "job_uuid", + "synthetic_data" + ], + "title": "SyntheticDataGenerationJobArtifactsResponse" + }, + "JobStatus": { + "type": "string", + "enum": [ + "completed", + "in_progress", + "failed", + "scheduled" + ], + "title": "JobStatus" + }, + "SyntheticDataGenerationJobStatusResponse": { + "type": "object", + "properties": { + "job_uuid": { + "type": "string" + }, + "status": { + "$ref": "#/components/schemas/JobStatus" + }, + "scheduled_at": { + "type": "string", + "format": "date-time" + }, + "started_at": { + "type": "string", + "format": "date-time" + }, + "completed_at": { + "type": "string", + "format": "date-time" + }, + "synthetic_data": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } + } + }, + "statistics": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + } + } + }, + "additionalProperties": false, + "required": [ + "job_uuid", + "status", + "synthetic_data" + ], + "title": "SyntheticDataGenerationJobStatusResponse", + "description": "Status of a synthetic data generation job." + }, + "SyntheticDataGenerationJob": { + "type": "object", + "properties": { + "job_uuid": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "job_uuid" + ], + "title": "SyntheticDataGenerationJob" + }, + "ListSyntheticDataGenerationJobsResponse": { + "type": "object", + "properties": { + "data": { + "type": "array", + "items": { + "$ref": "#/components/schemas/SyntheticDataGenerationJob" + } + } + }, + "additionalProperties": false, + "required": [ + "data" + ], + "title": "ListSyntheticDataGenerationJobsResponse" + }, "Tool": { "type": "object", "properties": { @@ -7533,16 +7903,6 @@ "title": "PostTrainingJobArtifactsResponse", "description": "Artifacts of a finetuning job." }, - "JobStatus": { - "type": "string", - "enum": [ - "completed", - "in_progress", - "failed", - "scheduled" - ], - "title": "JobStatus" - }, "PostTrainingJobStatusResponse": { "type": "object", "properties": { @@ -9781,70 +10141,6 @@ ], "title": "SyntheticDataGenerateRequest" }, - "SyntheticDataGenerationResponse": { - "type": "object", - "properties": { - "synthetic_data": { - "type": "array", - "items": { - "type": "object", - "additionalProperties": { - "oneOf": [ - { - "type": "null" - }, - { - "type": "boolean" - }, - { - "type": "number" - }, - { - "type": "string" - }, - { - "type": "array" - }, - { - "type": "object" - } - ] - } - } - }, - "statistics": { - "type": "object", - "additionalProperties": { - "oneOf": [ - { - "type": "null" - }, - { - "type": "boolean" - }, - { - "type": "number" - }, - { - "type": "string" - }, - { - "type": "array" - }, - { - "type": "object" - } - ] - } - } - }, - "additionalProperties": false, - "required": [ - "synthetic_data" - ], - "title": "SyntheticDataGenerationResponse", - "description": "Response from the synthetic data generation. Batch of (prompt, response, score) tuples that pass the threshold." - }, "VersionInfo": { "type": "object", "properties": { diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index d6001c00d..78a30757f 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -142,6 +142,31 @@ paths: schema: $ref: '#/components/schemas/BatchCompletionRequest' required: true + /v1/synthetic-data-generation/job/cancel: + post: + responses: + '200': + description: OK + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - SyntheticDataGeneration (Coming Soon) + description: '' + parameters: [] + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/CancelSyntheticDataGenerationJobRequest' + required: true /v1/post-training/job/cancel: post: responses: @@ -1006,6 +1031,89 @@ paths: schema: $ref: '#/components/schemas/GetSpanTreeRequest' required: true + /v1/synthetic-data-generation/job/artifacts: + get: + responses: + '200': + description: OK + content: + application/json: + schema: + oneOf: + - $ref: '#/components/schemas/SyntheticDataGenerationJobArtifactsResponse' + - type: 'null' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - SyntheticDataGeneration (Coming Soon) + description: '' + parameters: + - name: job_uuid + in: query + required: true + schema: + type: string + /v1/synthetic-data-generation/job/status: + get: + responses: + '200': + description: OK + content: + application/json: + schema: + oneOf: + - $ref: '#/components/schemas/SyntheticDataGenerationJobStatusResponse' + - type: 'null' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - SyntheticDataGeneration (Coming Soon) + description: '' + parameters: + - name: job_uuid + in: query + required: true + schema: + type: string + /v1/synthetic-data-generation/jobs: + get: + responses: + '200': + description: OK + content: + application/json: + schema: + $ref: '#/components/schemas/ListSyntheticDataGenerationJobsResponse' + '400': + $ref: '#/components/responses/BadRequest400' + '429': + $ref: >- + #/components/responses/TooManyRequests429 + '500': + $ref: >- + #/components/responses/InternalServerError500 + default: + $ref: '#/components/responses/DefaultError' + tags: + - SyntheticDataGeneration (Coming Soon) + description: '' + parameters: [] /v1/tools/{tool_name}: get: responses: @@ -2420,7 +2528,7 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/SyntheticDataGenerationResponse' + $ref: '#/components/schemas/SyntheticDataGenerationJob' '400': $ref: '#/components/responses/BadRequest400' '429': @@ -3155,6 +3263,15 @@ components: - stop_reason title: CompletionResponse description: Response from a completion request. + CancelSyntheticDataGenerationJobRequest: + type: object + properties: + job_uuid: + type: string + additionalProperties: false + required: + - job_uuid + title: CancelSyntheticDataGenerationJobRequest CancelTrainingJobRequest: type: object properties: @@ -5042,6 +5159,113 @@ components: required: - data title: QuerySpanTreeResponse + "SyntheticDataGenerationJobArtifactsResponse": + type: object + properties: + job_uuid: + type: string + synthetic_data: + type: array + items: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + statistics: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + additionalProperties: false + required: + - job_uuid + - synthetic_data + title: >- + SyntheticDataGenerationJobArtifactsResponse + JobStatus: + type: string + enum: + - completed + - in_progress + - failed + - scheduled + title: JobStatus + SyntheticDataGenerationJobStatusResponse: + type: object + properties: + job_uuid: + type: string + status: + $ref: '#/components/schemas/JobStatus' + scheduled_at: + type: string + format: date-time + started_at: + type: string + format: date-time + completed_at: + type: string + format: date-time + synthetic_data: + type: array + items: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + statistics: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + additionalProperties: false + required: + - job_uuid + - status + - synthetic_data + title: SyntheticDataGenerationJobStatusResponse + description: >- + Status of a synthetic data generation job. + SyntheticDataGenerationJob: + type: object + properties: + job_uuid: + type: string + additionalProperties: false + required: + - job_uuid + title: SyntheticDataGenerationJob + ListSyntheticDataGenerationJobsResponse: + type: object + properties: + data: + type: array + items: + $ref: '#/components/schemas/SyntheticDataGenerationJob' + additionalProperties: false + required: + - data + title: ListSyntheticDataGenerationJobsResponse Tool: type: object properties: @@ -5162,14 +5386,6 @@ components: - checkpoints title: PostTrainingJobArtifactsResponse description: Artifacts of a finetuning job. - JobStatus: - type: string - enum: - - completed - - in_progress - - failed - - scheduled - title: JobStatus PostTrainingJobStatusResponse: type: object properties: @@ -6577,38 +6793,6 @@ components: - dialogs - filtering_function title: SyntheticDataGenerateRequest - SyntheticDataGenerationResponse: - type: object - properties: - synthetic_data: - type: array - items: - type: object - additionalProperties: - oneOf: - - type: 'null' - - type: boolean - - type: number - - type: string - - type: array - - type: object - statistics: - type: object - additionalProperties: - oneOf: - - type: 'null' - - type: boolean - - type: number - - type: string - - type: array - - type: object - additionalProperties: false - required: - - synthetic_data - title: SyntheticDataGenerationResponse - description: >- - Response from the synthetic data generation. Batch of (prompt, response, score) - tuples that pass the threshold. VersionInfo: type: object properties: diff --git a/llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py b/llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py index 7b41192af..bbe728b87 100644 --- a/llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py +++ b/llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py @@ -4,11 +4,13 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. +from datetime import datetime from enum import Enum -from typing import Any, Dict, List, Optional, Protocol, Union +from typing import Any, Dict, List, Optional, Protocol from pydantic import BaseModel +from llama_stack.apis.common.job_types import JobStatus from llama_stack.apis.inference import Message from llama_stack.schema_utils import json_schema_type, webmethod @@ -34,18 +36,58 @@ class SyntheticDataGenerationRequest(BaseModel): @json_schema_type -class SyntheticDataGenerationResponse(BaseModel): - """Response from the synthetic data generation. Batch of (prompt, response, score) tuples that pass the threshold.""" +class SyntheticDataGenerationJob(BaseModel): + job_uuid: str + + +@json_schema_type +class SyntheticDataGenerationJobStatusResponse(BaseModel): + """Status of a synthetic data generation job.""" + + job_uuid: str + status: JobStatus + + scheduled_at: Optional[datetime] = None + started_at: Optional[datetime] = None + completed_at: Optional[datetime] = None + + synthetic_data: List[Dict[str, Any]] + statistics: Optional[Dict[str, Any]] = None + + +class ListSyntheticDataGenerationJobsResponse(BaseModel): + data: List[SyntheticDataGenerationJob] + + +@json_schema_type +class SyntheticDataGenerationJobArtifactsResponse(BaseModel): + job_uuid: str synthetic_data: List[Dict[str, Any]] statistics: Optional[Dict[str, Any]] = None class SyntheticDataGeneration(Protocol): - @webmethod(route="/synthetic-data-generation/generate") + @webmethod(route="/synthetic-data-generation/generate", method="POST") def synthetic_data_generate( self, dialogs: List[Message], filtering_function: FilteringFunction = FilteringFunction.none, model: Optional[str] = None, - ) -> Union[SyntheticDataGenerationResponse]: ... + ) -> SyntheticDataGenerationJob: ... + + @webmethod(route="/synthetic-data-generation/jobs", method="GET") + async def get_synthetic_data_generation_jobs(self) -> ListSyntheticDataGenerationJobsResponse: ... + + @webmethod(route="/synthetic-data-generation/job/status", method="GET") + async def get_synthetic_data_generation_job_status( + self, job_uuid: str + ) -> Optional[SyntheticDataGenerationJobStatusResponse]: ... + + @webmethod(route="/synthetic-data-generation/job/cancel", method="POST") + async def cancel_synthetic_data_generation_job(self, job_uuid: str) -> None: ... + + @webmethod(route="/synthetic-data-generation/job/artifacts", method="GET") + async def get_synthetic_data_generation_job_artifacts( + self, job_uuid: str + ) -> Optional[SyntheticDataGenerationJobArtifactsResponse]: ...