add nvidia distribution

llama_stack/templates/nvidia/build.yaml

@@ -20,6 +20,8 @@ distribution_spec:
     - inline::basic
     - inline::llm-as-judge
     - inline::braintrust
+    post_training:
+    - remote::nvidia
     tool_runtime:
     - inline::rag-runtime
 image_type: conda

llama_stack/templates/nvidia/nvidia.py

@@ -9,6 +9,7 @@ from pathlib import Path
 from llama_stack.distribution.datatypes import ModelInput, Provider, ShieldInput, ToolGroupInput
 from llama_stack.providers.remote.inference.nvidia import NVIDIAConfig
 from llama_stack.providers.remote.inference.nvidia.models import MODEL_ENTRIES
+from llama_stack.providers.remote.post_training.nvidia import NvidiaPostTrainingConfig
 from llama_stack.providers.remote.safety.nvidia import NVIDIASafetyConfig
 from llama_stack.templates.template import DistributionTemplate, RunConfigSettings, get_model_registry
 
@@ -18,6 +19,7 @@ def get_distribution_template() -> DistributionTemplate:
         "inference": ["remote::nvidia"],
         "vector_io": ["inline::faiss"],
         "safety": ["remote::nvidia"],
+        "post_training": ["remote::nvidia"],
         "agents": ["inline::meta-reference"],
         "telemetry": ["inline::meta-reference"],
         "eval": ["inline::meta-reference"],
@@ -31,6 +33,12 @@ def get_distribution_template() -> DistributionTemplate:
         provider_type="remote::nvidia",
         config=NVIDIAConfig.sample_run_config(),
     )
+
+    post_training_provider = Provider(
+        provider_id="nvidia",
+        provider_type="remote::nvidia",
+        config=NvidiaPostTrainingConfig.sample_run_config(),
+    )
     safety_provider = Provider(
         provider_id="nvidia",
         provider_type="remote::nvidia",
@@ -89,6 +97,31 @@ def get_distribution_template() -> DistributionTemplate:
                 "",
                 "NVIDIA API Key",
             ),
+            ## Nemo Customizer related variables
+            "NVIDIA_USER_ID": (
+                "llama-stack-user",
+                "NVIDIA User ID",
+            ),
+            "NVIDIA_DATASET_NAMESPACE": (
+                "default",
+                "NVIDIA Dataset Namespace",
+            ),
+            "NVIDIA_ACCESS_POLICIES": (
+                "{}",
+                "NVIDIA Access Policies",
+            ),
+            "NVIDIA_PROJECT_ID": (
+                "test-project",
+                "NVIDIA Project ID",
+            ),
+            "NVIDIA_CUSTOMIZER_URL": (
+                "https://customizer.api.nvidia.com",
+                "NVIDIA Customizer URL",
+            ),
+            "NVIDIA_OUTPUT_MODEL_DIR": (
+                "test-example-model@v1",
+                "NVIDIA Output Model Directory",
+            ),
             "GUARDRAILS_SERVICE_URL": (
                 "http://0.0.0.0:7331",
                 "URL for the NeMo Guardrails Service",

llama_stack/templates/nvidia/run.yaml

@@ -8,6 +8,7 @@ apis:
 - safety
 - scoring
 - telemetry
+- post_training
 - tool_runtime
 - vector_io
 providers:
@@ -73,6 +74,10 @@ providers:
     provider_type: inline::braintrust
     config:
       openai_api_key: ${env.OPENAI_API_KEY:}
+  post_training:
+  - provider_id: nvidia-customizer
+    provider_type: remote::nvidia
+    config: {}
   tool_runtime:
   - provider_id: rag-runtime
     provider_type: inline::rag-runtime