feat: associated models API with post_training

there are likely scenarios where admins of a stack only want to allow clients to fine-tune certain models, register certain models to be fine-tuned. etc
introduce the post_training router and post_training_models as the associated type. A different model type needs to be used for inference vs post_training due to the structure of the router currently.

Signed-off-by: Charlie Doern <cdoern@redhat.com>

llama_stack/distribution/routers/__init__.py

@@ -21,7 +21,8 @@ async def get_routing_table_impl(
 ) -> Any:
     from ..routing_tables.benchmarks import BenchmarksRoutingTable
     from ..routing_tables.datasets import DatasetsRoutingTable
-    from ..routing_tables.models import ModelsRoutingTable
+    from ..routing_tables.models import InferenceModelsRoutingTable
+    from ..routing_tables.post_training_models import PostTrainingModelsRoutingTable
     from ..routing_tables.scoring_functions import ScoringFunctionsRoutingTable
     from ..routing_tables.shields import ShieldsRoutingTable
     from ..routing_tables.toolgroups import ToolGroupsRoutingTable
@@ -29,7 +30,8 @@ async def get_routing_table_impl(
 
     api_to_tables = {
         "vector_dbs": VectorDBsRoutingTable,
-        "models": ModelsRoutingTable,
+        "models": InferenceModelsRoutingTable,
+        "post_training_models": PostTrainingModelsRoutingTable,
         "shields": ShieldsRoutingTable,
         "datasets": DatasetsRoutingTable,
         "scoring_functions": ScoringFunctionsRoutingTable,
@@ -40,7 +42,12 @@ async def get_routing_table_impl(
     if api.value not in api_to_tables:
         raise ValueError(f"API {api.value} not found in router map")
 
-    impl = api_to_tables[api.value](impls_by_provider_id, dist_registry)
+    # For post-training API, we want to use the post-training models routing table
+    if api == Api.post_training:
+        impl = PostTrainingModelsRoutingTable(impls_by_provider_id, dist_registry)
+    else:
+        impl = api_to_tables[api.value](impls_by_provider_id, dist_registry)
+
     await impl.initialize()
     return impl
 
@@ -51,6 +58,7 @@ async def get_auto_router_impl(
     from .datasets import DatasetIORouter
     from .eval_scoring import EvalRouter, ScoringRouter
     from .inference import InferenceRouter
+    from .post_training import PostTrainingRouter
     from .safety import SafetyRouter
     from .tool_runtime import ToolRuntimeRouter
     from .vector_io import VectorIORouter
@@ -63,6 +71,7 @@ async def get_auto_router_impl(
         "scoring": ScoringRouter,
         "eval": EvalRouter,
         "tool_runtime": ToolRuntimeRouter,
+        "post_training": PostTrainingRouter,
     }
     api_to_deps = {
         "inference": {"telemetry": Api.telemetry},

llama_stack/distribution/routers/post_training.py

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+from llama_stack.apis.models import Model
+from llama_stack.apis.post_training import (
+    AlgorithmConfig,
+    DPOAlignmentConfig,
+    ListPostTrainingJobsResponse,
+    PostTraining,
+    PostTrainingJob,
+    PostTrainingJobArtifactsResponse,
+    PostTrainingJobStatusResponse,
+    TrainingConfig,
+)
+from llama_stack.log import get_logger
+from llama_stack.providers.datatypes import RoutingTable
+
+logger = get_logger(name=__name__, category="core")
+
+
+class PostTrainingRouter(PostTraining):
+    """Routes to an provider based on the model"""
+
+    async def initialize(self) -> None:
+        pass
+
+    def __init__(
+        self,
+        routing_table: RoutingTable,
+    ) -> None:
+        logger.debug("Initializing InferenceRouter")
+        self.routing_table = routing_table
+
+    async def supervised_fine_tune(
+        self,
+        job_uuid: str,
+        training_config: TrainingConfig,
+        hyperparam_search_config: dict[str, Any],
+        logger_config: dict[str, Any],
+        model: str,
+        checkpoint_dir: str | None = None,
+        algorithm_config: AlgorithmConfig | None = None,
+    ) -> PostTrainingJob:
+        provider = self.routing_table.get_provider_impl(model)
+        params = dict(
+            job_uuid=job_uuid,
+            training_config=training_config,
+            hyperparam_search_config=hyperparam_search_config,
+            logger_config=logger_config,
+            model=model,
+            checkpoint_dir=checkpoint_dir,
+            algorithm_config=algorithm_config,
+        )
+        return provider.supervised_fine_tune(**params)
+
+    async def register_model(self, model: Model) -> Model:
+        try:
+            # get static list of models
+            model = await self.register_helper.register_model(model)
+        except ValueError:
+            # if model is NOT in the list, its probably ok, but warn the user.
+            #
+            logger.warning(
+                f"Model {model.identifier} is not in the model registry for this provider, there might be unexpected issues."
+            )
+        if model.provider_resource_id is None:
+            raise ValueError("Model provider_resource_id cannot be None")
+        provider_resource_id = self.register_helper.get_provider_model_id(model.provider_resource_id)
+        if provider_resource_id is None:
+            provider_resource_id = model.provider_resource_id
+        model.provider_resource_id = provider_resource_id
+
+        return model
+
+    async def preference_optimize(
+        self,
+        job_uuid: str,
+        finetuned_model: str,
+        algorithm_config: DPOAlignmentConfig,
+        training_config: TrainingConfig,
+        hyperparam_search_config: dict[str, Any],
+        logger_config: dict[str, Any],
+    ) -> PostTrainingJob:
+        pass
+
+    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
+        pass
+
+    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse | None:
+        pass
+
+    async def cancel_training_job(self, job_uuid: str) -> None:
+        pass
+
+    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
+        pass