phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-10 04:08:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

post training job api

Browse source
This commit is contained in:
Xi Yan 2025-03-12 01:24:00 -07:00
parent deb24f5abe
commit 961528eae1
1 changed files with 11 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

47
llama_stack/apis/post_training/post_training.py
View file

@ -12,7 +12,7 @@ from pydantic import BaseModel, Field
from typing_extensions import Annotated
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.common.job_types import JobStatus
from llama_stack.apis.common.job_types import JobCommonFields, JobStatus
from llama_stack.apis.common.training_types import Checkpoint
from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
@ -89,7 +89,9 @@ class QATFinetuningConfig(BaseModel):
AlgorithmConfig = register_schema(
Annotated[Union[LoraFinetuningConfig, QATFinetuningConfig], Field(discriminator="type")],
Annotated[
Union[LoraFinetuningConfig, QATFinetuningConfig], Field(discriminator="type")
],
name="AlgorithmConfig",
)
@ -137,23 +139,9 @@ class PostTrainingRLHFRequest(BaseModel):
logger_config: Dict[str, Any]
class PostTrainingJob(BaseModel):
job_uuid: str
@json_schema_type
class PostTrainingJobStatusResponse(BaseModel):
"""Status of a finetuning job."""
job_uuid: str
status: JobStatus
scheduled_at: Optional[datetime] = None
started_at: Optional[datetime] = None
completed_at: Optional[datetime] = None
class PostTrainingJob(JobCommonFields):
resources_allocated: Optional[Dict[str, Any]] = None
checkpoints: List[Checkpoint] = Field(default_factory=list)
@ -161,18 +149,8 @@ class ListPostTrainingJobsResponse(BaseModel):
data: List[PostTrainingJob]
@json_schema_type
class PostTrainingJobArtifactsResponse(BaseModel):
"""Artifacts of a finetuning job."""
job_uuid: str
checkpoints: List[Checkpoint] = Field(default_factory=list)
# TODO(ashwin): metrics, evals
class PostTraining(Protocol):
@webmethod(route="/post-training/supervised-fine-tune", method="POST")
@webmethod(route="/post-training/supervised-fine-tune/jobs", method="POST")
async def supervised_fine_tune(
self,
job_uuid: str,
@ -187,7 +165,7 @@ class PostTraining(Protocol):
algorithm_config: Optional[AlgorithmConfig] = None,
) -> PostTrainingJob: ...
@webmethod(route="/post-training/preference-optimize", method="POST")
@webmethod(route="/post-training/preference-optimize/jobs", method="POST")
async def preference_optimize(
self,
job_uuid: str,
@ -201,11 +179,8 @@ class PostTraining(Protocol):
@webmethod(route="/post-training/jobs", method="GET")
async def get_training_jobs(self) -> ListPostTrainingJobsResponse: ...
@webmethod(route="/post-training/job/status", method="GET")
async def get_training_job_status(self, job_uuid: str) -> Optional[PostTrainingJobStatusResponse]: ...
@webmethod(route="/post-training/jobs/{job_id}", method="GET")
async def get_training_job(self, job_id: str) -> Optional[PostTrainingJob]: ...
@webmethod(route="/post-training/job/cancel", method="POST")
async def cancel_training_job(self, job_uuid: str) -> None: ...
@webmethod(route="/post-training/job/artifacts", method="GET")
async def get_training_job_artifacts(self, job_uuid: str) -> Optional[PostTrainingJobArtifactsResponse]: ...
@webmethod(route="/post-training/jobs/{job_id}", method="DELETE")
async def cancel_training_job(self, job_id: str) -> None: ...