Deploy NVIDIA NIM for self-hosted Llama 4 with LiteLLM routing

Loading diagram...

• Line 1: Declares a Mermaid flowchart with left-to-right (LR) layout direction.
• Line 2: Defines node A (FastAPI Application) sending an OpenAI-compatible request to node B (LiteLLM Router), establishing the entry point where the app delegates LLM calls through a unified routing layer.
• Line 3: Routes from LiteLLM Router to node C (NVIDIA NIM Pod running Llama 4 on an L4 GPU) when the estimated cost is below a defined threshold — the cost-optimized self-hosted path.
• Line 4: Routes from LiteLLM Router to node D (hosted APIs like Claude or Gemini) as a fallback when cost exceeds the threshold or when the NIM pod is unavailable, ensuring resilience.
• Line 5: Connects the NIM Pod (C) to node E, showing it runs on a dedicated GKE node pool (g2-standard-8 instance with an NVIDIA L4 GPU), detailing the infrastructure backing self-hosted inference.
• Line 6: Connects the hosted API path (D) to node F (Cloud API Endpoint), indicating this route incurs per-token billing from the external provider.
• Line 7: Shows LiteLLM Router (B) emitting metrics to node G (Prometheus), specifically tracking cost_per_token and nim_queue_depth for observability.
• Line 8: Creates a feedback loop from Prometheus (G) back to the LiteLLM Router (B), triggering an alert when cost exceeds the budget — enabling the router to dynamically adjust its routing decisions.

1from kubernetes import client
2
3def build_nim_deployment(
4    model_name: str = "meta/llama-4-maverick",
5    gpu_type: str = "nvidia-l4",
6    gpu_count: int = 1,
7    replicas: int = 2,
8    namespace: str = "inference",
9) -> client.V1Deployment:
10    """Build a NIM deployment targeting a GPU node pool."""
11    gpu_resource = f"nvidia.com/gpu"
12    container = client.V1Container(
13        name="nim-llm",
14        image="nvcr.io/nim/meta/llama-4-maverick:latest",
15        ports=[client.V1ContainerPort(container_port=8000)],
16        resources=client.V1ResourceRequirements(
17            limits={gpu_resource: str(gpu_count), "memory": "32Gi"},
18            requests={gpu_resource: str(gpu_count), "memory": "24Gi"},
19        ),
20        env=[
21            client.V1EnvVar(name="NIM_MODEL_NAME", value=model_name),
22            client.V1EnvVar(name="NIM_MAX_BATCH_SIZE", value="64"),
23            client.V1EnvVar(name="NIM_LOG_LEVEL", value="INFO"),
24            client.V1EnvVar(
25                name="NGC_API_KEY",
26                value_from=client.V1EnvVarSource(
27                    secret_key_ref=client.V1SecretKeySelector(
28                        name="nim-credentials", key="ngc-api-key"
29                    )
30                ),
31            ),
32        ],
33        volume_mounts=[
34            client.V1VolumeMount(
35                name="model-cache", mount_path="/opt/nim/.cache"
36            )
37        ],
38        readiness_probe=client.V1Probe(
39            http_get=client.V1HTTPGetAction(path="/v1/health/ready", port=8000),
40            initial_delay_seconds=120,
41            period_seconds=10,
42        ),
43    )
44    node_selector = {"cloud.google.com/gke-accelerator": gpu_type}
45    toleration = client.V1Toleration(
46        key="nvidia.com/gpu", operator="Exists", effect="NoSchedule"
47    )
48    volume = client.V1Volume(
49        name="model-cache",
50        persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
51            claim_name="nim-model-cache"
52        ),
53    )
54    template = client.V1PodTemplateSpec(
55        metadata=client.V1ObjectMeta(labels={"app": "nim-llm"}),
56        spec=client.V1PodSpec(
57            containers=[container],
58            node_selector=node_selector,
59            tolerations=[toleration],
60            volumes=[volume],
61        ),
62    )
63    return client.V1Deployment(
64        metadata=client.V1ObjectMeta(name="nim-llm", namespace=namespace),
65        spec=client.V1DeploymentSpec(
66            replicas=replicas,
67            selector=client.V1LabelSelector(match_labels={"app": "nim-llm"}),
68            template=template,
69        ),
70    )

• Lines 1-1: Imports the kubernetes client library, which provides typed Python objects for every Kubernetes resource kind.
• Lines 3-9: Defines build_nim_deployment with parameters that vary across environments—gpu_type switches between "nvidia-l4" for cost-efficient inference and "nvidia-a100-80gb" for maximum throughput.
• Lines 10-11: Sets the GPU resource key. Kubernetes uses nvidia.com/gpu as the extended resource name registered by the NVIDIA device plugin DaemonSet.
• Lines 12-18: Constructs the container spec pointing to the NIM image on NVIDIA's NGC registry (nvcr.io). The port 8000 is NIM's default HTTP serving port.
• Lines 19-22: Configures GPU resource limits and requests. Setting limits and requests to the same GPU count ensures the pod gets exactly that many GPUs—Kubernetes does not support GPU overcommit.
• Lines 23-34: Injects environment variables. NIM_MAX_BATCH_SIZE controls in-flight batching—higher values increase throughput but consume more GPU memory. The NGC_API_KEY is sourced from a Kubernetes Secret rather than hardcoded, following the principle that secrets never appear in deployment manifests.
• Lines 35-38: Mounts a persistent volume at NIM's cache directory. Model weights for Llama 4 Maverick are approximately 20 GB; caching them on a PVC means pod restarts skip the multi-minute download from NGC.
• Lines 39-43: Configures a readiness probe with a 120-second initial delay. NIM takes 1-2 minutes to load model weights into GPU memory, compile TensorRT engines on first run, and begin serving. The probe prevents Kubernetes from routing traffic to a pod that is still initializing.
• Lines 44-48: Sets node_selector to target only nodes with the specified GPU accelerator. The toleration matches the taint applied to GPU nodes, ensuring non-GPU workloads do not accidentally schedule onto expensive GPU machines.
• Lines 49-63: Assembles the pod template, pod spec, and deployment spec. The replicas parameter defaults to 2 for high availability—if one pod is evicted during a node upgrade, the other continues serving.

