fix: Restore discriminator for AlgorithmConfig

It was lost during refactoring in #1548.

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>

docs/_static/llama-stack-spec.html

@@ -9863,6 +9863,23 @@
                 ],
                 "title": "ScoreBatchResponse"
             },
+            "AlgorithmConfig": {
+                "oneOf": [
+                    {
+                        "$ref": "#/components/schemas/LoraFinetuningConfig"
+                    },
+                    {
+                        "$ref": "#/components/schemas/QATFinetuningConfig"
+                    }
+                ],
+                "discriminator": {
+                    "propertyName": "type",
+                    "mapping": {
+                        "LoRA": "#/components/schemas/LoraFinetuningConfig",
+                        "QAT": "#/components/schemas/QATFinetuningConfig"
+                    }
+                }
+            },
             "LoraFinetuningConfig": {
                 "type": "object",
                 "properties": {
@@ -9998,14 +10015,7 @@
                         "type": "string"
                     },
                     "algorithm_config": {
-                        "oneOf": [
+                        "$ref": "#/components/schemas/AlgorithmConfig"
-                            {
-                                "$ref": "#/components/schemas/LoraFinetuningConfig"
-                            },
-                            {
-                                "$ref": "#/components/schemas/QATFinetuningConfig"
-                            }
-                        ]
                     }
                 },
                 "additionalProperties": false,

docs/_static/llama-stack-spec.yaml

@@ -6689,6 +6689,15 @@ components:
       required:
         - results
       title: ScoreBatchResponse
+    AlgorithmConfig:
+      oneOf:
+        - $ref: '#/components/schemas/LoraFinetuningConfig'
+        - $ref: '#/components/schemas/QATFinetuningConfig'
+      discriminator:
+        propertyName: type
+        mapping:
+          LoRA: '#/components/schemas/LoraFinetuningConfig'
+          QAT: '#/components/schemas/QATFinetuningConfig'
     LoraFinetuningConfig:
       type: object
       properties:
@@ -6772,9 +6781,7 @@ components:
         checkpoint_dir:
           type: string
         algorithm_config:
-          oneOf:
+          $ref: '#/components/schemas/AlgorithmConfig'
-            - $ref: '#/components/schemas/LoraFinetuningConfig'
-            - $ref: '#/components/schemas/QATFinetuningConfig'
       additionalProperties: false
       required:
         - job_uuid

llama_stack/apis/post_training/post_training.py

@@ -6,7 +6,7 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Protocol
+from typing import Any, Dict, List, Literal, Optional, Protocol, Union
 
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
@@ -88,7 +88,7 @@ class QATFinetuningConfig(BaseModel):
     group_size: int
 
 
-AlgorithmConfig = Annotated[LoraFinetuningConfig | QATFinetuningConfig, Field(discriminator="type")]
+AlgorithmConfig = Annotated[Union[LoraFinetuningConfig, QATFinetuningConfig], Field(discriminator="type")]
 register_schema(AlgorithmConfig, name="AlgorithmConfig")
 
 
@@ -182,7 +182,7 @@ class PostTraining(Protocol):
             description="Model descriptor from `llama model list`",
         ),
         checkpoint_dir: Optional[str] = None,
-        algorithm_config: Optional[LoraFinetuningConfig | QATFinetuningConfig] = None,
+        algorithm_config: Optional[AlgorithmConfig] = None,
     ) -> PostTrainingJob: ...
 
     @webmethod(route="/post-training/preference-optimize", method="POST")