diff --git a/source/api_definitions.py b/source/api_definitions.py index a02536453..e5697f86f 100644 --- a/source/api_definitions.py +++ b/source/api_definitions.py @@ -12,6 +12,7 @@ from agentic_system_types import ( ) from finetuning_types import ( + Checkpoint, Dataset, DoraFinetuningConfig, FinetuningAlgorithm, @@ -233,6 +234,7 @@ class MessageScore: message: Message score: float + @json_schema_type @dataclass class KScoredPromptGenerations: @@ -375,6 +377,19 @@ class FinetuningJobStatusResponse: resources_allocated: Optional[Dict[str, Any]] = None + checkpoints: List[Checkpoint] = field(default_factory=list) + + +@json_schema_type +@dataclass +class FinetuningJobArtifactsResponse: + """Artifacts of a finetuning job.""" + + job_uuid: str + checkpoints: List[Checkpoint] = field(default_factory=list) + + # TODO(ashwin): metrics, evals + class Finetuning(Protocol): @webmethod(route="/finetuning/text_generation/train") @@ -393,6 +408,11 @@ class Finetuning(Protocol): @webmethod(route="/finetuning/job/cancel") def cancel_training_job(self, job_uuid: str) -> None: ... + @webmethod(route="/finetuning/job/artifacts") + def get_training_job_artifacts( + self, job_uuid: str + ) -> FinetuningJobArtifactsResponse: ... + class LlamaStackEndpoints( Inference, diff --git a/source/finetuning_types.py b/source/finetuning_types.py index 99c3f112d..9aa897eac 100644 --- a/source/finetuning_types.py +++ b/source/finetuning_types.py @@ -96,3 +96,9 @@ class FinetuningJobStatus(Enum): completed = "completed" failed = "failed" scheduled = "scheduled" + + +@dataclass +class Checkpoint: + iters: int + path: URL diff --git a/source/openapi.html b/source/openapi.html index 6e51f862a..605ff235a 100644 --- a/source/openapi.html +++ b/source/openapi.html @@ -192,6 +192,35 @@ ] } }, + "/finetuning/job/artifacts": { + "get": { + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/FinetuningJobArtifactsResponse" + } + } + } + } + }, + "tags": [ + "Finetuning" + ], + "parameters": [ + { + "name": "job_uuid", + "in": "query", + "required": true, + "schema": { + "type": "string" + } + } + ] + } + }, "/finetuning/job/status": { "get": { "responses": { @@ -1276,6 +1305,39 @@ ], "title": "Dataset to be used for training or evaluating language models." }, + "FinetuningJobArtifactsResponse": { + "type": "object", + "properties": { + "job_uuid": { + "type": "string" + }, + "checkpoints": { + "type": "array", + "items": { + "type": "object", + "properties": { + "iters": { + "type": "integer" + }, + "path": { + "$ref": "#/components/schemas/URL" + } + }, + "additionalProperties": false, + "required": [ + "iters", + "path" + ] + } + } + }, + "additionalProperties": false, + "required": [ + "job_uuid", + "checkpoints" + ], + "title": "Artifacts of a finetuning job." + }, "FinetuningJobStatusResponse": { "type": "object", "properties": { @@ -1327,12 +1389,32 @@ } ] } + }, + "checkpoints": { + "type": "array", + "items": { + "type": "object", + "properties": { + "iters": { + "type": "integer" + }, + "path": { + "$ref": "#/components/schemas/URL" + } + }, + "additionalProperties": false, + "required": [ + "iters", + "path" + ] + } } }, "additionalProperties": false, "required": [ "job_uuid", - "status" + "status", + "checkpoints" ], "title": "Status of a finetuning job." }, @@ -2316,13 +2398,7 @@ ], "tags": [ { - "name": "RewardScoring" - }, - { - "name": "Inference" - }, - { - "name": "AgenticSystem" + "name": "Finetuning" }, { "name": "Datasets" @@ -2331,7 +2407,13 @@ "name": "SyntheticDataGeneration" }, { - "name": "Finetuning" + "name": "RewardScoring" + }, + { + "name": "AgenticSystem" + }, + { + "name": "Inference" }, { "name": "ShieldConfig", @@ -2381,6 +2463,10 @@ "name": "Dataset", "description": "Dataset to be used for training or evaluating language models.\n\n" }, + { + "name": "FinetuningJobArtifactsResponse", + "description": "Artifacts of a finetuning job.\n\n" + }, { "name": "FinetuningJobStatusResponse", "description": "Status of a finetuning job.\n\n" @@ -2484,6 +2570,7 @@ "CompletionResponseStreamChunk", "CreateDatasetRequest", "Dataset", + "FinetuningJobArtifactsResponse", "FinetuningJobLogStream", "FinetuningJobStatusResponse", "FinetuningTrainRequest", diff --git a/source/openapi.yaml b/source/openapi.yaml index 9acc76086..a0e592e12 100644 --- a/source/openapi.yaml +++ b/source/openapi.yaml @@ -747,6 +747,29 @@ components: - metadata title: Dataset to be used for training or evaluating language models. type: object + FinetuningJobArtifactsResponse: + additionalProperties: false + properties: + checkpoints: + items: + additionalProperties: false + properties: + iters: + type: integer + path: + $ref: '#/components/schemas/URL' + required: + - iters + - path + type: object + type: array + job_uuid: + type: string + required: + - job_uuid + - checkpoints + title: Artifacts of a finetuning job. + type: object FinetuningJobLogStream: additionalProperties: false properties: @@ -764,6 +787,19 @@ components: FinetuningJobStatusResponse: additionalProperties: false properties: + checkpoints: + items: + additionalProperties: false + properties: + iters: + type: integer + path: + $ref: '#/components/schemas/URL' + required: + - iters + - path + type: object + type: array completed_at: format: date-time type: string @@ -795,6 +831,7 @@ components: required: - job_uuid - status + - checkpoints title: Status of a finetuning job. type: object FinetuningTrainRequest: @@ -1325,6 +1362,23 @@ paths: description: OK tags: - Datasets + /finetuning/job/artifacts: + get: + parameters: + - in: query + name: job_uuid + required: true + schema: + type: string + responses: + '200': + content: + application/json: + schema: + $ref: '#/components/schemas/FinetuningJobArtifactsResponse' + description: OK + tags: + - Finetuning /finetuning/job/logs: get: parameters: @@ -1414,12 +1468,12 @@ security: servers: - url: http://llama.meta.com tags: -- name: RewardScoring -- name: Inference -- name: AgenticSystem +- name: Finetuning - name: Datasets - name: SyntheticDataGeneration -- name: Finetuning +- name: RewardScoring +- name: AgenticSystem +- name: Inference - description: name: ShieldConfig - description: ' name: Dataset +- description: 'Artifacts of a finetuning job. + + + ' + name: FinetuningJobArtifactsResponse - description: 'Status of a finetuning job. @@ -1581,6 +1641,7 @@ x-tagGroups: - CompletionResponseStreamChunk - CreateDatasetRequest - Dataset + - FinetuningJobArtifactsResponse - FinetuningJobLogStream - FinetuningJobStatusResponse - FinetuningTrainRequest