1import litellm
2from litellm import Router
3import httpx
4import logging
5
6logger = logging.getLogger(__name__)
7
8ROUTER_CONFIG = {
9    "model_list": [
10        {
11            "model_name": "llama4",
12            "litellm_params": {
13                "model": "openai/meta-llama-4-maverick",
14                "api_base": "http://nim-llm.inference.svc.cluster.local:8000/v1",
15                "api_key": "not-needed",
16            },
17            "model_info": {"id": "nim-local", "input_cost_per_token": 0.0},
18        },
19        {
20            "model_name": "llama4",
21            "litellm_params": {
22                "model": "anthropic/claude-sonnet-4-20250514",
23                "api_key": "os.environ/ANTHROPIC_API_KEY",
24            },
25            "model_info": {
26                "id": "claude-fallback",
27                "input_cost_per_token": 0.000003,
28            },
29        },
30    ],
31    "routing_strategy": "cost-based-routing",
32    "num_retries": 2,
33    "timeout": 30,
34}
35
36class CostThresholdRouter:
37    def __init__(self, cost_threshold_per_1k: float = 0.005):
38        self.router = Router(**ROUTER_CONFIG)
39        self.cost_threshold = cost_threshold_per_1k
40        self.nim_health_url = (
41            "http://nim-llm.inference.svc.cluster.local:8000/v1/health/ready"
42        )
43        self._http_client = httpx.AsyncClient(timeout=2.0)
44
45    async def _check_nim_health(self) -> bool:
46        try:
47            resp = await self._http_client.get(self.nim_health_url)
48            return resp.status_code == 200
49        except httpx.RequestError:
50            logger.warning("NIM health check failed, routing to fallback")
51            return False
52
53    async def route_request(self, messages: list[dict], **kwargs) -> dict:
54        nim_healthy = await self._check_nim_health()
55        if nim_healthy:
56            try:
57                response = await self.router.acompletion(
58                    model="llama4",
59                    messages=messages,
60                    specific_deployment={"id": "nim-local"},
61                    **kwargs,
62                )
63                logger.info("Request served by NIM (cost: $0.00)")
64                return response
65            except Exception as exc:
66                logger.error("NIM request failed: %s, falling back", exc)
67
68        response = await self.router.acompletion(
69            model="llama4",
70            messages=messages,
71            specific_deployment={"id": "claude-fallback"},
72            **kwargs,
73        )
74        logger.info("Request served by hosted fallback")
75        return response

• Lines 1-4: Imports litellm, its Router class, httpx for async HTTP calls, and the standard logging module. The httpx.AsyncClient is preferred over requests because the router operates in an async FastAPI context.
• Lines 8-35: Defines the router configuration dictionary. Two model deployments share the same model_name of "llama4", enabling LiteLLM to treat them as interchangeable backends. The NIM deployment at line 14 uses the Kubernetes service DNS name (nim-llm.inference.svc.cluster.local) for in-cluster routing, eliminating external network hops. The input_cost_per_token of 0.0 for NIM reflects that self-hosted inference has no per-token charge—you pay for the GPU node regardless.
• Lines 38-45: The CostThresholdRouter.__init__ method instantiates the LiteLLM Router with the config and stores the cost threshold. The httpx.AsyncClient is created once and reused across requests to leverage HTTP connection pooling, which matters when the health check runs on every request.
• Lines 47-53: The _check_nim_health coroutine sends a GET request to NIM's readiness endpoint. The 2-second timeout ensures a hung NIM pod does not block the entire request pipeline. On any httpx.RequestError (connection refused, timeout, DNS failure), the method returns False and logs a warning.
• Lines 55-64: The route_request method implements the routing decision. When NIM is healthy, it calls self.router.acompletion with specific_deployment={"id": "nim-local"} to bypass LiteLLM's built-in routing and force the request to the NIM backend. This explicit targeting is necessary because LiteLLM's cost-based routing might not account for the real-time health status.
• Lines 65-66: If the NIM call raises any exception—such as a CUDA out-of-memory error surfaced as an HTTP 500—the code catches it, logs the error, and falls through to the fallback path. This try/except block ensures no request is lost due to GPU-level failures.
• Lines 68-74: The fallback path explicitly routes to claude-fallback. In production, you would emit a metric here (e.g., nim_fallback_total.inc()) to track how often fallback occurs, which directly correlates with unexpected cost increases.