some finetuning changes

This commit is contained in:
Ashwin Bharambe 2024-07-10 22:08:42 -07:00
parent 6ec7c47938
commit c1f6816d76
4 changed files with 187 additions and 13 deletions

View file

@ -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,

View file

@ -96,3 +96,9 @@ class FinetuningJobStatus(Enum):
completed = "completed"
failed = "failed"
scheduled = "scheduled"
@dataclass
class Checkpoint:
iters: int
path: URL

View file

@ -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<SchemaDefinition schemaRef=\"#/components/schemas/Dataset\" />"
},
{
"name": "FinetuningJobArtifactsResponse",
"description": "Artifacts of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/FinetuningJobArtifactsResponse\" />"
},
{
"name": "FinetuningJobStatusResponse",
"description": "Status of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/FinetuningJobStatusResponse\" />"
@ -2484,6 +2570,7 @@
"CompletionResponseStreamChunk",
"CreateDatasetRequest",
"Dataset",
"FinetuningJobArtifactsResponse",
"FinetuningJobLogStream",
"FinetuningJobStatusResponse",
"FinetuningTrainRequest",

View file

@ -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: <SchemaDefinition schemaRef="#/components/schemas/ShieldConfig" />
name: ShieldConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/AgenticSystemCreateRequest"
@ -1468,6 +1522,12 @@ tags:
<SchemaDefinition schemaRef="#/components/schemas/Dataset" />'
name: Dataset
- description: 'Artifacts of a finetuning job.
<SchemaDefinition schemaRef="#/components/schemas/FinetuningJobArtifactsResponse"
/>'
name: FinetuningJobArtifactsResponse
- description: 'Status of a finetuning job.
@ -1581,6 +1641,7 @@ x-tagGroups:
- CompletionResponseStreamChunk
- CreateDatasetRequest
- Dataset
- FinetuningJobArtifactsResponse
- FinetuningJobLogStream
- FinetuningJobStatusResponse
- FinetuningTrainRequest