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 ( from finetuning_types import (
Checkpoint,
Dataset, Dataset,
DoraFinetuningConfig, DoraFinetuningConfig,
FinetuningAlgorithm, FinetuningAlgorithm,
@ -233,6 +234,7 @@ class MessageScore:
message: Message message: Message
score: float score: float
@json_schema_type @json_schema_type
@dataclass @dataclass
class KScoredPromptGenerations: class KScoredPromptGenerations:
@ -375,6 +377,19 @@ class FinetuningJobStatusResponse:
resources_allocated: Optional[Dict[str, Any]] = None 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): class Finetuning(Protocol):
@webmethod(route="/finetuning/text_generation/train") @webmethod(route="/finetuning/text_generation/train")
@ -393,6 +408,11 @@ class Finetuning(Protocol):
@webmethod(route="/finetuning/job/cancel") @webmethod(route="/finetuning/job/cancel")
def cancel_training_job(self, job_uuid: str) -> None: ... 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( class LlamaStackEndpoints(
Inference, Inference,

View file

@ -96,3 +96,9 @@ class FinetuningJobStatus(Enum):
completed = "completed" completed = "completed"
failed = "failed" failed = "failed"
scheduled = "scheduled" 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": { "/finetuning/job/status": {
"get": { "get": {
"responses": { "responses": {
@ -1276,6 +1305,39 @@
], ],
"title": "Dataset to be used for training or evaluating language models." "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": { "FinetuningJobStatusResponse": {
"type": "object", "type": "object",
"properties": { "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, "additionalProperties": false,
"required": [ "required": [
"job_uuid", "job_uuid",
"status" "status",
"checkpoints"
], ],
"title": "Status of a finetuning job." "title": "Status of a finetuning job."
}, },
@ -2316,13 +2398,7 @@
], ],
"tags": [ "tags": [
{ {
"name": "RewardScoring" "name": "Finetuning"
},
{
"name": "Inference"
},
{
"name": "AgenticSystem"
}, },
{ {
"name": "Datasets" "name": "Datasets"
@ -2331,7 +2407,13 @@
"name": "SyntheticDataGeneration" "name": "SyntheticDataGeneration"
}, },
{ {
"name": "Finetuning" "name": "RewardScoring"
},
{
"name": "AgenticSystem"
},
{
"name": "Inference"
}, },
{ {
"name": "ShieldConfig", "name": "ShieldConfig",
@ -2381,6 +2463,10 @@
"name": "Dataset", "name": "Dataset",
"description": "Dataset to be used for training or evaluating language models.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/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", "name": "FinetuningJobStatusResponse",
"description": "Status of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/FinetuningJobStatusResponse\" />" "description": "Status of a finetuning job.\n\n<SchemaDefinition schemaRef=\"#/components/schemas/FinetuningJobStatusResponse\" />"
@ -2484,6 +2570,7 @@
"CompletionResponseStreamChunk", "CompletionResponseStreamChunk",
"CreateDatasetRequest", "CreateDatasetRequest",
"Dataset", "Dataset",
"FinetuningJobArtifactsResponse",
"FinetuningJobLogStream", "FinetuningJobLogStream",
"FinetuningJobStatusResponse", "FinetuningJobStatusResponse",
"FinetuningTrainRequest", "FinetuningTrainRequest",

View file

@ -747,6 +747,29 @@ components:
- metadata - metadata
title: Dataset to be used for training or evaluating language models. title: Dataset to be used for training or evaluating language models.
type: object 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: FinetuningJobLogStream:
additionalProperties: false additionalProperties: false
properties: properties:
@ -764,6 +787,19 @@ components:
FinetuningJobStatusResponse: FinetuningJobStatusResponse:
additionalProperties: false additionalProperties: false
properties: properties:
checkpoints:
items:
additionalProperties: false
properties:
iters:
type: integer
path:
$ref: '#/components/schemas/URL'
required:
- iters
- path
type: object
type: array
completed_at: completed_at:
format: date-time format: date-time
type: string type: string
@ -795,6 +831,7 @@ components:
required: required:
- job_uuid - job_uuid
- status - status
- checkpoints
title: Status of a finetuning job. title: Status of a finetuning job.
type: object type: object
FinetuningTrainRequest: FinetuningTrainRequest:
@ -1325,6 +1362,23 @@ paths:
description: OK description: OK
tags: tags:
- Datasets - 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: /finetuning/job/logs:
get: get:
parameters: parameters:
@ -1414,12 +1468,12 @@ security:
servers: servers:
- url: http://llama.meta.com - url: http://llama.meta.com
tags: tags:
- name: RewardScoring - name: Finetuning
- name: Inference
- name: AgenticSystem
- name: Datasets - name: Datasets
- name: SyntheticDataGeneration - name: SyntheticDataGeneration
- name: Finetuning - name: RewardScoring
- name: AgenticSystem
- name: Inference
- description: <SchemaDefinition schemaRef="#/components/schemas/ShieldConfig" /> - description: <SchemaDefinition schemaRef="#/components/schemas/ShieldConfig" />
name: ShieldConfig name: ShieldConfig
- description: <SchemaDefinition schemaRef="#/components/schemas/AgenticSystemCreateRequest" - description: <SchemaDefinition schemaRef="#/components/schemas/AgenticSystemCreateRequest"
@ -1468,6 +1522,12 @@ tags:
<SchemaDefinition schemaRef="#/components/schemas/Dataset" />' <SchemaDefinition schemaRef="#/components/schemas/Dataset" />'
name: Dataset name: Dataset
- description: 'Artifacts of a finetuning job.
<SchemaDefinition schemaRef="#/components/schemas/FinetuningJobArtifactsResponse"
/>'
name: FinetuningJobArtifactsResponse
- description: 'Status of a finetuning job. - description: 'Status of a finetuning job.
@ -1581,6 +1641,7 @@ x-tagGroups:
- CompletionResponseStreamChunk - CompletionResponseStreamChunk
- CreateDatasetRequest - CreateDatasetRequest
- Dataset - Dataset
- FinetuningJobArtifactsResponse
- FinetuningJobLogStream - FinetuningJobLogStream
- FinetuningJobStatusResponse - FinetuningJobStatusResponse
- FinetuningTrainRequest - FinetuningTrainRequest