Repository: kubernetes-sigs/gateway-api-inference-extension

Issue:
InferencePool helm chart flags setting is not user friendly
the flags section of the InferencePool helm chart is currently modeled as an array instead of a map.
that means that if one wants to set a single flag, the options are either to use a values.yaml file, OR to use indexed array of values, e.g.,:
```
--set inferenceExtension.flags[0].name=<MY_FLAG> --set inferenceExtension.flags[0].value=true
```

This is an non user friendly UX and is not as expected.
the expectation is to be able to set a cmd-line flag intuitively using:
```
--set inferenceExtension.flags.<MY_FLAG>=true
```

Hints from issue comments:

Solution patch:
diff --git a/.gitignore b/.gitignore
index 82afc2e..18a475b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,7 @@
 bin/*
 Dockerfile.cross
 artifacts
+latencypredictor/__pycache__
 
 # Test binary, built with `go test -c`
 *.test
diff --git a/config/charts/inferencepool/README.md b/config/charts/inferencepool/README.md
index 9fbbfb9..da4a1a9 100644
--- a/config/charts/inferencepool/README.md
+++ b/config/charts/inferencepool/README.md
@@ -22,6 +22,28 @@ $ helm install vllm-llama3-8b-instruct \
 
 Note that the provider name is needed to deploy provider-specific resources. If no provider is specified, then only the InferencePool object and the EPP are deployed.
 
+### Install with Custom Cmd-line Flags
+
+To set cmd-line flags, you can use the `--set` option to set each flag, e.g.,:
+
+```txt
+$ helm install vllm-llama3-8b-instruct \
+  --set inferencePool.modelServers.matchLabels.app=vllm-llama3-8b-instruct \
+  --set inferenceExtension.flags.<FLAG_NAME>=<FLAG_VALUE>
+  --set provider.name=[none|gke|istio] \
+  oci://us-central1-docker.pkg.dev/k8s-staging-images/gateway-api-inference-extension/charts/inferencepool --version v0
+```
+
+Alternatively, you can define flags in the `values.yaml` file:
+
+```yaml
+inferenceExtension:
+  flags:
+    FLAG_NAME: <FLAG_VALUE>
+    v: 3 ## Log verbosity
+    ...
+```
+
 ### Install with Custom Environment Variables
 
 To set custom environment variables for the EndpointPicker deployment, you can define them as free-form YAML in the `values.yaml` file:
@@ -182,7 +204,7 @@ The following table list the configurable parameters of the chart.
 | `inferenceExtension.env`                                   | List of environment variables to set in the endpoint picker container as free-form YAML. Defaults to `[]`.                                                                                                                                         |
 | `inferenceExtension.extraContainerPorts`                   | List of additional container ports to expose. Defaults to `[]`.                                                                                                                                                                                    |
 | `inferenceExtension.extraServicePorts`                     | List of additional service ports to expose. Defaults to `[]`.                                                                                                                                                                                      |
-| `inferenceExtension.flags`                                 | List of flags which are passed through to endpoint picker. Example flags, enable-pprof, grpc-port etc. Refer [runner.go](https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/cmd/epp/runner/runner.go) for complete list. |
+| `inferenceExtension.flags`                                 | map of flags which are passed through to endpoint picker. Example flags, enable-pprof, grpc-port etc. Refer [runner.go](https://github.com/kubernetes-sigs/gateway-api-inference-extension/blob/main/cmd/epp/runner/runner.go) for complete list. |
 | `inferenceExtension.affinity`                              | Affinity for the endpoint picker. Defaults to `{}`.                                                                                                                                                                                                |
 | `inferenceExtension.tolerations`                           | Tolerations for the endpoint picker. Defaults to `[]`.                                                                                                                                                                                             |
 | `inferenceExtension.monitoring.interval`                   | Metrics scraping interval for monitoring. Defaults to `10s`.                                                                                                                                                                                       |
diff --git a/config/charts/inferencepool/templates/epp-deployment.yaml b/config/charts/inferencepool/templates/epp-deployment.yaml
index de89233..1a791c2 100644
--- a/config/charts/inferencepool/templates/epp-deployment.yaml
+++ b/config/charts/inferencepool/templates/epp-deployment.yaml
@@ -58,9 +58,9 @@ spec:
         - --ha-enable-leader-election
         {{- end }}
         # Pass additional flags via the inferenceExtension.flags field in values.yaml.
-        {{- range .Values.inferenceExtension.flags }}
-        - "--{{ .name }}"
-        - "{{ .value }}"
+        {{- range $key, $value := .Values.inferenceExtension.flags }}
+        - --{{ $key }}
+        - {{ $value }}
         {{- end }}
         {{- if .Values.inferenceExtension.tracing.enabled }}
         - --tracing=true
@@ -103,7 +103,6 @@ spec:
             service: inference-extension
           {{- end }}
           periodSeconds: 2
-
         env:
         - name: NAMESPACE
           valueFrom:
diff --git a/config/charts/inferencepool/values.yaml b/config/charts/inferencepool/values.yaml
index 8b3385a..daf2bfc 100644
--- a/config/charts/inferencepool/values.yaml
+++ b/config/charts/inferencepool/values.yaml
@@ -33,8 +33,7 @@ inferenceExtension:
 
   flags:
     # Log verbosity
-    - name: v
-      value: 1
+    v: 1
 
   affinity: {}
 
diff --git a/latencypredictor/Dockerfile-prediction b/latencypredictor/Dockerfile-prediction
new file mode 100644
index 0000000..6e3015f
--- /dev/null
+++ b/latencypredictor/Dockerfile-prediction
@@ -0,0 +1,26 @@
+# Use an official Python runtime as a parent image
+FROM python:3.11-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the requirements file and install dependencies
+# (It's good practice to manage dependencies in a requirements.txt file)
+
+
+RUN apt-get update && apt-get install -y \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+    
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application code
+COPY . .
+
+# Expose the port the app runs on
+EXPOSE 8001
+
+# Command to run the application using uvicorn
+# We use 0.0.0.0 to bind to all network interfaces inside the container
+CMD ["uvicorn", "prediction_server:app", "--host", "0.0.0.0", "--port", "8001"]
diff --git a/latencypredictor/Dockerfile-test b/latencypredictor/Dockerfile-test
new file mode 100644
index 0000000..ea2335d
--- /dev/null
+++ b/latencypredictor/Dockerfile-test
@@ -0,0 +1,38 @@
+# Dockerfile-test
+FROM python:3.9-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    curl \
+    wget \
+    jq \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements and install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install additional testing dependencies
+RUN pip install --no-cache-dir \
+    pytest \
+    pytest-asyncio \
+    requests \
+    httpx \
+    aiohttp
+
+# Copy test files
+COPY test_dual_server_client.py .
+
+
+# Create test results directory
+RUN mkdir -p /test-results
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV PYTHONUNBUFFERED=1
+
+# Default command runs the specific test
+CMD ["pytest", "-v", "-s", "test_dual_server_client.py"]
\ No newline at end of file
diff --git a/latencypredictor/Dockerfile-training b/latencypredictor/Dockerfile-training
new file mode 100644
index 0000000..4a2c2ef
--- /dev/null
+++ b/latencypredictor/Dockerfile-training
@@ -0,0 +1,27 @@
+# Use an official Python runtime as a parent image
+FROM python:3.11-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the requirements file and install dependencies
+# (It's good practice to manage dependencies in a requirements.txt file)
+
+
+RUN apt-get update && apt-get install -y \
+    libgomp1 \
+    && rm -rf /var/lib/apt/lists/*
+
+    
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application code
+COPY . .
+
+# Expose the port the app runs on
+EXPOSE 8000
+
+# Command to run the application using uvicorn
+# We use 0.0.0.0 to bind to all network interfaces inside the container
+CMD ["uvicorn", "training_server:app", "--host", "0.0.0.0", "--port", "8000"]
\ No newline at end of file
diff --git a/latencypredictor/README.md b/latencypredictor/README.md
new file mode 100644
index 0000000..3e608a3
--- /dev/null
+++ b/latencypredictor/README.md
@@ -0,0 +1,220 @@
+# Latency Predictor - Build Guide
+
+This directory contains the Latency Predictor component with dual server architecture (training and prediction servers). Use the provided `build-deploy.sh` script to build and deploy container images to Google Cloud Platform.
+
+## Prerequisites
+
+- Docker (latest version)
+- Google Cloud SDK (`gcloud`) configured and authenticated
+- Required files in directory:
+  - `training_server.py`
+  - `prediction_server.py` 
+  - `requirements.txt`
+  - `Dockerfile-training`
+  - `Dockerfile-prediction`
+  - `dual-server-deployment.yaml`
+
+**Optional (for deployment and testing):**
+- kubectl configured for GKE cluster access
+
+## Configuration
+
+Before running the script, update the configuration variables in `build-deploy.sh`:
+
+```bash
+# Edit these values in the script
+PROJECT_ID="your-gcp-project-id"
+REGION="your-gcp-region"
+REPOSITORY="your-artifact-registry-repo"
+TRAINING_IMAGE="latencypredictor-training-server"
+PREDICTION_IMAGE="latencypredictor-prediction-server"
+TAG="latest"
+```
+
+## Usage
+
+### Build Images Only
+
+```bash
+# Make script executable
+chmod +x build-deploy.sh
+
+# Build and push images to registry
+./build-deploy.sh build
+./build-deploy.sh push
+```
+
+### Complete Build and Deploy (Optional)
+
+```bash
+# Run complete process (build, push, deploy, test)
+# Note: This requires GKE cluster access
+./build-deploy.sh all
+```
+
+### Individual Commands
+
+```bash
+# Check if all required files exist
+./build-deploy.sh check
+
+# Build Docker images only
+./build-deploy.sh build
+
+# Push images to Google Artifact Registry
+./build-deploy.sh push
+
+# Optional: Deploy to GKE cluster (requires cluster access)
+./build-deploy.sh deploy
+
+# Optional: Get service information and IPs
+./build-deploy.sh info
+
+# Optional: Test the deployed services
+./build-deploy.sh test
+```
+
+## What the Script Does
+
+### Build Phase (`./build-deploy.sh build`)
+- Builds training server image from `Dockerfile-training`
+- Builds prediction server image from `Dockerfile-prediction` 
+- Tags images for Google Artifact Registry
+- Images created:
+  - `latencypredictor-training-server:latest`
+  - `latencypredictor-prediction-server:latest`
+
+### Push Phase (`./build-deploy.sh push`)
+- Configures Docker for Artifact Registry authentication
+- Pushes both images to:
+  - `us-docker.pkg.dev/PROJECT_ID/REPOSITORY/latencypredictor-training-server:latest`
+  - `us-docker.pkg.dev/PROJECT_ID/REPOSITORY/latencypredictor-prediction-server:latest`
+
+### Deploy Phase (`./build-deploy.sh deploy`) - Optional
+- Applies Kubernetes manifests from `dual-server-deployment.yaml`
+- Waits for deployments to be ready (5-minute timeout)
+- Creates services:
+  - `training-service-external` (LoadBalancer)
+  - `prediction-service` (LoadBalancer)
+
+### Test Phase (`./build-deploy.sh test`) - Optional
+- Tests health endpoint: `/healthz`
+- Tests prediction endpoint: `/predict` with sample data
+- Sample prediction request:
+  ```json
+  {
+    "kv_cache_percentage": 0.3,
+    "input_token_length": 100,
+    "num_request_waiting": 2,
+    "num_request_running": 1,
+    "num_tokens_generated": 50
+  }
+  ```
+
+## Setup Instructions
+
+1. **Configure GCP Authentication**:
+   ```bash
+   gcloud auth login
+   gcloud config set project YOUR_PROJECT_ID
+   ```
+
+2. **Configure kubectl for GKE (Optional - only needed for deployment)**:
+   ```bash
+   gcloud container clusters get-credentials CLUSTER_NAME --zone ZONE
+   ```
+
+3. **Update Script Configuration**:
+   ```bash
+   # Edit build-deploy.sh with your project details
+   nano build-deploy.sh
+   ```
+
+4. **Build Images**:
+   ```bash
+   ./build-deploy.sh build
+   ./build-deploy.sh push
+   ```
+
+5. **Optional: Deploy and Test**:
+   ```bash
+   ./build-deploy.sh deploy
+   ./build-deploy.sh test
+   # Or run everything at once
+   ./build-deploy.sh all
+   ```
+
+## Troubleshooting
+
+### Permission Issues
+```bash
+chmod +x build-deploy.sh
+```
+
+### GCP Authentication
+```bash
+gcloud auth configure-docker us-docker.pkg.dev
+```
+
+### Check Cluster Access
+```bash
+kubectl cluster-info
+kubectl get nodes
+```
+
+### View Service Status
+```bash
+./build-deploy.sh info
+kubectl get services
+kubectl get pods
+```
+
+### Check Logs
+```bash
+# Training server logs
+kubectl logs -l app=training-server
+
+# Prediction server logs  
+kubectl logs -l app=prediction-server
+```
+
+## Development Workflow
+
+1. **Make code changes** to `training_server.py` or `prediction_server.py`
+2. **Test locally** (optional):
+   ```bash
+   python training_server.py
+   python prediction_server.py
+   ```
+3. **Build and push images**:
+   ```bash
+   ./build-deploy.sh build
+   ./build-deploy.sh push
+   ```
+
+4. **Optional: Deploy and test**:
+   ```bash
+   ./build-deploy.sh deploy
+   ./build-deploy.sh test
+   ```
+
+## Service Endpoints
+
+After successful deployment:
+
+- **Training Service**: External LoadBalancer IP (check with `./build-deploy.sh info`)
+- **Prediction Service**: External LoadBalancer IP (check with `./build-deploy.sh info`)
+- **Health Check**: `http://PREDICTION_IP/healthz`
+- **Prediction API**: `http://PREDICTION_IP/predict` (POST)
+
+## Manual Build (Alternative)
+
+If you need to build manually:
+
+```bash
+# Build training server
+docker build -f Dockerfile-training -t training-server .
+
+# Build prediction server  
+docker build -f Dockerfile-prediction -t prediction-server .
+```
\ No newline at end of file
diff --git a/latencypredictor/build-deploy.sh b/latencypredictor/build-deploy.sh
new file mode 100755
index 0000000..27aaa56
--- /dev/null
+++ b/latencypredictor/build-deploy.sh
@@ -0,0 +1,238 @@
+#!/bin/bash
+
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+set -e
+
+# Configuration
+PROJECT_ID="your-gcp-project-id"
+REGION="your-gcp-region"
+REPOSITORY="your-artifact-registry-repo"
+TRAINING_IMAGE="latencypredictor-training-server"
+PREDICTION_IMAGE="latencypredictor-prediction-server"
+TAG="latest"
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No Color
+
+echo_status() {
+    echo -e "${GREEN}[INFO]${NC} $1"
+}
+
+echo_warning() {
+    echo -e "${YELLOW}[WARN]${NC} $1"
+}
+
+echo_error() {
+    echo -e "${RED}[ERROR]${NC} $1"
+}
+
+# Check if required files exist
+check_files() {
+    echo_status "Checking required files..."
+    
+    local files=("training_server.py" "prediction_server.py" "requirements.txt" "Dockerfile-training" "Dockerfile-prediction")
+    for file in "${files[@]}"; do
+        if [[ ! -f "$file" ]]; then
+            echo_error "Required file $file not found!"
+            exit 1
+        fi
+    done
+    
+    echo_status "All required files found."
+}
+
+# Build Docker images
+build_images() {
+    echo_status "Building Docker images..."
+    
+    # Build training server image
+    echo_status "Building training server image..."
+    docker build -f Dockerfile-training  -t ${TRAINING_IMAGE}:${TAG} .
+
+    # Tag for training server
+    docker tag ${TRAINING_IMAGE}:${TAG} \
+        us-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${TRAINING_IMAGE}:${TAG}
+    
+    # Build prediction server image
+    echo_status "Building prediction server image..."
+    docker build -f Dockerfile-prediction -t ${PREDICTION_IMAGE}:${TAG} .
+
+    # Tag for prediction server
+    docker tag ${PREDICTION_IMAGE}:${TAG} \
+        us-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${PREDICTION_IMAGE}:${TAG}
+    
+    echo_status "Images built successfully."
+}
+
+# Push images to Artifact Registry
+push_images() {
+    echo_status "Pushing images to Artifact Registry..."
+    
+    # Configure Docker for Artifact Registry
+    gcloud auth configure-docker us-docker.pkg.dev --quiet
+    
+    # Push training server
+    echo_status "Pushing training server image..."
+    docker push us-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${TRAINING_IMAGE}:${TAG}
+    
+    # Push prediction server
+    echo_status "Pushing prediction server image..."
+    docker push us-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${PREDICTION_IMAGE}:${TAG}
+    
+    echo_status "Images pushed successfully."
+}
+
+# Deploy to GKE
+deploy_to_gke() {
+    echo_status "Deploying to GKE..."
+    
+    # Apply the Kubernetes manifests
+    kubectl apply -f dual-server-deployment.yaml
+    
+    # Wait for deployments to be ready
+    echo_status "Waiting for training server deployment..."
+    kubectl rollout status deployment/training-server-deployment --timeout=300s
+    
+    echo_status "Waiting for prediction server deployment..."
+    kubectl rollout status deployment/prediction-server-deployment --timeout=300s
+    
+    echo_status "Deployment completed successfully."
+}
+
+# Get service information
+get_service_info() {
+    echo_status "Getting service information..."
+    
+    echo_status "Training Service:"
+    kubectl get service training-service-external -o wide
+    
+    echo_status "Prediction Service:"
+    kubectl get service prediction-service -o wide
+    
+    echo_status "Getting external IPs (may take a few minutes)..."
+    
+    # Wait for external IPs
+    echo_status "Waiting for training service external IP..."
+    kubectl get service training-service-external --watch --timeout=300s &
+    TRAINING_PID=$!
+    
+    echo_status "Waiting for prediction service external IP..."
+    kubectl get service prediction-service --watch --timeout=300s &
+    PREDICTION_PID=$!
+    
+    # Kill background processes after timeout
+    sleep 10
+    kill $TRAINING_PID $PREDICTION_PID 2>/dev/null || true
+    
+    echo_status "Current service status:"
+    kubectl get services
+}
+
+# Test the deployment
+test_deployment() {
+    echo_status "Testing deployment..."
+    
+    # Get prediction service external IP
+    PREDICTION_IP=$(kubectl get service prediction-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' 2>/dev/null || echo "")
+    
+    if [[ -n "$PREDICTION_IP" ]]; then
+        echo_status "Testing prediction endpoint at http://${PREDICTION_IP}/"
+        
+        # Test health endpoint
+        if curl -f -s "http://${PREDICTION_IP}/healthz" > /dev/null; then
+            echo_status "Health check passed!"
+        else
+            echo_warning "Health check failed or service not ready yet."
+        fi
+        
+        # Test prediction endpoint
+        echo_status "Testing prediction with sample data..."
+        curl -X POST "http://${PREDICTION_IP}/predict" \
+            -H "Content-Type: application/json" \
+            -d '{
+                "kv_cache_percentage": 0.3,
+                "input_token_length": 100,
+                "num_request_waiting": 2,
+                "num_request_running": 1,
+                "num_tokens_generated": 50
+            }' || echo_warning "Prediction test failed or service not ready yet."
+    else
+        echo_warning "External IP not assigned yet. You can test later using:"
+        echo "kubectl get services"
+    fi
+}
+
+# Cleanup function
+cleanup() {
+    echo_status "Cleaning up..."
+    docker system prune -f
+}
+
+# Main execution
+main() {
+    echo_status "Starting build and deployment process..."
+    
+    case "${1:-all}" in
+        "check")
+            check_files
+            ;;
+        "build")
+            check_files
+            build_images
+            ;;
+        "push")
+            push_images
+            ;;
+        "deploy")
+            deploy_to_gke
+            ;;
+        "info")
+            get_service_info
+            ;;
+        "test")
+            test_deployment
+            ;;
+        "all")
+            check_files
+            build_images
+            push_images
+            deploy_to_gke
+            get_service_info
+            test_deployment
+            cleanup
+            ;;
+        *)
+            echo "Usage: $0 {check|build|push|deploy|info|test|all}"
+            echo ""
+            echo "Commands:"
+            echo "  check  - Check if required files exist"
+            echo "  build  - Build Docker images"
+            echo "  push   - Push images to Artifact Registry"
+            echo "  deploy - Deploy to GKE"
+            echo "  info   - Get service information"
+            echo "  test   - Test the deployment"
+            echo "  all    - Run complete build and deployment process"
+            exit 1
+            ;;
+    esac
+    
+    echo_status "Process completed successfully!"
+}
+
+# Run main function
+main "$@"
\ No newline at end of file
diff --git a/latencypredictor/manifests/dual-server-deployment.yaml b/latencypredictor/manifests/dual-server-deployment.yaml
new file mode 100644
index 0000000..761cf79
--- /dev/null
+++ b/latencypredictor/manifests/dual-server-deployment.yaml
@@ -0,0 +1,261 @@
+# Simple deployment using HTTP for model sharing - No ReadWriteMany needed!
+
+# --- 1. ConfigMaps ---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: latency-predictor-config
+  namespace: default
+data:
+  LATENCY_RETRAINING_INTERVAL_SEC: "1"
+  LATENCY_MIN_SAMPLES_FOR_RETRAIN: "100"
+  LATENCY_TTFT_MODEL_PATH: "/models/ttft.joblib"
+  LATENCY_TPOT_MODEL_PATH: "/models/tpot.joblib"
+  LATENCY_TTFT_SCALER_PATH: "/models/ttft_scaler.joblib"
+  LATENCY_TPOT_SCALER_PATH: "/models/tpot_scaler.joblib"
+  LATENCY_MODEL_TYPE: "xgboost"
+
+---
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: prediction-server-config
+  namespace: default
+data:
+  MODEL_SYNC_INTERVAL_SEC: "10"  # Download models every 5 seconds
+  LATENCY_MODEL_TYPE: "xgboost"
+  PREDICT_HOST: "0.0.0.0"
+  PREDICT_PORT: "8001"
+  TRAINING_SERVER_URL: "http://training-service:8000"
+  LOCAL_TTFT_MODEL_PATH: "/local_models/ttft.joblib"
+  LOCAL_TPOT_MODEL_PATH: "/local_models/tpot.joblib"
+  LOCAL_TTFT_SCALER_PATH: "/local_models/ttft_scaler.joblib"
+  LOCAL_TPOT_SCALER_PATH: "/local_models/tpot_scaler.joblib"
+  HTTP_TIMEOUT: "30"
+
+---
+# --- 2. StorageClass for Hyperdisk ---
+apiVersion: storage.k8s.io/v1
+kind: StorageClass
+metadata:
+  name: hyperdisk-balanced-sc
+provisioner: pd.csi.storage.gke.io
+parameters:
+  type: hyperdisk-balanced
+reclaimPolicy: Delete
+allowVolumeExpansion: true
+volumeBindingMode: WaitForFirstConsumer
+
+---
+# --- 3. Persistent Volume Claim (PVC) ---
+# Requests persistent storage for the models using the custom StorageClass.
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: training-models-pvc
+  namespace: default
+spec:
+  storageClassName: hyperdisk-balanced-sc # Explicitly use the compatible StorageClass
+  accessModes:
+    - ReadWriteOnce # Sufficient since only the leader pod writes to the volume.
+  resources:
+    requests:
+      storage: 100Gi
+---
+# --- 3. Training Server Deployment ---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: training-server-deployment
+  namespace: default
+  labels:
+    app: training-server
+    component: training
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app: training-server
+      component: training
+  template:
+    metadata:
+      labels:
+        app: training-server
+        component: training
+    spec:
+      nodeSelector:
+        cloud.google.com/gke-nodepool: "pool-2"
+      containers:
+      - name: training-server
+        image: placeholder-for-training-server-image:latest # TODO Replace with actual image
+
+        imagePullPolicy: Always
+        ports:
+        - containerPort: 8000
+          name: training-port
+        livenessProbe:
+          httpGet:
+            path: /healthz
+            port: 8000
+          initialDelaySeconds: 30
+          periodSeconds: 20
+        readinessProbe:
+          httpGet:
+            path: /readyz
+            port: 8000
+          initialDelaySeconds: 45
+          periodSeconds: 10
+        resources:
+          # Increased CPU & memory
+          requests:
+            cpu: "1000m"            # was 500m
+            memory: "2Gi"           # was 512Mi
+            #ephemeral-storage: "50Gi" # new: reserve 5Gi of scratch space
+          limits:
+            cpu: "2000m"            # was 1000m
+            memory: "4Gi"           # was 1Gi
+            #ephemeral-storage: "100Gi" # new: cap at 10Gi of scratch space
+
+        envFrom:
+        - configMapRef:
+            name: latency-predictor-config
+        env:
+        - name: POD_NAME
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.name
+        - name: SERVER_TYPE
+          value: "training"
+        volumeMounts:
+        - name: model-storage
+          mountPath: /models
+      volumes:
+      - name: model-storage
+        persistentVolumeClaim:
+          claimName: training-models-pvc
+
+---
+# --- 4. Prediction Server Deployment ---
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: prediction-server-deployment
+  namespace: default
+  labels:
+    app: prediction-server
+    component: prediction
+spec:
+  replicas: 10
+  selector:
+    matchLabels:
+      app: prediction-server
+      component: prediction
+  template:
+    metadata:
+      labels:
+        app: prediction-server
+        component: prediction
+    spec:
+      nodeSelector:
+        cloud.google.com/gke-nodepool: "pool-2"
+      containers:
+      - name: prediction-server
+        image: placeholder-for-prediction-server-image:latest # TODO Replace with actual image
+        imagePullPolicy: Always
+        ports:
+        - containerPort: 8001
+          name: predict-port
+        livenessProbe:
+          httpGet:
+            path: /healthz
+            port: 8001
+          initialDelaySeconds: 15
+          periodSeconds: 15
+        readinessProbe:
+          httpGet:
+            path: /readyz
+            port: 8001
+          initialDelaySeconds: 10
+          periodSeconds: 5
+          failureThreshold: 10  # Allow more failures while downloading models
+        resources:
+          requests:
+            cpu: "250m"
+            memory: "512Mi"
+          limits:
+            cpu: "1000m"
+            memory: "2Gi"
+        envFrom:
+        - configMapRef:
+            name: prediction-server-config
+        env:
+        - name: POD_NAME
+          valueFrom:
+            fieldRef:
+              fieldPath: metadata.name
+        - name: SERVER_TYPE
+          value: "prediction"
+        volumeMounts:
+        # Only local storage needed - no shared volumes!
+        - name: local-model-storage
+          mountPath: /local_models
+      volumes:
+      - name: local-model-storage
+        emptyDir: {}  # Each pod gets its own local storage
+
+---
+# --- 5. Services ---
+apiVersion: v1
+kind: Service
+metadata:
+  name: training-service
+  namespace: default
+  labels:
+    component: training
+spec:
+  type: ClusterIP
+  selector:
+    app: training-server
+    component: training
+  ports:
+  - protocol: TCP
+    port: 8000
+    targetPort: 8000
+    name: training
+
+---
+apiVersion: v1
+kind: Service
+metadata:
+  name: prediction-service
+  namespace: default
+  labels:
+    component: prediction
+spec:
+  type: LoadBalancer
+  selector:
+    app: prediction-server
+    component: prediction
+  ports:
+  - protocol: TCP
+    port: 80
+    targetPort: 8001
+    name: prediction
+
+---
+# --- 6. Optional: External Training Service ---
+apiVersion: v1
+kind: Service
+metadata:
+  name: training-service-external
+  namespace: default
+spec:
+  type: LoadBalancer
+  selector:
+    app: training-server
+    component: training
+  ports:
+  - protocol: TCP
+    port: 8080
+    targetPort: 8000
+
diff --git a/latencypredictor/manifests/test-dual-server-deployment.yaml b/latencypredictor/manifests/test-dual-server-deployment.yaml
new file mode 100644
index 0000000..59a76aa
--- /dev/null
+++ b/latencypredictor/manifests/test-dual-server-deployment.yaml
@@ -0,0 +1,56 @@
+# --- Test Job ---
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: latency-predictor-test
+  namespace: default
+  labels:
+    app: latency-predictor-test
+    component: test
+spec:
+  template:
+    metadata:
+      labels:
+        app: latency-predictor-test
+        component: test
+    spec:
+      # Use the same node pool as your services
+      nodeSelector:
+        cloud.google.com/gke-nodepool: "pool-2"
+      restartPolicy: Never
+      containers:
+      - name: test-runner
+        # Use your test image here
+        image: placeholder-for-test-image:latest # TODO Replace with actual test image
+        imagePullPolicy: Always
+        resources:
+          requests:
+            cpu: "500m"
+            memory: "50Gi"
+          limits:
+            cpu: "1000m"
+            memory: "100Gi"
+        env:
+        # Point to your internal services
+        - name: TRAINING_SERVER_URL
+          value: "http://training-service:8000"
+        - name: PREDICTION_SERVER_URL
+          value: "http://prediction-service:80"
+        - name: TEST_TIMEOUT
+          value: "300"  # 5 minutes
+        - name: TARGET_QPS
+          value: "1000"  # Match the target QPS in your test script
+        envFrom:
+        - configMapRef:
+            name: prediction-server-config  # Reuse existing config if needed
+        # Override the default command to run specific pytest
+        command: ["pytest"]
+        args: ["-v", "-s", "test_dual_server_client.py"]
+        # If your tests need to store results or logs
+        volumeMounts:
+        - name: test-results
+          mountPath: /test-results
+      volumes:
+      - name: test-results
+        emptyDir: {}
+  backoffLimit: 3  # Retry up to 3 times if test fails
\ No newline at end of file
diff --git a/latencypredictor/prediction_server.py b/latencypredictor/prediction_server.py
new file mode 100644
index 0000000..581f834
--- /dev/null
+++ b/latencypredictor/prediction_server.py
@@ -0,0 +1,776 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+import time
+import logging
+import threading
+import requests
+from datetime import datetime, timezone
+from typing import Tuple, Optional, List
+from enum import Enum
+
+import joblib
+import uvicorn
+import numpy as np
+import pandas as pd
+from fastapi import FastAPI, HTTPException, status
+from pydantic import BaseModel, Field
+
+# Try to import XGBoost; fall back if unavailable
+try:
+    import xgboost as xgb
+    XGBOOST_AVAILABLE = True
+except ImportError:
+    XGBOOST_AVAILABLE = False
+    logging.warning("XGBoost not available. Install with: pip install xgboost")
+
+try:
+    import lightgbm as lgb
+    LIGHTGBM_AVAILABLE = True
+except ImportError:
+    LIGHTGBM_AVAILABLE = False
+    logging.warning("LightGBM not available. Install with: pip install lightgbm")
+    
+class ModelType(str, Enum):
+    BAYESIAN_RIDGE = "bayesian_ridge"
+    XGBOOST = "xgboost"
+    LIGHTGBM = "lightgbm"
+
+
+class PredictSettings:
+    """Configuration for the prediction server."""
+
+    # Training server URL
+    TRAINING_SERVER_URL: str = os.getenv("TRAINING_SERVER_URL", "http://training-service:8000")
+
+    # Local model paths
+    LOCAL_TTFT_MODEL_PATH: str = os.getenv("LOCAL_TTFT_MODEL_PATH", "/local_models/ttft.joblib")
+    LOCAL_TPOT_MODEL_PATH: str = os.getenv("LOCAL_TPOT_MODEL_PATH", "/local_models/tpot.joblib")
+    LOCAL_TTFT_SCALER_PATH: str = os.getenv("LOCAL_TTFT_SCALER_PATH", "/local_models/ttft_scaler.joblib")
+    LOCAL_TPOT_SCALER_PATH: str = os.getenv("LOCAL_TPOT_SCALER_PATH", "/local_models/tpot_scaler.joblib")
+
+    # Sync interval and model type
+    MODEL_SYNC_INTERVAL_SEC: int = int(os.getenv("MODEL_SYNC_INTERVAL_SEC", "10"))
+    MODEL_TYPE: ModelType = ModelType(os.getenv("LATENCY_MODEL_TYPE", "xgboost"))
+    
+    # Quantile configuration (should match training server)
+    QUANTILE_ALPHA: float = float(os.getenv("LATENCY_QUANTILE_ALPHA", "0.9"))  # p90 quantile
+
+    # Server host/port
+    HOST: str = os.getenv("PREDICT_HOST", "0.0.0.0")
+    PORT: int = int(os.getenv("PREDICT_PORT", "8001"))
+
+    # HTTP timeout
+    HTTP_TIMEOUT: int = int(os.getenv("HTTP_TIMEOUT", "30"))
+
+
+settings = PredictSettings()
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+class ModelSyncer:
+    """Downloads models from a training server via HTTP."""
+
+    def __init__(self):
+        self._shutdown_event = threading.Event()
+        self._sync_thread: Optional[threading.Thread] = None
+        self._sync_lock = threading.Lock()
+
+        # Ensure local directories
+        for path in [
+            settings.LOCAL_TTFT_MODEL_PATH,
+            settings.LOCAL_TPOT_MODEL_PATH,
+            settings.LOCAL_TTFT_SCALER_PATH,
+            settings.LOCAL_TPOT_SCALER_PATH,
+        ]:
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    def _download_model_if_newer(self, name: str, dest: str) -> bool:
+        try:
+            info_url = f"{settings.TRAINING_SERVER_URL}/model/{name}/info"
+            r = requests.get(info_url, timeout=settings.HTTP_TIMEOUT)
+            if r.status_code != 200:
+                return False
+            info = r.json()
+            mtime = info.get("last_modified")
+            if not mtime:
+                return False
+            server_time = datetime.fromisoformat(mtime.replace('Z', '+00:00'))
+
+            if os.path.exists(dest):
+                local_time = datetime.fromtimestamp(os.path.getmtime(dest), tz=timezone.utc)
+                if local_time >= server_time:
+                    logging.info(f"Model {name} is up-to-date: {dest}")
+                    return False
+
+            dl_url = f"{settings.TRAINING_SERVER_URL}/model/{name}/download"
+            dl = requests.get(dl_url, timeout=settings.HTTP_TIMEOUT, stream=True)
+            if dl.status_code != 200:
+                logging.error(f"Failed download {name}: {dl.status_code}")
+                return False
+
+            tmp = dest + ".tmp"
+            with open(tmp, 'wb') as f:
+                for chunk in dl.iter_content(8192):
+                    if chunk:
+                        f.write(chunk)
+            if os.path.getsize(tmp) == 0:
+                os.remove(tmp)
+                return False
+
+            # Atomic replace
+            os.replace(tmp, dest)
+            logging.info(f"Downloaded {name} -> {dest}")
+            return True
+
+        except requests.RequestException as e:
+            logging.error(f"Network error for {name}: {e}")
+            return False
+        except OSError as e:
+            logging.error(f"Filesystem error for {name}: {e}")
+            return False
+
+    def sync_models(self) -> bool:
+        """Sync all relevant models; returns True if any updated."""
+        with self._sync_lock:
+            updated = False
+            to_sync = [
+                ("ttft", settings.LOCAL_TTFT_MODEL_PATH),
+                ("tpot", settings.LOCAL_TPOT_MODEL_PATH),
+            ]
+            if settings.MODEL_TYPE == ModelType.BAYESIAN_RIDGE:
+                to_sync += [
+                    ("ttft_scaler", settings.LOCAL_TTFT_SCALER_PATH),
+                    ("tpot_scaler", settings.LOCAL_TPOT_SCALER_PATH),
+                ]
+            for name, path in to_sync:
+                if self._download_model_if_newer(name, path):
+                    updated = True
+            return updated
+
+    def _sync_loop(self):
+        while not self._shutdown_event.is_set():
+            try:
+                if self.sync_models():
+                    predictor.load_models()
+            except Exception as e:
+                logging.error(f"Error in sync loop: {e}")
+            self._shutdown_event.wait(timeout=settings.MODEL_SYNC_INTERVAL_SEC)
+        logging.info("Model sync loop exited")
+
+    def start(self):
+        if self._sync_thread:
+            return
+        self._sync_thread = threading.Thread(target=self._sync_loop, daemon=True)
+        self._sync_thread.start()
+        logging.info(f"Sync thread started (interval {settings.MODEL_SYNC_INTERVAL_SEC}s)")
+
+    def shutdown(self):
+        self._shutdown_event.set()
+        if self._sync_thread:
+            self._sync_thread.join()
+
+
+class LightweightPredictor:
+    """Handles inference using loaded quantile regression models."""
+
+    def __init__(self):
+        mt = settings.MODEL_TYPE
+        self.prefix_buckets = 4
+        
+        # Add LightGBM fallback logic
+        if mt == ModelType.XGBOOST and not XGBOOST_AVAILABLE:
+            logging.warning("XGBoost not available. Falling back to Bayesian Ridge")
+            mt = ModelType.BAYESIAN_RIDGE
+        elif mt == ModelType.LIGHTGBM and not LIGHTGBM_AVAILABLE:
+            logging.warning("LightGBM not available. Falling back to Bayesian Ridge")
+            mt = ModelType.BAYESIAN_RIDGE
+            
+        self.model_type = mt
+        self.quantile = settings.QUANTILE_ALPHA
+        self.ttft_model = None
+        self.tpot_model = None
+        self.ttft_scaler = None
+        self.tpot_scaler = None
+        self.lock = threading.RLock()
+        self.last_load: Optional[datetime] = None
+        logging.info(f"Predictor type: {self.model_type}, quantile: {self.quantile}")
+
+    @property
+    def is_ready(self) -> bool:
+        with self.lock:
+            if self.model_type == ModelType.BAYESIAN_RIDGE:
+                return all([self.ttft_model, self.tpot_model, self.ttft_scaler, self.tpot_scaler])
+            else:  # XGBoost or LightGBM
+                return all([self.ttft_model, self.tpot_model])
+
+    def _prepare_features_with_interaction(self, df: pd.DataFrame, model_type: str) -> pd.DataFrame:
+        """
+        Prepare features with interaction terms to match training server.
+        
+        Args:
+            df: DataFrame with raw features
+            model_type: 'ttft' or 'tpot'
+        
+        Returns:
+            DataFrame with engineered features including interactions
+        """
+        if model_type == "ttft":
+            # Create interaction: prefix score * input length
+            df['effective_input_tokens'] = (1-df['prefix_cache_score']) * df['input_token_length']
+            df['prefill_score_bucket'] = (
+            (df['prefix_cache_score'].clip(0, 1) * self.prefix_buckets)
+            .astype(int)
+            .clip(upper=self.prefix_buckets - 1)
+        )
+
+            # make it categorical for tree models (safe for LGB, XGB with enable_categorical)
+            df['prefill_score_bucket'] = pd.Categorical(df['prefill_score_bucket'], categories=[0,1,2,3], ordered=True)
+ 
+            
+            # Return TTFT features with interaction
+            feature_cols = [
+                'kv_cache_percentage',
+                'input_token_length',
+                'num_request_waiting',
+                'num_request_running',
+                'prefix_cache_score',
+                'effective_input_tokens',
+                'prefill_score_bucket'
+            ]
+            
+            return df[feature_cols]
+            
+        else:  # tpot
+            # TPOT doesn't use prefix_cache_score, so no interaction needed
+            feature_cols = [
+                'kv_cache_percentage',
+                'input_token_length',
+                'num_request_waiting',
+                'num_request_running',
+                'num_tokens_generated'
+            ]
+            
+            return df[feature_cols]
+
+    def load_models(self) -> bool:
+        try:
+            with self.lock:
+                new_ttft = joblib.load(settings.LOCAL_TTFT_MODEL_PATH) if os.path.exists(settings.LOCAL_TTFT_MODEL_PATH) else None
+                new_tpot = joblib.load(settings.LOCAL_TPOT_MODEL_PATH) if os.path.exists(settings.LOCAL_TPOT_MODEL_PATH) else None
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    new_ttft_scaler = joblib.load(settings.LOCAL_TTFT_SCALER_PATH) if os.path.exists(settings.LOCAL_TTFT_SCALER_PATH) else None
+                    new_tpot_scaler = joblib.load(settings.LOCAL_TPOT_SCALER_PATH) if os.path.exists(settings.LOCAL_TPOT_SCALER_PATH) else None
+                else:
+                    new_ttft_scaler = new_tpot_scaler = None
+
+                if new_ttft: self.ttft_model = new_ttft
+                if new_tpot: self.tpot_model = new_tpot
+                if new_ttft_scaler: self.ttft_scaler = new_ttft_scaler
+                if new_tpot_scaler: self.tpot_scaler = new_tpot_scaler
+                self.last_load = datetime.now(timezone.utc)
+                if self.is_ready:
+                    logging.info("Models loaded")
+                    return True
+                logging.warning("Models missing after load")
+                return False
+        except Exception as e:
+            logging.error(f"Load error: {e}")
+            return False
+
+    def predict(self, features: dict) -> Tuple[float, float]:
+        """Make quantile predictions using the loaded models."""
+        try:
+            with self.lock:
+                if not self.is_ready:
+                    raise HTTPException(status_code=503, detail="Models not ready")
+                
+                # Validation
+                required = ['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 
+                           'num_request_running', 'num_tokens_generated', 'prefix_cache_score']
+                for f in required:
+                    if f not in features:
+                        raise ValueError(f"Missing required feature: {f}")
+                    if not isinstance(features[f], (int, float)):
+                        raise ValueError(f"Invalid type for feature {f}: expected number")
+
+                # Create raw DataFrames (without interaction)
+                ttft_raw_data = {
+                    'kv_cache_percentage': features['kv_cache_percentage'],
+                    'input_token_length': features['input_token_length'],
+                    'num_request_waiting': features['num_request_waiting'],
+                    'num_request_running': features['num_request_running'],
+                    'prefix_cache_score': features['prefix_cache_score']
+                }
+                
+                tpot_raw_data = {
+                    'kv_cache_percentage': features['kv_cache_percentage'],
+                    'input_token_length': features['input_token_length'],
+                    'num_request_waiting': features['num_request_waiting'],
+                    'num_request_running': features['num_request_running'],
+                    'num_tokens_generated': features['num_tokens_generated']
+                }
+                
+                # Prepare features with interactions
+                df_ttft_raw = pd.DataFrame([ttft_raw_data])
+                df_ttft = self._prepare_features_with_interaction(df_ttft_raw, "ttft")
+               
+                
+                df_tpot_raw = pd.DataFrame([tpot_raw_data])
+                df_tpot = self._prepare_features_with_interaction(df_tpot_raw, "tpot")
+                #df_tpot = pd.DataFrame([tpot_raw_data])
+
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    
+                    ttft_for_scale = df_ttft.drop(columns=['prefill_score_bucket'], errors='ignore')
+                    ttft_scaled = self.ttft_scaler.transform(ttft_for_scale)
+                    tpot_scaled = self.tpot_scaler.transform(df_tpot)
+
+                    ttft_pred_mean, ttft_std = self.ttft_model.predict(ttft_scaled, return_std=True)
+                    tpot_pred_mean, tpot_std = self.tpot_model.predict(tpot_scaled, return_std=True)
+
+                    std_factor = 1.28 if self.quantile == 0.9 else (2.0 if self.quantile == 0.95 else 0.674)
+                    ttft_pred = ttft_pred_mean[0] + std_factor * ttft_std[0]
+                    tpot_pred = tpot_pred_mean[0] + std_factor * tpot_std[0]
+                    
+                    return ttft_pred, tpot_pred
+                
+                elif self.model_type == ModelType.XGBOOST:
+                    ttft_pred = self.ttft_model.predict(df_ttft)
+                    tpot_pred = self.tpot_model.predict(df_tpot)
+                    
+                    return ttft_pred[0], tpot_pred[0]
+                
+                else:  # LightGBM
+                    ttft_pred = self.ttft_model.predict(df_ttft)
+                    tpot_pred = self.tpot_model.predict(df_tpot)
+                    
+                    return ttft_pred[0], tpot_pred[0]
+                    
+        except ValueError as ve:
+            logging.warning(f"Client error in predict(): {ve}")
+            raise HTTPException(status_code=400, detail=str(ve))
+        except HTTPException:
+            raise
+        except Exception as e:
+            logging.error("Error in predict():", exc_info=True)
+            raise HTTPException(status_code=500, detail="Internal error during prediction")
+
+    def predict_batch(self, features_list: List[dict]) -> Tuple[np.ndarray, np.ndarray]:
+        """Make batch quantile predictions using the loaded models."""
+        try:
+            with self.lock:
+                if not self.is_ready:
+                    raise HTTPException(status_code=503, detail="Models not ready")
+                
+                # Validation
+                required = ['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 
+                           'num_request_running', 'num_tokens_generated', 'prefix_cache_score']
+                for i, features in enumerate(features_list):
+                    for f in required:
+                        if f not in features:
+                            raise ValueError(f"Missing required feature '{f}' in request {i}")
+                        if not isinstance(features[f], (int, float)):
+                            raise ValueError(f"Invalid type for feature '{f}' in request {i}: expected number")
+
+                # Create raw feature data (without interaction)
+                ttft_raw_data = []
+                tpot_raw_data = []
+                
+                for features in features_list:
+                    ttft_raw_data.append({
+                        'kv_cache_percentage': features['kv_cache_percentage'],
+                        'input_token_length': features['input_token_length'],
+                        'num_request_waiting': features['num_request_waiting'],
+                        'num_request_running': features['num_request_running'],
+                        'prefix_cache_score': features['prefix_cache_score']
+                    })
+                    
+                    tpot_raw_data.append({
+                        'kv_cache_percentage': features['kv_cache_percentage'],
+                        'input_token_length': features['input_token_length'],
+                        'num_request_waiting': features['num_request_waiting'],
+                        'num_request_running': features['num_request_running'],
+                        'num_tokens_generated': features['num_tokens_generated']
+                    })
+                
+                # Prepare features with interactions
+                df_ttft_raw = pd.DataFrame(ttft_raw_data)
+                df_ttft_batch = self._prepare_features_with_interaction(df_ttft_raw, "ttft")
+                #df_ttft_batch = pd.DataFrame(ttft_raw_data)
+                
+                df_tpot_raw = pd.DataFrame(tpot_raw_data)
+                df_tpot_batch = self._prepare_features_with_interaction(df_tpot_raw, "tpot")
+                #df_tpot_batch = pd.DataFrame(tpot_raw_data)
+
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    ttft_for_scale = df_ttft_batch.drop(columns=['prefill_score_bucket'], errors='ignore')
+                    ttft_scaled = self.ttft_scaler.transform(ttft_for_scale)
+                    tpot_scaled = self.tpot_scaler.transform(df_tpot_batch)
+
+                    ttft_pred_mean, ttft_std = self.ttft_model.predict(ttft_scaled, return_std=True)
+                    tpot_pred_mean, tpot_std = self.tpot_model.predict(tpot_scaled, return_std=True)
+                    
+                    std_factor = 1.28 if self.quantile == 0.9 else (2.0 if self.quantile == 0.95 else 0.674)
+                    ttft_pred = ttft_pred_mean + std_factor * ttft_std
+                    tpot_pred = tpot_pred_mean + std_factor * tpot_std
+                    
+                    return ttft_pred, tpot_pred
+                
+                elif self.model_type == ModelType.XGBOOST:
+                    ttft_pred = self.ttft_model.predict(df_ttft_batch)
+                    tpot_pred = self.tpot_model.predict(df_tpot_batch)
+                    
+                    return ttft_pred, tpot_pred
+                
+                else:  # LightGBM
+                    ttft_pred = self.ttft_model.predict(df_ttft_batch)
+                    tpot_pred = self.tpot_model.predict(df_tpot_batch)
+                    
+                    return ttft_pred, tpot_pred
+                    
+        except ValueError as ve:
+            logging.warning(f"Client error in predict_batch(): {ve}")
+            raise HTTPException(status_code=400, detail=str(ve))
+        except HTTPException:
+            raise
+        except Exception as e:
+            logging.error("Error in predict_batch():", exc_info=True)
+            raise HTTPException(status_code=500, detail="Internal error during batch prediction")
+
+
+# Instantiate
+model_syncer = ModelSyncer()
+predictor = LightweightPredictor()
+
+# FastAPI app
+app = FastAPI(
+    title="HTTP-based Quantile Latency Predictor",
+    description="A prediction service that downloads quantile regression models from training server via HTTP.",
+    version="1.0.0"
+)
+
+
+# Pydantic models
+class PredictionRequest(BaseModel):
+    kv_cache_percentage: float = Field(..., ge=0.0, le=1.0)
+    input_token_length: int = Field(..., ge=0)
+    num_request_waiting: int = Field(..., ge=0)
+    num_request_running: int = Field(..., ge=0)
+    num_tokens_generated: int = Field(..., ge=0)
+    prefix_cache_score: float = Field(..., ge=0.0, le=1.0, description="Prefix cache hit ratio score (0.0 to 1.0)")
+
+
+class PredictionResponse(BaseModel):
+    ttft_ms: float = Field(..., description=f"Predicted {settings.QUANTILE_ALPHA:.0%} quantile TTFT in milliseconds")
+    tpot_ms: float = Field(..., description=f"Predicted {settings.QUANTILE_ALPHA:.0%} quantile TPOT in milliseconds")
+    predicted_at: datetime
+    model_type: str = Field(..., description="Type of model used for prediction")
+    quantile: float = Field(..., description="Quantile being predicted")
+    last_model_load: Optional[datetime]
+
+
+class StatusResponse(BaseModel):
+    is_ready: bool
+    model_type: str
+    quantile: float = Field(..., description="Quantile being predicted")
+    last_model_load: Optional[datetime]
+    training_server_url: str
+    models_exist: dict
+
+
+class BulkPredictionRequest(BaseModel):
+    requests: List[PredictionRequest] = Field(..., min_items=1, max_items=10000, description="List of prediction requests (max 10000)")
+
+class BulkPredictionResponse(BaseModel):
+    predictions: List[PredictionResponse] = Field(..., description="List of prediction responses")
+    total_requests: int = Field(..., description="Total number of requests processed")
+    successful_predictions: int = Field(..., description="Number of successful predictions")
+    failed_predictions: int = Field(..., description="Number of failed predictions")
+    processing_time_ms: float = Field(..., description="Total processing time in milliseconds")
+
+class BulkPredictionError(BaseModel):
+    index: int = Field(..., description="Index of the failed request in the original batch")
+    error: str = Field(..., description="Error message")
+    request: PredictionRequest = Field(..., description="The original request that failed")
+
+class BulkPredictionResponseWithErrors(BaseModel):
+    predictions: List[Optional[PredictionResponse]] = Field(..., description="List of prediction responses (None for failed predictions)")
+    errors: List[BulkPredictionError] = Field(..., description="List of errors for failed predictions")
+    total_requests: int = Field(..., description="Total number of requests processed")
+    successful_predictions: int = Field(..., description="Number of successful predictions")
+    failed_predictions: int = Field(..., description="Number of failed predictions")
+    processing_time_ms: float = Field(..., description="Total processing time in milliseconds")
+
+
+# API endpoints
+
+@app.get("/status", response_model=StatusResponse)
+async def status_endpoint():
+    """Get server status and model information."""
+    models_exist = {
+        "ttft_model": os.path.exists(settings.LOCAL_TTFT_MODEL_PATH),
+        "tpot_model": os.path.exists(settings.LOCAL_TPOT_MODEL_PATH),
+    }
+    
+    if predictor.model_type == ModelType.BAYESIAN_RIDGE:
+        models_exist.update({
+            "ttft_scaler": os.path.exists(settings.LOCAL_TTFT_SCALER_PATH),
+            "tpot_scaler": os.path.exists(settings.LOCAL_TPOT_SCALER_PATH),
+        })
+    
+    return StatusResponse(
+        is_ready=predictor.is_ready,
+        model_type=predictor.model_type.value,
+        quantile=predictor.quantile,
+        last_model_load=predictor.last_load,
+        training_server_url=settings.TRAINING_SERVER_URL,
+        models_exist=models_exist
+    )
+
+@app.post("/predict", response_model=PredictionResponse)
+async def predict_endpoint(request: PredictionRequest):
+    """Make quantile latency predictions."""
+    try:
+        ttft_pred, tpot_pred = predictor.predict(request.dict())
+        
+        # Ensure non-negative predictions
+        ttft_pred = max(0, ttft_pred)
+        tpot_pred = max(0, tpot_pred)
+        
+        return PredictionResponse(
+            ttft_ms=ttft_pred,
+            tpot_ms=tpot_pred,
+            predicted_at=datetime.now(timezone.utc),
+            model_type=predictor.model_type.value,
+            quantile=predictor.quantile,
+            last_model_load=predictor.last_load 
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logging.error(f"Prediction failed: {e}")
+        raise HTTPException(status_code=500, detail="An internal error occurred during prediction")
+
+
+@app.post("/predict/bulk/strict", response_model=BulkPredictionResponse)
+async def predict_bulk_strict_endpoint(request: BulkPredictionRequest):
+    """Make bulk quantile latency predictions using batch processing (fails on any single error)."""
+    start_time = time.time()
+    
+    try:
+        # Convert all requests to dict format
+        features_list = [pred_request.dict() for pred_request in request.requests]
+        
+        # Make batch prediction
+        ttft_preds, tpot_preds = predictor.predict_batch(features_list)
+        
+        # Build response list
+        predictions = []
+        current_time = datetime.now(timezone.utc)
+        
+        for i in range(len(request.requests)):
+            # Ensure non-negative predictions
+            ttft_pred = max(0, ttft_preds[i])
+            tpot_pred = max(0, tpot_preds[i])
+            
+            prediction_response = PredictionResponse(
+                ttft_ms=ttft_pred,
+                tpot_ms=tpot_pred,
+                predicted_at=current_time,
+                model_type=predictor.model_type.value,
+                quantile=predictor.quantile,
+                last_model_load=predictor.last_load 
+            )
+            
+            predictions.append(prediction_response)
+        
+        processing_time_ms = (time.time() - start_time) * 1000
+        
+        return BulkPredictionResponse(
+            predictions=predictions,
+            total_requests=len(request.requests),
+            successful_predictions=len(predictions),
+            failed_predictions=0,
+            processing_time_ms=processing_time_ms
+        )
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logging.error(f"Bulk prediction failed: {e}")
+        raise HTTPException(status_code=500, detail="Bulk prediction failed")
+
+
+@app.post("/predict/bulk", response_model=BulkPredictionResponseWithErrors)
+async def predict_bulk_endpoint(request: BulkPredictionRequest):
+    """Make bulk quantile latency predictions using batch processing with error handling."""
+    start_time = time.time()
+    
+    # Separate valid and invalid requests
+    valid_requests = []
+    valid_indices = []
+    errors = []
+    
+    # Pre-validate all requests
+    required = ['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 'num_request_running', 'num_tokens_generated', 'prefix_cache_score']
+    
+    for i, pred_request in enumerate(request.requests):
+        try:
+            features = pred_request.dict()
+            # Validate features
+            for f in required:
+                if f not in features:
+                    raise ValueError(f"Missing required feature: {f}")
+                if not isinstance(features[f], (int, float)):
+                    raise ValueError(f"Invalid type for feature {f}: expected number")
+            
+            valid_requests.append(features)
+            valid_indices.append(i)
+            
+        except Exception as e:
+            errors.append(BulkPredictionError(
+                index=i,
+                error=str(e),
+                request=pred_request
+            ))
+    
+    # Initialize predictions list with None values
+    predictions = [None] * len(request.requests)
+    successful_count = len(valid_requests)
+    failed_count = len(errors)
+    
+    # Process valid requests in batch if any exist
+    if valid_requests:
+        try:
+            # Make batch prediction for all valid requests
+            ttft_preds, tpot_preds = predictor.predict_batch(valid_requests)
+            
+            current_time = datetime.now(timezone.utc)
+            
+            # Fill in predictions for valid requests
+            for batch_idx, original_idx in enumerate(valid_indices):
+                # Ensure non-negative predictions
+                ttft_pred = max(0, ttft_preds[batch_idx])
+                tpot_pred = max(0, tpot_preds[batch_idx])
+                
+                prediction_response = PredictionResponse(
+                    ttft_ms=ttft_pred,
+                    tpot_ms=tpot_pred,
+                    predicted_at=current_time,
+                    model_type=predictor.model_type.value,
+                    quantile=predictor.quantile,
+                    last_model_load=predictor.last_load 
+                )
+                
+                predictions[original_idx] = prediction_response
+                
+        except Exception as e:
+            # If batch prediction fails, mark all valid requests as failed
+            for original_idx in valid_indices:
+                errors.append(BulkPredictionError(
+                    index=original_idx,
+                    error=f"Batch prediction error: {str(e)}",
+                    request=request.requests[original_idx]
+                ))
+                predictions[original_idx] = None
+            
+            successful_count = 0
+            failed_count = len(request.requests)
+    
+    processing_time_ms = (time.time() - start_time) * 1000
+    
+    return BulkPredictionResponseWithErrors(
+        predictions=predictions,
+        errors=errors,
+        total_requests=len(request.requests),
+        successful_predictions=successful_count,
+        failed_predictions=failed_count,
+        processing_time_ms=processing_time_ms
+    )
+
+@app.post("/reload")
+async def reload_models():
+    """Manually trigger model reload."""
+    try:
+        # First sync from training server
+        synced = model_syncer.sync_models()
+        
+        # Then load models
+        loaded = predictor.load_models()
+        
+        return {
+            "synced": synced,
+            "loaded": loaded,
+            "is_ready": predictor.is_ready,
+            "model_type": predictor.model_type.value,
+            "quantile": predictor.quantile,
+            "last_load_time": predictor.last_load 
+        }
+    except Exception as e:
+        logging.error(f"Error reloading models: {e}")
+        raise HTTPException(status_code=500, detail=f"Error reloading models: {str(e)}")
+
+@app.get("/healthz", status_code=status.HTTP_200_OK)
+async def health_check():
+    """Health check endpoint."""
+    return {"status": "ok", "service": "http-based-quantile-latency-predictor"}
+
+
+@app.get("/readyz", status_code=status.HTTP_200_OK)
+async def readiness_check():
+    """Readiness check endpoint."""
+    if not predictor.is_ready:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE, 
+            detail="Models are not ready"
+        )
+    return {
+        "status": "ready", 
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile
+    }
+
+
+@app.get("/", include_in_schema=False)
+async def root():
+    """Root endpoint."""
+    return {
+        "message": "HTTP-based Quantile Latency Predictor is running",
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile,
+        "description": f"Predicting {predictor.quantile:.0%} quantile for TTFT and TPOT latencies",
+        "is_ready": predictor.is_ready,
+        "sync_interval": settings.MODEL_SYNC_INTERVAL_SEC,
+        "training_server": settings.TRAINING_SERVER_URL
+    }
+
+
+@app.on_event("startup")
+async def startup():
+    logging.info("Starting up...")
+    # initial sync & load
+    model_syncer.sync_models()
+    predictor.load_models()
+    model_syncer.start()
+
+@app.on_event("shutdown")
+async def shutdown():
+    logging.info("Shutting down...")
+    model_syncer.shutdown()
+
+
+if __name__ == "__main__":
+    uvicorn.run("__main__:app", host=settings.HOST, port=settings.PORT, reload=True)
\ No newline at end of file
diff --git a/latencypredictor/requirements.txt b/latencypredictor/requirements.txt
new file mode 100644
index 0000000..30b6c54
--- /dev/null
+++ b/latencypredictor/requirements.txt
@@ -0,0 +1,12 @@
+fastapi
+uvicorn[standard]
+scikit-learn
+numpy
+pandas
+joblib
+river
+pydantic
+requests
+xgboost
+aiohttp
+lightgbm
diff --git a/latencypredictor/test_dual_server_client.py b/latencypredictor/test_dual_server_client.py
new file mode 100644
index 0000000..5ab9669
--- /dev/null
+++ b/latencypredictor/test_dual_server_client.py
@@ -0,0 +1,1412 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import time
+import asyncio
+import aiohttp
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from collections import defaultdict
+import random
+import requests
+import pytest
+import joblib
+import numpy as np
+import tempfile
+
+# Base URLs for the dual-server architecture
+
+PREDICTION_URL = os.getenv("PREDICTION_SERVER_URL", "http://<PREDICTION_IP>:80")  # Update this
+TRAINING_URL = os.getenv("TRAINING_SERVER_URL", "http://<TRAINING_IP>:8080")  # Update this
+
+TARGET_QPS = float(os.getenv("TARGET_QPS", 1000))  # Update this
+TARGET_QPS_LARGE_BATCH = float(os.getenv("TARGET_QPS_LARGE_BATCH", 100))  # Update this
+# Helper to wait until the servers are ready
+def wait_for_ready(url: str, timeout: float = 30.0, interval: float = 1.0):
+    start = time.time()
+    while True:
+        try:
+            r = requests.get(f"{url}/readyz", timeout=2.0)
+            if r.status_code == 200:
+                return
+        except requests.RequestException:
+            pass
+        if time.time() - start > timeout:
+            pytest.skip(f"Server at {url} did not become ready in time")
+        time.sleep(interval)
+
+@pytest.fixture(scope="module", autouse=True)
+def ensure_servers_ready():
+    """Wait for both servers to be ready before running tests."""
+    print("Waiting for prediction server...")
+    wait_for_ready(PREDICTION_URL)
+    print("Waiting for training server...")
+    wait_for_ready(TRAINING_URL)
+
+
+def test_prediction_server_healthz():
+    """Test prediction server health endpoint."""
+    r = requests.get(f"{PREDICTION_URL}/healthz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ok"
+
+
+def test_training_server_healthz():
+    """Test training server health endpoint."""
+    r = requests.get(f"{TRAINING_URL}/healthz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ok"
+
+
+def test_prediction_server_readyz():
+    """Test prediction server readiness."""
+    r = requests.get(f"{PREDICTION_URL}/readyz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ready"
+
+
+def test_training_server_readyz():
+    """Test training server readiness."""
+    r = requests.get(f"{TRAINING_URL}/readyz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ready"
+
+
+def test_prediction_server_status():
+    """Test prediction server status endpoint."""
+    r = requests.get(f"{PREDICTION_URL}/status")
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert "is_ready" in data
+    assert "model_type" in data
+    assert "models_exist" in data
+    assert "quantile" in data
+    assert data["model_type"] in ["bayesian_ridge", "xgboost", "lightgbm"]
+    assert 0 < data["quantile"] <= 1.0
+    
+    print(f"Prediction server using model type: {data['model_type']}")
+    print(f"Quantile: {data['quantile']}")
+    print(f"Models ready: {data['is_ready']}")
+    print(f"Models exist: {data['models_exist']}")
+
+
+def test_training_server_model_info():
+    """Test training server model info endpoint."""
+    r = requests.get(f"{TRAINING_URL}/model/download/info")
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert "model_type" in data
+    assert "available_endpoints" in data
+    assert data["model_type"] in ["bayesian_ridge", "xgboost", "lightgbm"]
+    
+    print(f"Training server using model type: {data['model_type']}")
+
+
+def test_training_server_models_list():
+    """Test training server models list endpoint."""
+    r = requests.get(f"{TRAINING_URL}/models/list")
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert "models" in data
+    assert "model_type" in data
+    assert "server_time" in data
+    
+    models = data["models"]
+    expected_models = ["ttft", "tpot"]
+    if data["model_type"] == "bayesian_ridge":
+        expected_models.extend(["ttft_scaler", "tpot_scaler"])
+    
+    for model_name in expected_models:
+        assert model_name in models, f"Model {model_name} should be listed"
+        print(f"Model {model_name}: exists={models[model_name]['exists']}, size={models[model_name]['size_bytes']} bytes")
+
+
+def test_model_download_from_training_server():
+    """Test downloading models from training server."""
+    # First check what models are available
+    models_r = requests.get(f"{TRAINING_URL}/models/list")
+    models_data = models_r.json()
+    
+    for model_name in ["ttft", "tpot"]:
+        if models_data["models"][model_name]["exists"]:
+            # Test model info endpoint
+            info_r = requests.get(f"{TRAINING_URL}/model/{model_name}/info")
+            assert info_r.status_code == 200
+            info_data = info_r.json()
+            assert info_data["exists"] == True
+            assert info_data["size_bytes"] > 0
+            
+            # Test model download with retry and streaming
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    download_r = requests.get(
+                        f"{TRAINING_URL}/model/{model_name}/download", 
+                        timeout=30,
+                        stream=True  # Use streaming to handle large files better
+                    )
+                    if download_r.status_code == 200:
+                        # Read content in chunks to avoid memory issues
+                        content_length = 0
+                        for chunk in download_r.iter_content(chunk_size=8192):
+                            content_length += len(chunk)
+                        
+                        assert content_length > 0, f"Downloaded {model_name} model is empty"
+                        print(f"Successfully downloaded {model_name} model ({content_length} bytes)")
+                        break
+                except requests.exceptions.ChunkedEncodingError as e:
+                    print(f"Download attempt {attempt + 1}/{max_retries} failed for {model_name}: {e}")
+                    if attempt == max_retries - 1:
+                        print(f"⚠️ Model download test skipped for {model_name} due to connection issues")
+                        # Don't fail the test - this might be a network/server issue
+                        continue
+                    time.sleep(2)  # Wait before retry
+
+def test_lightgbm_endpoints_on_training_server():
+    """Test LightGBM endpoints on training server if LightGBM is being used."""
+    model_info_r = requests.get(f"{TRAINING_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    if model_type != "lightgbm":
+        print("Skipping LightGBM endpoint tests - not using LightGBM model")
+        return
+    
+    print("Testing LightGBM endpoints on training server...")
+    
+    # Test TTFT model text format
+    ttft_txt_response = requests.get(f"{TRAINING_URL}/model/ttft/lgb/txt")
+    if ttft_txt_response.status_code == 200:
+        print("✓ TTFT LightGBM text model available")
+        assert ttft_txt_response.headers.get('content-type') == 'text/plain; charset=utf-8'
+    else:
+        print(f"TTFT LightGBM text model not yet available (status: {ttft_txt_response.status_code})")
+    
+    # Test TPOT model text format
+    tpot_txt_response = requests.get(f"{TRAINING_URL}/model/tpot/lgb/txt")
+    if tpot_txt_response.status_code == 200:
+        print("✓ TPOT LightGBM text model available")
+        assert tpot_txt_response.headers.get('content-type') == 'text/plain; charset=utf-8'
+    else:
+        print(f"TPOT LightGBM text model not yet available (status: {tpot_txt_response.status_code})")
+    
+    # Test TTFT feature importances
+    ttft_imp_response = requests.get(f"{TRAINING_URL}/model/ttft/lgb/importances")
+    if ttft_imp_response.status_code == 200:
+        ttft_importances = ttft_imp_response.json()
+        assert isinstance(ttft_importances, dict), "TTFT importances should be a dict"
+        
+        # Check for expected features including prefix_cache_score
+        expected_features = ["kv_cache_percentage", "input_token_length", "num_request_waiting", 
+                           "num_request_running", "prefix_cache_score"]
+        for feature in expected_features:
+            assert feature in ttft_importances, f"Missing feature importance: {feature}"
+        
+        print(f"✓ TTFT LightGBM importances available with {len(ttft_importances)} features")
+    else:
+        print(f"TTFT LightGBM importances not yet available (status: {ttft_imp_response.status_code})")
+    
+    # Test TPOT feature importances
+    tpot_imp_response = requests.get(f"{TRAINING_URL}/model/tpot/lgb/importances")
+    if tpot_imp_response.status_code == 200:
+        tpot_importances = tpot_imp_response.json()
+        assert isinstance(tpot_importances, dict), "TPOT importances should be a dict"
+        
+        # Check for expected features
+        expected_features = ["kv_cache_percentage", "input_token_length", "num_request_waiting", 
+                           "num_request_running", "num_tokens_generated"]
+        for feature in expected_features:
+            assert feature in tpot_importances, f"Missing feature importance: {feature}"
+        
+        print(f"✓ TPOT LightGBM importances available with {len(tpot_importances)} features")
+    else:
+        print(f"TPOT LightGBM importances not yet available (status: {tpot_imp_response.status_code})")
+        
+        
+def test_add_training_data_to_training_server():
+    """
+    Send training data to the training server.
+    The prediction server should eventually sync these models.
+    """
+    entries = []
+    
+    # Generate 50 training samples with known pattern
+    for i in range(1, 51):
+        waiting = i % 10 + 1
+        tokens = waiting
+        inp_len = 10 * i
+        kv = 0.5
+        running = 1
+        prefix_cache = random.uniform(0.1, 0.9)  # Added prefix_cache_score
+        
+        entries.append({
+            "kv_cache_percentage": kv,
+            "input_token_length": inp_len,
+            "num_request_waiting": waiting,
+            "num_request_running": running,
+            "actual_ttft_ms": (inp_len*2.0 + waiting*3.0 + running*4.0 + kv*50.0 + prefix_cache*30.0) + 95,  # Include prefix_cache effect
+            "actual_tpot_ms": (kv*100.0 + inp_len*0.5 + tokens*1.0 + running*5.0) + 9,
+            "num_tokens_generated": tokens,
+            "prefix_cache_score": prefix_cache,  # Added prefix_cache_score field
+        })
+
+    payload = {"entries": entries}
+    r = requests.post(f"{TRAINING_URL}/add_training_data_bulk", json=payload)
+    assert r.status_code == 202, f"Expected 202, got {r.status_code}"
+    assert r.json().get("message") == "Accepted 50 training samples."
+    
+    print("Successfully sent training data to training server")
+
+
+def test_prediction_server_model_sync():
+    """
+    Test that the prediction server can sync models from the training server.
+    This may take some time as models need to be downloaded.
+    """
+    # Trigger a manual reload on the prediction server
+    reload_r = requests.post(f"{PREDICTION_URL}/reload")
+    assert reload_r.status_code == 200
+    
+    reload_data = reload_r.json()
+    print(f"Model reload result: synced={reload_data.get('synced')}, loaded={reload_data.get('loaded')}")
+    
+    # Check status after reload
+    status_r = requests.get(f"{PREDICTION_URL}/status")
+    status_data = status_r.json()
+    
+    # Wait a bit for models to sync if they're not ready yet
+    max_wait = 60  # 60 seconds max wait
+    start_time = time.time()
+    
+    while not status_data.get("is_ready") and (time.time() - start_time) < max_wait:
+        print("Waiting for prediction server models to be ready...")
+        time.sleep(5)
+        
+        # Try reload again
+        requests.post(f"{PREDICTION_URL}/reload")
+        
+        status_r = requests.get(f"{PREDICTION_URL}/status")
+        status_data = status_r.json()
+    
+    assert status_data.get("is_ready"), f"Prediction server models not ready after {max_wait}s"
+    print("Prediction server models are ready!")
+
+
+def test_prediction_via_prediction_server():
+    """Test making predictions via the prediction server."""
+    features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 200,
+        "num_request_waiting": 4,
+        "num_request_running": 1,
+        "num_tokens_generated": 4,
+        "prefix_cache_score": 0.7,  # Added prefix_cache_score field
+    }
+    
+    r = requests.post(f"{PREDICTION_URL}/predict", json=features)
+    assert r.status_code == 200
+    
+    data = r.json()
+    required_fields = [
+        "ttft_ms", "tpot_ms", 
+        "predicted_at", "model_type", "last_model_load"
+    ]
+    
+    for field in required_fields:
+        assert field in data, f"Missing required field: {field}"
+    
+    # Verify predictions are reasonable
+    assert data["ttft_ms"] > 0
+    assert data["tpot_ms"] > 0
+    #assert data["ttft_uncertainty"] >= 0
+    #assert data["tpot_uncertainty"] >= 0
+    
+    print(f"Prediction successful: TTFT={data['ttft_ms']:.2f}ms, TPOT={data['tpot_ms']:.2f}ms")
+    print(f"Model type: {data['model_type']}")
+
+
+def test_bulk_prediction_strict():
+    """Test bulk predictions with strict error handling."""
+    print("Testing bulk prediction strict endpoint...")
+    
+    requests_data = [
+        {
+            "kv_cache_percentage": 0.5,
+            "input_token_length": 200,
+            "num_request_waiting": 4,
+            "num_request_running": 1,
+            "num_tokens_generated": 4,
+            "prefix_cache_score": 0.7,
+        },
+        {
+            "kv_cache_percentage": 0.3,
+            "input_token_length": 150,
+            "num_request_waiting": 2,
+            "num_request_running": 1,
+            "num_tokens_generated": 5,
+            "prefix_cache_score": 0.5,
+        }
+    ]
+    
+    bulk_request = {"requests": requests_data}
+    
+    r = requests.post(f"{PREDICTION_URL}/predict/bulk/strict", json=bulk_request)
+    assert r.status_code == 200
+    
+    data = r.json()
+    
+    # Check bulk response structure
+    assert "predictions" in data
+    assert "total_requests" in data
+    assert "successful_predictions" in data
+    assert "failed_predictions" in data
+    assert "processing_time_ms" in data
+    
+    assert len(data["predictions"]) == 2
+    assert data["total_requests"] == 2
+    assert data["successful_predictions"] == 2
+    assert data["failed_predictions"] == 0
+    
+    # Check individual prediction structure
+    for prediction in data["predictions"]:
+        assert "ttft_ms" in prediction
+        assert "tpot_ms" in prediction
+        #assert "ttft_uncertainty" in prediction
+        #assert "tpot_uncertainty" in prediction
+       #assert "ttft_prediction_bounds" in prediction
+        #assert "tpot_prediction_bounds" in prediction
+        assert "predicted_at" in prediction
+        assert "model_type" in prediction
+        assert "quantile" in prediction
+        
+    print("✓ Bulk prediction strict endpoint test passed")
+
+
+def test_bulk_prediction_with_validation_errors():
+    """Test that bulk predictions fail completely when any request has validation errors."""
+    print("Testing bulk prediction validation error handling...")
+    
+    requests_data = [
+        # Valid request
+        {
+            "kv_cache_percentage": 0.5,
+            "input_token_length": 200,
+            "num_request_waiting": 4,
+            "num_request_running": 1,
+            "num_tokens_generated": 4,
+            "prefix_cache_score": 0.7,
+        },
+        # Invalid request (missing prefix_cache_score)
+        {
+            "kv_cache_percentage": 0.3,
+            "input_token_length": 150,
+            "num_request_waiting": 2,
+            "num_request_running": 1,
+            "num_tokens_generated": 5,
+            # Missing prefix_cache_score
+        }
+    ]
+    
+    bulk_request = {"requests": requests_data}
+    
+    r = requests.post(f"{PREDICTION_URL}/predict/bulk", json=bulk_request)
+    assert r.status_code == 422  # Validation error expected
+    
+    # Check that error response contains validation details
+    error_data = r.json()
+    assert "detail" in error_data
+    
+    print("✓ Bulk prediction correctly failed when any request had validation errors")
+
+
+def test_bulk_prediction_all_valid():
+    """Test bulk predictions when all requests are valid."""
+    print("Testing bulk prediction with all valid requests...")
+    
+    requests_data = [
+        {
+            "kv_cache_percentage": 0.5,
+            "input_token_length": 200,
+            "num_request_waiting": 4,
+            "num_request_running": 1,
+            "num_tokens_generated": 4,
+            "prefix_cache_score": 0.7,
+        },
+        {
+            "kv_cache_percentage": 0.3,
+            "input_token_length": 150,
+            "num_request_waiting": 2,
+            "num_request_running": 1,
+            "num_tokens_generated": 5,
+            "prefix_cache_score": 0.5,  # Include required field
+        }
+    ]
+    
+    bulk_request = {"requests": requests_data}
+    
+    r = requests.post(f"{PREDICTION_URL}/predict/bulk", json=bulk_request)
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert data["total_requests"] == 2
+    assert data["successful_predictions"] == 2
+    assert data["failed_predictions"] == 0
+    
+    print("✓ Bulk prediction succeeded with all valid requests")
+
+def test_prediction_missing_prefix_cache_score():
+    """Test that predictions fail when prefix_cache_score is missing."""
+    features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 200,
+        "num_request_waiting": 4,
+        "num_request_running": 1,
+        "num_tokens_generated": 4,
+        # Missing prefix_cache_score
+    }
+    
+    r = requests.post(f"{PREDICTION_URL}/predict", json=features)
+    assert r.status_code == 422  # Should fail validation
+    
+    print("✓ Prediction correctly failed when prefix_cache_score was missing")
+
+
+def test_training_server_metrics():
+    """Test training server metrics endpoint."""
+    r = requests.get(f"{TRAINING_URL}/metrics")
+    assert r.status_code == 200
+    
+    content = r.text
+    
+    # Should contain model type metric
+    assert "model_type{" in content
+    
+    # Should contain either coefficients (Bayesian Ridge) or importance (XGBoost)
+    has_coef = "ttft_coef{" in content or "tpot_coef{" in content
+    has_importance = "ttft_importance{" in content or "tpot_importance{" in content
+    
+    assert has_coef or has_importance, "Should have either coefficients or feature importance metrics"
+    
+    # Should have standard metrics
+    assert "training_samples_count" in content
+    
+    # Check for prefix_cache_score in TTFT metrics
+    if has_coef:
+        assert 'feature="prefix_cache_score"' in content, "Should have prefix_cache_score coefficient for TTFT model"
+    if has_importance:
+        assert 'feature="prefix_cache_score"' in content, "Should have prefix_cache_score importance for TTFT model"
+    
+    print("Training server metrics endpoint working correctly")
+    print("✓ Prefix cache score feature found in metrics")
+
+
+def test_model_consistency_between_servers():
+    """Test that both servers report the same model type."""
+    # Get model type from training server
+    training_info_r = requests.get(f"{TRAINING_URL}/model/download/info")
+    training_model_type = training_info_r.json().get("model_type")
+    
+    # Get model type from prediction server
+    prediction_status_r = requests.get(f"{PREDICTION_URL}/status")
+    prediction_model_type = prediction_status_r.json().get("model_type")
+    
+    assert training_model_type == prediction_model_type, (
+        f"Model type mismatch: training={training_model_type}, prediction={prediction_model_type}"
+    )
+    
+    print(f"Model type consistent across servers: {training_model_type}")
+
+
+# 6. Update test_xgboost_tree_endpoints_on_training_server function name and add both
+def test_model_specific_endpoints_on_training_server():
+    """Test model-specific endpoints on training server based on model type."""
+    model_info_r = requests.get(f"{TRAINING_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    if model_type == "xgboost":
+        print("Testing XGBoost tree endpoints on training server...")
+        
+        # Test TTFT trees
+        ttft_response = requests.get(f"{TRAINING_URL}/model/ttft/xgb/json")
+        if ttft_response.status_code == 200:
+            ttft_trees = ttft_response.json()
+            assert isinstance(ttft_trees, list), "TTFT trees should be a list"
+            print(f"✓ TTFT XGBoost trees available: {len(ttft_trees)} trees")
+        else:
+            print(f"TTFT XGBoost trees not yet available (status: {ttft_response.status_code})")
+        
+        # Test TPOT trees  
+        tpot_response = requests.get(f"{TRAINING_URL}/model/tpot/xgb/json")
+        if tpot_response.status_code == 200:
+            tpot_trees = tpot_response.json()
+            assert isinstance(tpot_trees, list), "TPOT trees should be a list"
+            print(f"✓ TPOT XGBoost trees available: {len(tpot_trees)} trees")
+        else:
+            print(f"TPOT XGBoost trees not yet available (status: {tpot_response.status_code})")
+    
+    elif model_type == "lightgbm":
+        test_lightgbm_endpoints_on_training_server()
+    
+    else:
+        print(f"No model-specific endpoints to test for {model_type}")
+
+
+async def async_predict_request(session, payload, request_id):
+    """Make an async prediction request."""
+    start_time = time.time()
+    try:
+        async with session.post(f"{PREDICTION_URL}/predict", json=payload, timeout=aiohttp.ClientTimeout(total=5)) as response:
+            end_time = time.time()
+            response_data = await response.json()
+            return {
+                'request_id': request_id,
+                'status_code': response.status,
+                'response_time': end_time - start_time,
+                'success': response.status == 200,
+                'response_data': response_data,
+                'model_type': response_data.get('model_type') if response.status == 200 else None
+            }
+    except Exception as e:
+        end_time = time.time()
+        return {
+            'request_id': request_id,
+            'status_code': 0,
+            'response_time': end_time - start_time,
+            'success': False,
+            'error': str(e),
+            'model_type': None
+        }
+
+
+async def async_bulk_predict_request(session, payload, request_id):
+    """Make an async bulk prediction request."""
+    start_time = time.time()
+    try:
+        async with session.post(f"{PREDICTION_URL}/predict/bulk/strict", json=payload, timeout=aiohttp.ClientTimeout(total=10)) as response:
+            end_time = time.time()
+            response_data = await response.json()
+            return {
+                'request_id': request_id,
+                'status_code': response.status,
+                'response_time': end_time - start_time,
+                'success': response.status == 200,
+                'response_data': response_data,
+                'batch_size': len(payload.get('requests', [])),
+                'predictions_count': len(response_data.get('predictions', [])) if response.status == 200 else 0
+            }
+    except Exception as e:
+        end_time = time.time()
+        return {
+            'request_id': request_id,
+            'status_code': 0,
+            'response_time': end_time - start_time,
+            'success': False,
+            'error': str(e),
+            'batch_size': len(payload.get('requests', [])),
+            'predictions_count': 0
+        }
+
+
+def generate_random_prediction_payload():
+    """Generate a random prediction payload."""
+    return {
+        "kv_cache_percentage": random.uniform(0.1, 0.9),
+        "input_token_length": random.randint(10, 1000),
+        "num_request_waiting": random.randint(1, 20),
+        "num_request_running": random.randint(1, 10),
+        "num_tokens_generated": random.randint(1, 20),
+        "prefix_cache_score": random.uniform(0.0, 1.0),
+    }
+
+
+def generate_bulk_prediction_payload(batch_size=10):
+    """Generate a bulk prediction payload with specified batch size."""
+    requests_data = []
+    for _ in range(batch_size):
+        requests_data.append({
+            "kv_cache_percentage": random.uniform(0.1, 0.9),
+            "input_token_length": random.randint(10, 1000),
+            "num_request_waiting": random.randint(1, 20),
+            "num_request_running": random.randint(1, 10),
+            "num_tokens_generated": random.randint(1, 20),
+            "prefix_cache_score": random.uniform(0.0, 1.0),
+        })
+    return {"requests": requests_data}
+
+
+def generate_random_training_payload():
+    """Generate a random training payload."""
+    input_tokens = random.randint(10, 1000)
+    waiting_requests = random.randint(1, 20)
+    running_requests = random.randint(1, 10)
+    kv = random.uniform(0.01, 0.99)
+    tokens_generated = random.randint(1, 20)
+    prefix_cache = random.uniform(0.0, 1.0)  # Added prefix cache score
+    
+    return {
+        "kv_cache_percentage": kv,
+        "input_token_length": input_tokens,
+        "num_request_waiting": waiting_requests,
+        "num_request_running": running_requests,
+        "actual_ttft_ms": (
+            input_tokens * 2.0
+            + waiting_requests * 3.0
+            + running_requests * 4.0
+            + kv * 50.0
+            + prefix_cache * 30.0  # Added prefix cache effect
+            + 95 + random.uniform(-10, 10)
+        ),
+        "actual_tpot_ms": (
+            kv * 100.0
+            + input_tokens * 0.5
+            + tokens_generated * 1.0
+            + running_requests * 5.0
+            + 9 + random.uniform(-5, 5)
+        ),
+        "num_tokens_generated": tokens_generated,
+        "prefix_cache_score": prefix_cache,  # Added prefix cache score
+    }
+
+
+def test_dual_server_quantile_regression_learns_distribution():
+    """
+    Quantile regression should learn the q-quantile of a Gaussian residual model
+    with fixed sigma, verified by (a) relative error vs μ+zσ and (b) empirical coverage.
+    """
+    import random, time, math
+    import numpy as np
+    import requests
+    from scipy.stats import norm
+
+    RNG_SEED = 42
+    random.seed(RNG_SEED)
+    np.random.seed(RNG_SEED)
+
+    # Config
+    TRAIN_N = 3000
+    TEST_N  = 200
+    TTFT_STD, TPOT_STD = 20.0, 10.0
+    REL_ERR_TOL = 0.15  # 15%
+    COVERAGE_TOL = 0.05 # ±5% around target quantile
+    MAX_WAIT_S = 180
+    POLL_INTERVAL_S = 3
+
+    # 1) Confirm server mode
+    r = requests.get(f"{TRAINING_URL}/model/download/info", timeout=10)
+    assert r.status_code == 200, "model info endpoint failed"
+    model_type = r.json().get("model_type", "unknown")
+
+    s = requests.get(f"{PREDICTION_URL}/status", timeout=10)
+    assert s.status_code == 200, "prediction status endpoint failed"
+    target_quantile = float(s.json().get("quantile", 0.9))
+
+    assert "xgboost" in model_type.lower() or "lightgbm" in model_type.lower(), f"Model not in quantile mode: {model_type}"
+
+    z = norm.ppf(target_quantile)
+
+    # 2) Generate training data (vectorized)
+    kv = np.random.uniform(0.1, 0.9, size=TRAIN_N)
+    input_len = np.random.randint(50, 801, size=TRAIN_N)
+    waiting = np.random.randint(0, 9, size=TRAIN_N)
+    running = np.random.randint(1, 5, size=TRAIN_N)
+    tokens_gen = np.random.randint(1, 26, size=TRAIN_N)
+    prefix = np.random.uniform(0.0, 1.0, size=TRAIN_N)
+
+    ttft_mu = (input_len*2.0 + waiting*3.0 + running*4.0 + kv*50.0 + prefix*30.0 + 95)
+    tpot_mu = (kv*100.0 + input_len*0.5 + tokens_gen*1.0 + running*5.0 + 9)
+
+    ttft_y = np.maximum(1.0, ttft_mu + np.random.normal(0, TTFT_STD, size=TRAIN_N))
+    tpot_y = np.maximum(1.0, tpot_mu + np.random.normal(0, TPOT_STD, size=TRAIN_N))
+
+    entries = [dict(
+        kv_cache_percentage=float(kv[i]),
+        input_token_length=int(input_len[i]),
+        num_request_waiting=int(waiting[i]),
+        num_request_running=int(running[i]),
+        actual_ttft_ms=float(ttft_y[i]),
+        actual_tpot_ms=float(tpot_y[i]),
+        num_tokens_generated=int(tokens_gen[i]),
+        prefix_cache_score=float(prefix[i]),
+    ) for i in range(TRAIN_N)]
+
+    # 3) Submit training data (with a couple retries)
+    for _ in range(3):
+        tr = requests.post(f"{TRAINING_URL}/add_training_data_bulk", json={"entries": entries}, timeout=60)
+        if tr.status_code == 202:
+            break
+        time.sleep(2)
+    assert tr.status_code == 202, f"training submit failed: {tr.status_code}"
+
+    # 4) Wait for training to complete
+    time.sleep(30)
+    # 5) Sync models to prediction server
+    synced = False
+    for _ in range(10):
+        rr = requests.post(f"{PREDICTION_URL}/reload", timeout=20)
+        if rr.status_code == 200 and rr.json().get("is_ready"):
+            synced = True
+            break
+        time.sleep(3)
+    assert synced, "Failed to sync models"
+
+    # 6) Build test set + expected quantiles
+    kv_t = np.random.uniform(0.1, 0.9, size=TEST_N)
+    in_t = np.random.randint(100, 601, size=TEST_N)
+    wait_t = np.random.randint(1, 9, size=TEST_N)
+    run_t = np.random.randint(1, 5, size=TEST_N)
+    tok_t = np.random.randint(5, 21, size=TEST_N)
+    pre_t = np.random.uniform(0.0, 1.0, size=TEST_N)
+
+    ttft_mu_t = (in_t*2.0 + wait_t*3.0 + run_t*4.0 + kv_t*50.0 + pre_t*30.0 + 95)
+    tpot_mu_t = (kv_t*100.0 + in_t*0.5 + tok_t*1.0 + run_t*5.0 + 9)
+    ttft_q_exp = ttft_mu_t + z*TTFT_STD
+    tpot_q_exp = tpot_mu_t + z*TPOT_STD
+
+    test_cases = [dict(
+        kv_cache_percentage=float(kv_t[i]),
+        input_token_length=int(in_t[i]),
+        num_request_waiting=int(wait_t[i]),
+        num_request_running=int(run_t[i]),
+        num_tokens_generated=int(tok_t[i]),
+        prefix_cache_score=float(pre_t[i]),
+    ) for i in range(TEST_N)]
+
+    # 7) Predict (bulk)
+    pr = requests.post(f"{PREDICTION_URL}/predict/bulk/strict", json={"requests": test_cases}, timeout=60)
+    assert pr.status_code == 200, f"predict failed: {pr.status_code}"
+    jd = pr.json()
+    assert jd["total_requests"] == TEST_N and jd["successful_predictions"] == TEST_N and jd["failed_predictions"] == 0
+    preds = jd["predictions"]
+
+    ttft_pred = np.array([p["ttft_ms"] for p in preds], dtype=float)
+    tpot_pred = np.array([p["tpot_ms"] for p in preds], dtype=float)
+
+    # 8) Relative error vs μ + zσ
+    ttft_rel_err = np.abs(ttft_pred - ttft_q_exp) / ttft_q_exp
+    tpot_rel_err = np.abs(tpot_pred - tpot_q_exp) / tpot_q_exp
+    acc_mask = (ttft_rel_err <= REL_ERR_TOL) & (tpot_rel_err <= REL_ERR_TOL)
+    rel_accuracy = acc_mask.mean()
+    print(f"Relative-err accuracy (≤{int(REL_ERR_TOL*100)}%): {rel_accuracy*100:.1f}%")
+
+    # 9) Coverage calibration (simulate actuals for the same test X)
+    # Generate fresh noise so it's an *unseen* draw from the same D|X:
+    ttft_actual = np.maximum(1.0, ttft_mu_t + np.random.normal(0, TTFT_STD, size=TEST_N))
+    tpot_actual = np.maximum(1.0, tpot_mu_t + np.random.normal(0, TPOT_STD, size=TEST_N))
+
+    ttft_cov = (ttft_actual <= ttft_pred).mean()
+    tpot_cov = (tpot_actual <= tpot_pred).mean()
+    print(f"Coverage: TTFT={ttft_cov:.3f}, TPOT={tpot_cov:.3f} (target {target_quantile:.3f} ± {COVERAGE_TOL})")
+
+    # 10) Monotonic sanity checks on a few random pairs (no hard fail, just helpful asserts)
+    # pick one sample index and perturb input_token_length upward
+    idx = 0
+    base = test_cases[idx].copy(); up = test_cases[idx].copy(); up["input_token_length"] += 100
+    br = requests.post(f"{PREDICTION_URL}/predict/bulk/strict", json={"requests":[base, up]}, timeout=30)
+    if br.status_code == 200:
+        _bp = br.json()["predictions"]
+        assert _bp[1]["ttft_ms"] >= _bp[0]["ttft_ms"] - 1e-6, "TTFT should not decrease with longer input"
+
+    # 11) Final assertions
+    assert rel_accuracy >= 0.70, f"Only {rel_accuracy*100:.1f}% within ±{int(REL_ERR_TOL*100)}% (expected ≥70%)"
+    assert abs(ttft_cov - target_quantile) <= COVERAGE_TOL, f"TTFT coverage {ttft_cov:.3f} not within ±{COVERAGE_TOL} of {target_quantile:.3f}"
+    assert abs(tpot_cov - target_quantile) <= COVERAGE_TOL, f"TPOT coverage {tpot_cov:.3f} not within ±{COVERAGE_TOL} of {target_quantile:.3f}"
+
+
+
+
+async def run_prediction_stress_test(duration_seconds=30, target_qps=1000):
+    """Run stress test against the prediction server only."""
+    interval = 1.0 / target_qps
+    start = time.time()
+    connector = aiohttp.TCPConnector(limit=1000, limit_per_host=1000)
+    
+    async with aiohttp.ClientSession(connector=connector) as session:
+        tasks = []
+        req_id = 0
+        next_time = start
+        
+        while time.time() - start < duration_seconds:
+            now = time.time()
+            while next_time <= now:
+                req_id += 1
+                payload = generate_random_prediction_payload()
+                tasks.append(asyncio.create_task(async_predict_request(session, payload, req_id)))
+                next_time += interval
+            
+            await asyncio.sleep(0.001)
+        
+        print(f"Waiting for {len(tasks)} prediction requests to complete...")
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        valid_results = [r for r in results if isinstance(r, dict)]
+        
+        if valid_results:
+            actual_qps = len(valid_results) / duration_seconds
+            print(f"Target QPS: {target_qps}, Actual QPS: {actual_qps:.1f}")
+        
+        return valid_results
+
+
+async def run_bulk_prediction_stress_test(duration_seconds=30, target_rps=100, batch_size=10):
+    """Run stress test against the bulk prediction endpoint."""
+    interval = 1.0 / target_rps  # requests per second
+    start = time.time()
+    connector = aiohttp.TCPConnector(limit=200, limit_per_host=200)
+    
+    async with aiohttp.ClientSession(connector=connector) as session:
+        tasks = []
+        req_id = 0
+        next_time = start
+        
+        while time.time() - start < duration_seconds:
+            now = time.time()
+            while next_time <= now:
+                req_id += 1
+                payload = generate_bulk_prediction_payload(batch_size)
+                tasks.append(asyncio.create_task(async_bulk_predict_request(session, payload, req_id)))
+                next_time += interval
+            
+            await asyncio.sleep(0.01)  # Slightly longer sleep for bulk requests
+        
+        print(f"Waiting for {len(tasks)} bulk prediction requests to complete...")
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        valid_results = [r for r in results if isinstance(r, dict)]
+        
+        if valid_results:
+            actual_rps = len(valid_results) / duration_seconds
+            total_predictions = sum(r.get('predictions_count', 0) for r in valid_results)
+            actual_pps = total_predictions / duration_seconds  # predictions per second
+            print(f"Target RPS: {target_rps}, Actual RPS: {actual_rps:.1f}")
+            print(f"Total Predictions: {total_predictions}, Predictions/sec: {actual_pps:.1f}")
+        
+        return valid_results
+
+
+def analyze_prediction_stress_results(results):
+    """Analyze prediction stress test results."""
+    if not results:
+        print("No results to analyze")
+        return
+    
+    total_requests = len(results)
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    failed_requests = total_requests - successful_requests
+    
+    response_times = [r['response_time'] for r in results if r.get('response_time')]
+    avg_response_time = sum(response_times) / len(response_times) if response_times else 0
+    
+    status_codes = defaultdict(int)
+    for r in results:
+        status_codes[r.get('status_code', 0)] += 1
+    
+    model_types = defaultdict(int)
+    for r in results:
+        if r.get('model_type'):
+            model_types[r['model_type']] += 1
+    
+    print(f"\n{'='*50}")
+    print("PREDICTION SERVER STRESS TEST RESULTS")
+    print(f"{'='*50}")
+    print(f"Total Requests: {total_requests}")
+    print(f"Successful: {successful_requests} ({successful_requests/total_requests*100:.1f}%)")
+    print(f"Failed: {failed_requests} ({failed_requests/total_requests*100:.1f}%)")
+    print(f"Average Response Time: {avg_response_time*1000:.2f}ms")
+    
+    if model_types:
+        print(f"\nModel Types in Predictions:")
+        for model_type, count in model_types.items():
+            print(f"  {model_type}: {count}")
+    
+    print(f"\nStatus Code Distribution:")
+    for status, count in status_codes.items():
+        print(f"  {status}: {count}")
+    
+    if response_times:
+        sorted_times = sorted(response_times)
+        p50 = sorted_times[int(len(sorted_times) * 0.5)] * 1000
+        p95 = sorted_times[int(len(sorted_times) * 0.95)] * 1000
+        p99 = sorted_times[int(len(sorted_times) * 0.99)] * 1000
+        print(f"\nResponse Time Percentiles:")
+        print(f"  P50: {p50:.2f}ms")
+        print(f"  P95: {p95:.2f}ms")
+        print(f"  P99: {p99:.2f}ms")
+
+
+def analyze_bulk_prediction_stress_results(results):
+    """Analyze bulk prediction stress test results."""
+    if not results:
+        print("No results to analyze")
+        return
+    
+    total_requests = len(results)
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    failed_requests = total_requests - successful_requests
+    
+    total_predictions = sum(r.get('predictions_count', 0) for r in results)
+    total_batch_size = sum(r.get('batch_size', 0) for r in results)
+    
+    response_times = [r['response_time'] for r in results if r.get('response_time')]
+    avg_response_time = sum(response_times) / len(response_times) if response_times else 0
+    
+    status_codes = defaultdict(int)
+    for r in results:
+        status_codes[r.get('status_code', 0)] += 1
+    
+    print(f"\n{'='*50}")
+    print("BULK PREDICTION STRESS TEST RESULTS")
+    print(f"{'='*50}")
+    print(f"Total Bulk Requests: {total_requests}")
+    print(f"Successful: {successful_requests} ({successful_requests/total_requests*100:.1f}%)")
+    print(f"Failed: {failed_requests} ({failed_requests/total_requests*100:.1f}%)")
+    print(f"Total Individual Predictions: {total_predictions}")
+    print(f"Total Batch Size: {total_batch_size}")
+    print(f"Average Response Time: {avg_response_time*1000:.2f}ms")
+    
+    if total_batch_size > 0:
+        print(f"Average Batch Size: {total_batch_size/total_requests:.1f}")
+        print(f"Prediction Success Rate: {total_predictions/total_batch_size*100:.1f}%")
+    
+    print(f"\nStatus Code Distribution:")
+    for status, count in status_codes.items():
+        print(f"  {status}: {count}")
+    
+    if response_times:
+        sorted_times = sorted(response_times)
+        p50 = sorted_times[int(len(sorted_times) * 0.5)] * 1000
+        p95 = sorted_times[int(len(sorted_times) * 0.95)] * 1000
+        p99 = sorted_times[int(len(sorted_times) * 0.99)] * 1000
+        print(f"\nResponse Time Percentiles:")
+        print(f"  P50: {p50:.2f}ms")
+        print(f"  P95: {p95:.2f}ms")
+        print(f"  P99: {p99:.2f}ms")
+
+
+def test_prediction_server_stress_test():
+    """Stress test the prediction server."""
+    print("Running prediction server stress test...")
+    
+    results = asyncio.run(run_prediction_stress_test(duration_seconds=100, target_qps=TARGET_QPS))
+    
+    analyze_prediction_stress_results(results)
+    
+    assert len(results) > 0, "No requests were made"
+    
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    success_rate = successful_requests / len(results)
+    
+    assert success_rate > 0.8, f"Success rate too low: {success_rate*100:.1f}%"
+    
+    print(f"Prediction server stress test completed with {success_rate*100:.1f}% success rate")
+
+
+def test_bulk_prediction_stress_test():
+    """Stress test the bulk prediction endpoint."""
+    print("Running bulk prediction stress test...")
+    
+    # Test with different batch sizes
+    batch_sizes = [5, 10, 25]
+    for batch_size in batch_sizes:
+        print(f"\nTesting with batch size {batch_size}...")
+        results = asyncio.run(run_bulk_prediction_stress_test(
+            duration_seconds=100, 
+            target_rps=TARGET_QPS,  # Lower RPS for bulk requests
+            batch_size=batch_size
+        ))
+        
+        analyze_bulk_prediction_stress_results(results)
+        
+        assert len(results) > 0, f"No bulk requests were made for batch size {batch_size}"
+        
+        successful_requests = sum(1 for r in results if r.get('success', False))
+        success_rate = successful_requests / len(results)
+        
+        assert success_rate > 0.7, f"Bulk success rate too low for batch size {batch_size}: {success_rate*100:.1f}%"
+        
+        print(f"Bulk prediction stress test (batch size {batch_size}) completed with {success_rate*100:.1f}% success rate")
+
+def test_large_batch_prediction_stress_test():
+    """Stress test the bulk prediction endpoint."""
+    print("Running bulk prediction stress test...")
+    
+    # Test with different batch sizes
+    batch_sizes = [1000]
+    for batch_size in batch_sizes:
+        print(f"\nTesting with batch size {batch_size}...")
+        results = asyncio.run(run_bulk_prediction_stress_test(
+            duration_seconds=100, 
+            target_rps=TARGET_QPS_LARGE_BATCH,  # Lower RPS for bulk requests
+            batch_size=batch_size
+        ))
+        
+        analyze_bulk_prediction_stress_results(results)
+        
+        assert len(results) > 0, f"No bulk requests were made for batch size {batch_size}"
+        
+        successful_requests = sum(1 for r in results if r.get('success', False))
+        success_rate = successful_requests / len(results)
+        
+        assert success_rate > 0.7, f"Bulk success rate too low for batch size {batch_size}: {success_rate*100:.1f}%"
+        
+        print(f"Bulk prediction stress test (batch size {batch_size}) completed with {success_rate*100:.1f}% success rate")
+
+
+def test_end_to_end_workflow():
+    """Test the complete end-to-end workflow with robust error handling."""
+    print("Testing end-to-end workflow...")
+    
+    # 1. Send training data to training server
+    print("Step 1: Sending training data to training server...")
+    training_payload = {"entries": [generate_random_training_payload() for _ in range(20)]}
+    
+    try:
+        training_r = requests.post(f"{TRAINING_URL}/add_training_data_bulk", json=training_payload, timeout=30)
+        assert training_r.status_code == 202
+    except requests.exceptions.RequestException as e:
+        pytest.skip(f"Training server not accessible: {e}")
+
+    # 2. Wait a bit for training
+    print("Step 2: Waiting for training...")
+    time.sleep(10)
+
+    # 3. Trigger model sync on prediction server
+    print("Step 3: Syncing models to prediction server...")
+    try:
+        reload_r = requests.post(f"{PREDICTION_URL}/reload", timeout=30)
+        assert reload_r.status_code == 200
+        time.sleep(5)  # Allow some time for models to sync
+    except requests.exceptions.RequestException as e:
+        pytest.skip(f"Prediction server not accessible for reload: {e}")
+
+    # 4. Make predictions with retry logic
+    print("Step 4: Making predictions...")
+    successful_predictions = 0
+    
+    for i in range(5):
+        payload = generate_random_prediction_payload()
+        max_retries = 3
+        
+        for attempt in range(max_retries):
+            try:
+                pred_r = requests.post(f"{PREDICTION_URL}/predict", json=payload, timeout=15)
+                if pred_r.status_code == 200:
+                    successful_predictions += 1
+                    pred_data = pred_r.json()
+                    print(f"  Prediction {i+1}: TTFT={pred_data['ttft_ms']:.2f}ms, TPOT={pred_data['tpot_ms']:.2f}ms (prefix_cache={payload['prefix_cache_score']:.2f})")
+                    break
+                else:
+                    print(f"  Prediction {i+1} attempt {attempt+1} failed with status {pred_r.status_code}")
+            except requests.exceptions.ConnectTimeout:
+                print(f"  Prediction {i+1} attempt {attempt+1} timed out")
+                if attempt < max_retries - 1:
+                    time.sleep(2)  # Wait before retry
+                else:
+                    print(f"  Prediction {i+1} failed after {max_retries} attempts")
+            except requests.exceptions.RequestException as e:
+                print(f"  Prediction {i+1} attempt {attempt+1} failed: {e}")
+                break
+    
+    # Accept partial success if servers are having issues
+    if successful_predictions == 0:
+        pytest.skip("All prediction requests failed - servers may be down")
+    elif successful_predictions < 5:
+        print(f"⚠️ Partial success: {successful_predictions}/5 predictions succeeded")
+    else:
+        print("✓ End-to-end workflow completed successfully!")
+
+
+def test_server_configuration():
+    """Test server configuration and setup."""
+    print("Testing server configuration...")
+    
+    # Test prediction server root endpoint
+    pred_root_r = requests.get(f"{PREDICTION_URL}/")
+    assert pred_root_r.status_code == 200
+    pred_root_data = pred_root_r.json()
+    print(f"Prediction server: {pred_root_data.get('message')}")
+    print(f"  Model type: {pred_root_data.get('model_type')}")
+    print(f"  Is ready: {pred_root_data.get('is_ready')}")
+    print(f"  Sync interval: {pred_root_data.get('sync_interval')}s")
+    print(f"  Training server URL: {pred_root_data.get('training_server')}")
+    
+    # Test training server root endpoint  
+    train_root_r = requests.get(f"{TRAINING_URL}/")
+    assert train_root_r.status_code == 200
+    train_root_data = train_root_r.json()
+    print(f"Training server: {train_root_data.get('message')}")
+    print(f"  Model type: {train_root_data.get('model_type')}")
+
+def test_training_server_flush_api():
+    """Test the training server flush API and data status endpoint."""
+    print("Testing training server flush API...")
+    
+    # 1. Check initial data status
+    print("Step 1: Checking initial data status...")
+    initial_status_r = requests.get(f"{TRAINING_URL}/data/status")
+    assert initial_status_r.status_code == 200
+    initial_status = initial_status_r.json()
+    
+    print(f"  Initial training samples: TTFT={initial_status['training_data']['ttft_samples']}, "
+          f"TPOT={initial_status['training_data']['tpot_samples']}")
+    print(f"  Initial test samples: TTFT={initial_status['test_data']['ttft_samples']}, "
+          f"TPOT={initial_status['test_data']['tpot_samples']}")
+    
+    # 2. Add training data
+    print("Step 2: Adding training data...")
+    training_entries = [generate_random_training_payload() for _ in range(100)]
+    training_payload = {"entries": training_entries}
+    
+    add_r = requests.post(f"{TRAINING_URL}/add_training_data_bulk", json=training_payload)
+    assert add_r.status_code == 202
+    print(f"  Added 100 training samples")
+    
+    # Wait a bit for data to be processed
+    time.sleep(2)
+    
+    # 3. Verify data was added
+    print("Step 3: Verifying data was added...")
+    after_add_status_r = requests.get(f"{TRAINING_URL}/data/status")
+    assert after_add_status_r.status_code == 200
+    after_add_status = after_add_status_r.json()
+    
+    total_samples_after = after_add_status['training_data']['total_samples'] + after_add_status['test_data']['total_samples']
+    print(f"  After adding - Training: {after_add_status['training_data']['total_samples']}, "
+          f"Test: {after_add_status['test_data']['total_samples']}, Total: {total_samples_after}")
+    
+    # Should have more data now (some goes to training, some to test based on TEST_TRAIN_RATIO)
+    assert total_samples_after > 0, "No samples were added"
+    
+    # 4. Test flush with only training data
+    print("Step 4: Testing flush with only training data...")
+    flush_training_only = {
+        "flush_training_data": True,
+        "flush_test_data": False,
+        "flush_metrics": False,
+        "reason": "Test flush training data only"
+    }
+    
+    flush_r = requests.post(f"{TRAINING_URL}/flush", json=flush_training_only)
+    assert flush_r.status_code == 200
+    flush_response = flush_r.json()
+    
+    assert flush_response["success"] == True
+    assert flush_response["metrics_cleared"] == False
+    assert flush_response["reason"] == "Test flush training data only"
+    
+    print(f"  Flushed {flush_response['ttft_training_samples_flushed']} TTFT training samples")
+    print(f"  Flushed {flush_response['tpot_training_samples_flushed']} TPOT training samples")
+    print(f"  Test samples flushed: {flush_response['ttft_test_samples_flushed']} TTFT, "
+          f"{flush_response['tpot_test_samples_flushed']} TPOT (should be 0)")
+    
+    # Verify training data was flushed but test data remains
+    after_flush_training_r = requests.get(f"{TRAINING_URL}/data/status")
+    after_flush_training = after_flush_training_r.json()
+    
+    assert after_flush_training['training_data']['total_samples'] == 0, "Training data should be empty"
+    # Test data should still exist if any was added
+    print(f"  After training flush - Training: {after_flush_training['training_data']['total_samples']}, "
+          f"Test: {after_flush_training['test_data']['total_samples']}")
+    
+    # 5. Add more data
+    print("Step 5: Adding more training data...")
+    more_entries = [generate_random_training_payload() for _ in range(50)]
+    requests.post(f"{TRAINING_URL}/add_training_data_bulk", json={"entries": more_entries})
+    time.sleep(2)
+    
+    # 6. Test flush everything
+    print("Step 6: Testing flush everything...")
+    flush_all = {
+        "flush_training_data": True,
+        "flush_test_data": True,
+        "flush_metrics": True,
+        "reason": "Complete flush test"
+    }
+    
+    flush_all_r = requests.post(f"{TRAINING_URL}/flush", json=flush_all)
+    assert flush_all_r.status_code == 200
+    flush_all_response = flush_all_r.json()
+    
+    assert flush_all_response["success"] == True
+    assert flush_all_response["metrics_cleared"] == True
+    assert "Successfully flushed" in flush_all_response["message"]
+    
+    print(f"  Complete flush message: {flush_all_response['message']}")
+    
+    # Verify everything was flushed
+    after_flush_all_r = requests.get(f"{TRAINING_URL}/data/status")
+    after_flush_all = after_flush_all_r.json()
+    
+    assert after_flush_all['training_data']['total_samples'] == 0, "Training data should be empty"
+    assert after_flush_all['test_data']['total_samples'] == 0, "Test data should be empty"
+    
+    print(f"  After complete flush - Training: {after_flush_all['training_data']['total_samples']}, "
+          f"Test: {after_flush_all['test_data']['total_samples']}")
+    
+    # 7. Test flush with default parameters (should flush everything)
+    print("Step 7: Testing default flush (no body)...")
+    
+    # Add some data first
+    requests.post(f"{TRAINING_URL}/add_training_data_bulk", 
+                 json={"entries": [generate_random_training_payload() for _ in range(20)]})
+    time.sleep(1)
+    
+    # Flush with empty body (uses defaults)
+    default_flush_r = requests.post(f"{TRAINING_URL}/flush")
+    assert default_flush_r.status_code == 200
+    default_flush_response = default_flush_r.json()
+    
+    assert default_flush_response["success"] == True
+    print(f"  Default flush result: {default_flush_response['message']}")
+    
+    # 8. Test flush with only test data
+    print("Step 8: Testing flush with only test data...")
+    
+    # Add data
+    requests.post(f"{TRAINING_URL}/add_training_data_bulk",
+                 json={"entries": [generate_random_training_payload() for _ in range(50)]})
+    time.sleep(2)
+    
+    # Get status before
+    before_test_flush_r = requests.get(f"{TRAINING_URL}/data/status")
+    before_test_flush = before_test_flush_r.json()
+    
+    # Flush only test data
+    flush_test_only = {
+        "flush_training_data": False,
+        "flush_test_data": True,
+        "flush_metrics": False,
+        "reason": "Test flush test data only"
+    }
+    
+    flush_test_r = requests.post(f"{TRAINING_URL}/flush", json=flush_test_only)
+    assert flush_test_r.status_code == 200
+    flush_test_response = flush_test_r.json()
+    
+    print(f"  Test data flush: {flush_test_response['ttft_test_samples_flushed']} TTFT, "
+          f"{flush_test_response['tpot_test_samples_flushed']} TPOT")
+    
+    # Verify only test data was flushed
+    after_test_flush_r = requests.get(f"{TRAINING_URL}/data/status")
+    after_test_flush = after_test_flush_r.json()
+    
+    assert after_test_flush['test_data']['total_samples'] == 0, "Test data should be empty"
+    # Training data should still exist
+    print(f"  After test flush - Training: {after_test_flush['training_data']['total_samples']}, "
+          f"Test: {after_test_flush['test_data']['total_samples']}")
+    
+    # 9. Test bucket distribution in status
+    print("Step 9: Testing bucket distribution in status...")
+    if "bucket_distribution" in after_flush_all:
+        print(f"  Bucket distribution available: {len(after_flush_all.get('bucket_distribution', {}))} buckets with data")
+    
+    print("✓ Flush API tests passed!")
+
+
+def test_training_server_flush_error_handling():
+    """Test error handling in flush API."""
+    print("Testing flush API error handling...")
+    
+    # Test with invalid JSON
+    invalid_json = '{"flush_training_data": "not_a_boolean"}'
+    headers = {'Content-Type': 'application/json'}
+    
+    try:
+        r = requests.post(f"{TRAINING_URL}/flush", data=invalid_json, headers=headers)
+        # Should get validation error
+        assert r.status_code in [400, 422], f"Expected 400 or 422, got {r.status_code}"
+        print("✓ Invalid JSON handled correctly")
+    except Exception as e:
+        print(f"⚠️ Error handling test skipped: {e}")
+    
+    # Test with valid parameters
+    valid_flush = {
+        "flush_training_data": False,
+        "flush_test_data": False,
+        "flush_metrics": True,
+        "reason": "Metrics only flush"
+    }
+    
+    r = requests.post(f"{TRAINING_URL}/flush", json=valid_flush)
+    assert r.status_code == 200
+    response = r.json()
+    assert response["metrics_cleared"] == True
+    assert response["ttft_training_samples_flushed"] == 0
+    assert response["tpot_training_samples_flushed"] == 0
+    
+    print("✓ Flush error handling tests passed!")
+
+if __name__ == "__main__":
+    print("Running dual-server architecture tests with prefix cache score support...")
+    print(f"Prediction server: {PREDICTION_URL}")
+    print(f"Training server: {TRAINING_URL}")
+    
+    # Update these URLs before running!
+    if "<PREDICTION_EXTERNAL_IP>" in PREDICTION_URL or "<TRAINING_EXTERNAL_IP>" in TRAINING_URL:
+        print("\n❌ ERROR: Please update the server URLs at the top of this file!")
+        print("Get external IPs with: kubectl get services")
+        exit(1)
+    
+    # Run individual tests
+    print("\n" + "="*50)
+    print("RUNNING DUAL-SERVER TESTS WITH PREFIX CACHE SCORE")
+    print("="*50)
+    
+    tests = [
+        ("Server Health Checks", lambda: (test_prediction_server_healthz(), test_training_server_healthz())),
+        ("Server Readiness", lambda: (test_prediction_server_readyz(), test_training_server_readyz())),
+        ("Server Configuration", test_server_configuration),
+        ("Prediction Server Status", test_prediction_server_status),
+        ("Training Server Model Info", test_training_server_model_info),
+        ("Training Server Models List", test_training_server_models_list),
+        ("Model Download", test_model_download_from_training_server),
+        ("Send Training Data", test_add_training_data_to_training_server),
+        ("Model Sync", test_prediction_server_model_sync),
+        ("Predictions", test_prediction_via_prediction_server),
+        ("Bulk Prediction Strict", test_bulk_prediction_strict),
+        ("Bulk Prediction With Errors", test_bulk_prediction_all_valid),
+        ("Bulk predictions all valid", test_bulk_prediction_with_validation_errors),
+        ("Prediction Missing Prefix Cache", test_prediction_missing_prefix_cache_score),
+        ("Training Metrics", test_training_server_metrics),
+        ("Model Consistency", test_model_consistency_between_servers),
+        ("XGBoost Trees", test_model_specific_endpoints_on_training_server),
+        ("Flush API", test_training_server_flush_api),  
+        ("Flush Error Handling", test_training_server_flush_error_handling),
+        
+        ("Dual Server Model Learns Equation", test_dual_server_quantile_regression_learns_distribution),
+        ("End-to-End Workflow", test_end_to_end_workflow),
+        ("Prediction Stress Test", test_prediction_server_stress_test),
+        ("Bulk Prediction Stress Test", test_bulk_prediction_stress_test),
+        ("Large Batch Prediction Stress Test", test_large_batch_prediction_stress_test),
+    ]
+    
+    passed = 0
+    failed = 0
+    
+    for test_name, test_func in tests:
+        try:
+            test_func()
+            print(f"✓ {test_name} passed")
+            passed += 1
+        except Exception as e:
+            print(f"✗ {test_name} failed: {e}")
+            failed += 1
+    
+    print(f"\n{'='*50}")
+    print(f"FINAL RESULTS: {passed} passed, {failed} failed")
+    print(f"{'='*50}")
+    
+    if failed == 0:
+        print("🎉 All tests passed! Your dual-server architecture with prefix cache score is working correctly.")
+    else:
+        print(f"⚠️  {failed} tests failed. Check the issues above.")
\ No newline at end of file
diff --git a/latencypredictor/test_latency_predictor_client.py b/latencypredictor/test_latency_predictor_client.py
new file mode 100644
index 0000000..def7a79
--- /dev/null
+++ b/latencypredictor/test_latency_predictor_client.py
@@ -0,0 +1,1257 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import time
+import asyncio
+import aiohttp
+import threading
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from collections import defaultdict
+import random
+
+import pytest
+import requests
+
+import joblib
+import numpy as np
+import tempfile
+import xgboost
+
+# Base URL of your running FastAPI server
+BASE_URL = os.getenv("TRAINING_SERVER_URL", "http://34.143.221.122:80")
+
+# Helper to wait until the server is ready
+def wait_for_ready(timeout: float = 30.0, interval: float = 1.0):
+    start = time.time()
+    while True:
+        try:
+            r = requests.get(f"{BASE_URL}/readyz", timeout=2.0)
+            if r.status_code == 200:
+                return
+        except requests.RequestException:
+            pass
+        if time.time() - start > timeout:
+            pytest.skip("Server did not become ready in time")
+        time.sleep(interval)
+
+@pytest.fixture(scope="module", autouse=True)
+def ensure_server_ready():
+    """Wait for the /readyz endpoint before running tests."""
+    wait_for_ready()
+
+
+def test_healthz():
+    r = requests.get(f"{BASE_URL}/healthz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ok"
+
+
+def test_readyz():
+    r = requests.get(f"{BASE_URL}/readyz")
+    assert r.status_code == 200
+    assert r.json().get("status") == "ready"
+
+
+def test_model_info():
+    """Test the simplified /model/download/info endpoint."""
+    r = requests.get(f"{BASE_URL}/model/download/info")
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert "model_type" in data
+    assert "model_status" in data
+    assert "available_endpoints" in data
+    assert data["model_type"] in ["bayesian_ridge", "xgboost"]
+    assert isinstance(data["model_status"], dict)
+    
+    print(f"Server using model type: {data['model_type']}")
+    
+    if data["model_type"] == "bayesian_ridge":
+        assert "coefficients_info" in data
+        assert data["available_endpoints"]["coefficients"] == "/metrics"
+    else:  # XGBoost
+        assert "trees" in data["available_endpoints"]
+
+
+def test_root_endpoint_enhanced():
+    """Test the enhanced root endpoint that now includes model info."""
+    r = requests.get(f"{BASE_URL}/")
+    assert r.status_code == 200
+    
+    data = r.json()
+    assert "message" in data
+    assert "model_type" in data
+    assert data["model_type"] in ["bayesian_ridge", "xgboost"]
+
+
+def test_add_training_data_bulk():
+    """
+    Send 120 training samples in one bulk request so the server can retrain:
+    Updated equations with prefix cache score:
+      actual_ttft_ms = 2*input_token_length + 3*num_request_waiting +
+                       4*num_request_running + 50*kv_cache_percentage + 
+                       30*prefix_cache_score + 95
+      actual_tpot_ms = 100*kv_cache_percentage + 0.5*input_token_length + 1*num_tokens_generated +
+                       5*num_request_running + 9
+    """
+    entries = []
+    common = {
+        "kv_cache_percentage": 0.5,
+        "num_request_running": 1,
+    }
+
+    for i in range(1, 121):
+        waiting = i % 10 + 1
+        tokens = waiting
+        inp_len = 10 * i
+        kv = common["kv_cache_percentage"]
+        running = common["num_request_running"]
+        prefix_cache = random.uniform(0.1, 0.9)  # Added prefix cache score
+        
+        entries.append({
+            "kv_cache_percentage": kv,
+            "input_token_length": inp_len,
+            "num_request_waiting": waiting,
+            "num_request_running": running,
+            # Updated TTFT formula to include prefix_cache_score
+            "actual_ttft_ms": (inp_len*2.0 + waiting*3.0 + running*4.0 + kv*50.0 + prefix_cache*30.0) + 95,
+            # TPOT formula remains unchanged
+            "actual_tpot_ms": (kv*100.0 + inp_len*0.5 + tokens*1.0 + running*5.0) + 9,
+            "num_tokens_generated": tokens,
+            "prefix_cache_score": prefix_cache,  # Added prefix cache score
+            "timestamp": time.time()  # FastAPI will coerce to datetime
+        })
+
+    payload = {"entries": entries}
+    r = requests.post(f"{BASE_URL}/add_training_data_bulk", json=payload)
+    assert r.status_code == 202, f"Expected 202, got {r.status_code}"
+    assert r.json().get("message") == "Accepted 120 training samples."
+
+
+def test_model_learns_equation():
+    """
+    After sending bulk data, poll /predict until the model's predictions
+    match our linear equations within tolerance, or fail after 60s.
+    Updated to include prefix_cache_score in the test equation.
+    """
+    # First check what model type we're using
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type", "unknown")
+    
+    features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 200,
+        "num_request_waiting": 4,
+        "num_request_running": 1,
+        "num_tokens_generated": 4,
+        "prefix_cache_score": 0.7,  # Added prefix cache score
+    }
+    
+    # Updated expected TTFT to include prefix cache score
+    expected_ttft = (
+        features["input_token_length"] * 2.0
+        + features["num_request_waiting"] * 3.0
+        + features["num_request_running"] * 4.0
+        + features["kv_cache_percentage"] * 50.0
+        + features["prefix_cache_score"] * 30.0  # New term
+        + 95
+    )
+    # TPOT formula remains unchanged
+    expected_tpot = (
+        features["kv_cache_percentage"] * 100.0
+        + features["input_token_length"] * 0.5
+        + features["num_tokens_generated"] * 1.0
+        + features["num_request_running"] * 5.0 + 9
+    )
+
+    # Adjust tolerance based on model type
+    # XGBoost might need more tolerance for tree-based predictions
+    tolerance = 0.15 if model_type == "xgboost" else 0.1
+    
+    deadline = time.time() + 60.0
+    last_ttft, last_tpot = None, None
+
+    while time.time() < deadline:
+        r = requests.post(f"{BASE_URL}/predict", json=features)
+        if r.status_code != 200:
+            time.sleep(1)
+            continue
+
+        body = r.json()
+        last_ttft = body["ttft_ms"]
+        last_tpot = body["tpot_ms"]
+        
+        # Verify the response includes model_type
+        assert "model_type" in body, "Response should include model_type"
+        assert body["model_type"] == model_type
+
+        ttft_ok = abs(last_ttft - expected_ttft) <= tolerance * expected_ttft
+        tpot_ok = abs(last_tpot - expected_tpot) <= tolerance * expected_tpot
+        if ttft_ok and tpot_ok:
+            print(f"Model converged with {model_type} in {60.0 - (deadline - time.time()):.1f}s")
+            print(f"  Expected TTFT: {expected_ttft:.1f}, Got: {last_ttft:.1f}")
+            print(f"  Expected TPOT: {expected_tpot:.1f}, Got: {last_tpot:.1f}")
+            break
+
+        time.sleep(1)
+
+    assert last_ttft is not None, "Never got a successful prediction."
+    assert abs(last_ttft - expected_ttft) <= tolerance * expected_ttft, (
+        f"TTFT={last_ttft:.1f} not within ±{tolerance*100}% of {expected_ttft:.1f} (model: {model_type})"
+    )
+    assert abs(last_tpot - expected_tpot) <= tolerance * expected_tpot, (
+        f"TPOT={last_tpot:.1f} not within ±{tolerance*100}% of {expected_tpot:.1f} (model: {model_type})"
+    )
+
+
+def test_prediction_missing_prefix_cache_score():
+    """Test that predictions fail when prefix_cache_score is missing."""
+    features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 200,
+        "num_request_waiting": 4,
+        "num_request_running": 1,
+        "num_tokens_generated": 4,
+        # Missing prefix_cache_score
+    }
+    
+    r = requests.post(f"{BASE_URL}/predict", json=features)
+    assert r.status_code == 422  # Should fail validation
+    
+    print("✓ Prediction correctly failed when prefix_cache_score was missing")
+
+
+def test_prefix_cache_score_impact_on_ttft():
+    """
+    Test that prefix_cache_score has the expected impact on TTFT predictions.
+    Since our test equation has +30*prefix_cache_score, higher scores should increase TTFT.
+    """
+    print("Testing prefix cache score impact on TTFT predictions...")
+    
+    base_features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 300,
+        "num_request_waiting": 4,
+        "num_request_running": 2,
+        "num_tokens_generated": 15,
+    }
+    
+    prefix_cache_scores = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
+    predictions = []
+    
+    for prefix_score in prefix_cache_scores:
+        test_features = {**base_features, "prefix_cache_score": prefix_score}
+        
+        pred_r = requests.post(f"{BASE_URL}/predict", json=test_features, timeout=10)
+        assert pred_r.status_code == 200
+        
+        pred_data = pred_r.json()
+        predictions.append({
+            "prefix_cache_score": prefix_score,
+            "ttft_ms": pred_data["ttft_ms"],
+            "tpot_ms": pred_data["tpot_ms"]
+        })
+        
+        print(f"  Prefix cache {prefix_score:.1f}: TTFT={pred_data['ttft_ms']:.1f}ms, TPOT={pred_data['tpot_ms']:.1f}ms")
+    
+    # Check that TTFT increases as prefix cache score increases
+    # (since our test equation has +30*prefix_cache_score)
+    ttft_values = [p["ttft_ms"] for p in predictions]
+    
+    # Calculate correlation between prefix cache score and TTFT
+    first_half_avg = sum(ttft_values[:3]) / 3  # Low prefix cache scores
+    second_half_avg = sum(ttft_values[3:]) / 3  # High prefix cache scores
+    
+    print(f"Low prefix cache avg TTFT: {first_half_avg:.1f}ms")
+    print(f"High prefix cache avg TTFT: {second_half_avg:.1f}ms")
+    
+    # Since our training equation has +30*prefix_cache_score, higher prefix cache should increase TTFT
+    ttft_difference = second_half_avg - first_half_avg
+    print(f"TTFT difference (high - low prefix cache): {ttft_difference:.1f}ms")
+    
+    # Should be positive difference (higher prefix cache = higher TTFT in our test equation)
+    assert ttft_difference > 10, f"Expected TTFT to increase with prefix cache score, got difference: {ttft_difference:.1f}ms"
+    
+    # TPOT should not be significantly affected by prefix cache score
+    tpot_values = [p["tpot_ms"] for p in predictions]
+    tpot_first_half = sum(tpot_values[:3]) / 3
+    tpot_second_half = sum(tpot_values[3:]) / 3
+    tpot_difference = abs(tpot_second_half - tpot_first_half)
+    
+    print(f"TPOT difference (should be small): {tpot_difference:.1f}ms")
+    assert tpot_difference < 5, f"TPOT should not be significantly affected by prefix cache, got difference: {tpot_difference:.1f}ms"
+    
+    print("✓ Prefix cache score impact test passed")
+
+
+def test_prediction_response_format():
+    """Test that prediction responses include all expected fields including new model_type."""
+    features = generate_random_prediction_payload()
+    
+    r = requests.post(f"{BASE_URL}/predict", json=features)
+    assert r.status_code == 200
+    
+    data = r.json()
+    required_fields = [
+        "ttft_ms", "tpot_ms", "ttft_uncertainty", "tpot_uncertainty",
+        "ttft_prediction_bounds", "tpot_prediction_bounds", 
+        "predicted_at", "model_type"
+    ]
+    
+    for field in required_fields:
+        assert field in data, f"Missing required field: {field}"
+    
+    # Verify model_type is valid
+    assert data["model_type"] in ["bayesian_ridge", "xgboost"]
+    
+    # Verify numeric fields are reasonable
+    assert data["ttft_ms"] >= 0
+    assert data["tpot_ms"] >= 0
+    assert data["ttft_uncertainty"] >= 0
+    assert data["tpot_uncertainty"] >= 0
+    
+    # Verify bounds are tuples
+    assert len(data["ttft_prediction_bounds"]) == 2
+    assert len(data["tpot_prediction_bounds"]) == 2
+
+
+def test_metrics_endpoint_enhanced():
+    """Test that metrics endpoint includes model-specific information with proper coefficients."""
+    r = requests.get(f"{BASE_URL}/metrics")
+    assert r.status_code == 200
+    
+    content = r.text
+    
+    # Should contain model type metric
+    assert "model_type{" in content
+    
+    # Should contain either coefficients (Bayesian Ridge) or importance (XGBoost)
+    has_coef = "ttft_coef{" in content or "tpot_coef{" in content
+    has_importance = "ttft_importance{" in content or "tpot_importance{" in content
+    
+    assert has_coef or has_importance, "Should have either coefficients or feature importance metrics"
+    
+    # Should have standard metrics
+    assert "ttft_r2_score{" in content
+    assert "tpot_r2_score{" in content
+    assert "training_samples_count" in content
+    
+    # Check for prefix_cache_score in TTFT metrics
+    if has_coef:
+        assert 'feature="prefix_cache_score"' in content, "Should have prefix_cache_score coefficient for TTFT model"
+    if has_importance:
+        assert 'feature="prefix_cache_score"' in content, "Should have prefix_cache_score importance for TTFT model"
+    
+    # Parse and validate coefficient values for Bayesian Ridge
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    if model_type == "bayesian_ridge":
+        # Check that coefficients are present and reasonable
+        lines = content.split('\n')
+        ttft_intercept = None
+        ttft_coefs = {}
+        tpot_intercept = None
+        tpot_coefs = {}
+        
+        for line in lines:
+            if line.startswith('ttft_intercept{'):
+                ttft_intercept = float(line.split('}')[1].strip())
+            elif line.startswith('ttft_coef{'):
+                feature = line.split('feature="')[1].split('"')[0]
+                value = float(line.split('}')[1].strip())
+                ttft_coefs[feature] = value
+            elif line.startswith('tpot_intercept{'):
+                tpot_intercept = float(line.split('}')[1].strip())
+            elif line.startswith('tpot_coef{'):
+                feature = line.split('feature="')[1].split('"')[0]
+                value = float(line.split('}')[1].strip())
+                tpot_coefs[feature] = value
+        
+        # Validate coefficients are present
+        assert ttft_intercept is not None, "TTFT intercept should be present"
+        assert tpot_intercept is not None, "TPOT intercept should be present"
+        
+        # Updated expected features to include prefix_cache_score for TTFT
+        expected_ttft_features = ["kv_cache_percentage", "input_token_length", "num_request_waiting", "num_request_running", "prefix_cache_score"]
+        expected_tpot_features = ["kv_cache_percentage", "input_token_length", "num_request_waiting", "num_request_running", "num_tokens_generated"]
+        
+        for feature in expected_ttft_features:
+            assert feature in ttft_coefs, f"TTFT coefficient for {feature} should be present"
+            
+        for feature in expected_tpot_features:
+            assert feature in tpot_coefs, f"TPOT coefficient for {feature} should be present"
+        
+        print(f"✓ Bayesian Ridge coefficients validated:")
+        print(f"  TTFT intercept: {ttft_intercept:.4f}")
+        print(f"  TTFT coefficients: {ttft_coefs}")
+        print(f"  TPOT intercept: {tpot_intercept:.4f}")
+        print(f"  TPOT coefficients: {tpot_coefs}")
+        
+        # Validate prefix_cache_score coefficient is reasonable
+        if "prefix_cache_score" in ttft_coefs:
+            prefix_coef = ttft_coefs["prefix_cache_score"]
+            print(f"  Prefix cache coefficient: {prefix_coef:.4f}")
+            # Should be positive and reasonably close to our training value of 30
+            assert 10 < prefix_coef < 50, f"Prefix cache coefficient should be reasonable: {prefix_coef}"
+
+    print("✓ Training server metrics endpoint working correctly with prefix cache support")
+
+
+def test_xgboost_tree_endpoints():
+    """Test XGBoost tree endpoints if XGBoost is being used."""
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    if model_type != "xgboost":
+        print("Skipping XGBoost tree tests - not using XGBoost model")
+        return
+    
+    print("Testing XGBoost tree endpoints...")
+    
+    # Test TTFT trees
+    ttft_response = requests.get(f"{BASE_URL}/model/ttft/xgb/json")
+    assert ttft_response.status_code == 200, "TTFT XGBoost trees should be available"
+    ttft_trees = ttft_response.json()
+    assert isinstance(ttft_trees, list), "TTFT trees should be a list"
+    assert len(ttft_trees) > 0, "Should have TTFT trees"
+    assert isinstance(ttft_trees[0], dict), "Each tree should be a dict"
+    
+    # Test TPOT trees
+    tpot_response = requests.get(f"{BASE_URL}/model/tpot/xgb/json")
+    assert tpot_response.status_code == 200, "TPOT XGBoost trees should be available"
+    tpot_trees = tpot_response.json()
+    assert isinstance(tpot_trees, list), "TPOT trees should be a list"
+    assert len(tpot_trees) > 0, "Should have TPOT trees"
+    assert isinstance(tpot_trees[0], dict), "Each tree should be a dict"
+    
+    print(f"✓ XGBoost trees available: {len(ttft_trees)} TTFT trees, {len(tpot_trees)} TPOT trees")
+
+
+def test_bayesian_ridge_coefficients():
+    """Test that Bayesian Ridge coefficients are properly descaled and stored."""
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    if model_type != "bayesian_ridge":
+        print("Skipping Bayesian Ridge coefficient tests - not using Bayesian Ridge model")
+        return
+    
+    print("Testing Bayesian Ridge coefficient storage and retrieval...")
+    
+    # Get coefficients from metrics
+    r = requests.get(f"{BASE_URL}/metrics")
+    assert r.status_code == 200
+    content = r.text
+    
+    # Parse coefficients from metrics
+    lines = content.split('\n')
+    ttft_coefs = {}
+    tpot_coefs = {}
+    
+    for line in lines:
+        if line.startswith('ttft_coef{'):
+            feature = line.split('feature="')[1].split('"')[0]
+            value = float(line.split('}')[1].strip())
+            ttft_coefs[feature] = value
+        elif line.startswith('tpot_coef{'):
+            feature = line.split('feature="')[1].split('"')[0]
+            value = float(line.split('}')[1].strip())
+            tpot_coefs[feature] = value
+    
+    # Test a prediction to see if coefficients make sense
+    test_features = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 100,
+        "num_request_waiting": 2,
+        "num_request_running": 1,
+        "num_tokens_generated": 5,
+        "prefix_cache_score": 0.8,  # Added prefix cache score
+    }
+    
+    # Make prediction via API
+    pred_response = requests.post(f"{BASE_URL}/predict", json=test_features)
+    assert pred_response.status_code == 200
+    api_prediction = pred_response.json()
+    
+    print(f"✓ Coefficients extracted from metrics:")
+    print(f"  TTFT coefficients: {ttft_coefs}")
+    print(f"  TPOT coefficients: {tpot_coefs}")
+    print(f"  API TTFT prediction: {api_prediction['ttft_ms']:.2f}")
+    print(f"  API TPOT prediction: {api_prediction['tpot_ms']:.2f}")
+    
+    # Verify prefix_cache_score coefficient exists for TTFT
+    assert "prefix_cache_score" in ttft_coefs, "prefix_cache_score should be in TTFT coefficients"
+    assert "prefix_cache_score" not in tpot_coefs, "prefix_cache_score should NOT be in TPOT coefficients"
+
+
+def test_model_endpoints_by_type():
+    """Test the appropriate endpoints based on model type."""
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_info = model_info_r.json()
+    model_type = model_info["model_type"]
+    
+    print(f"Testing endpoints for model type: {model_type}")
+    
+    if model_type == "bayesian_ridge":
+        # For Bayesian Ridge, we should have coefficients in metrics
+        test_bayesian_ridge_coefficients()
+        
+        # XGBoost endpoints should return 404
+        ttft_xgb_response = requests.get(f"{BASE_URL}/model/ttft/xgb/json")
+        assert ttft_xgb_response.status_code == 404, "XGBoost endpoints should not be available for Bayesian Ridge"
+        
+        print("✓ Bayesian Ridge: coefficients available in metrics, XGBoost endpoints properly blocked")
+        
+    else:  # XGBoost
+        # For XGBoost, we should have tree endpoints
+        test_xgboost_tree_endpoints()
+        
+        print("✓ XGBoost: tree endpoints available")
+
+
+def generate_random_prediction_payload():
+    """Generate a random prediction payload for stress testing including prefix_cache_score."""
+    return {
+        "kv_cache_percentage": random.uniform(0.1, 0.9),
+        "input_token_length": random.randint(10, 1000),
+        "num_request_waiting": random.randint(1, 20),
+        "num_request_running": random.randint(1, 10),
+        "num_tokens_generated": random.randint(1, 20),
+        "prefix_cache_score": random.uniform(0.0, 1.0),  # Added prefix cache score
+    }
+
+
+def generate_random_training_payload():
+    """Generate a random training data payload for stress testing with updated TTFT formula."""
+    input_tokens = random.randint(10, 1000)
+    waiting_requests = random.randint(1, 20)
+    running_requests = random.randint(1, 10)
+    kv = random.uniform(0.01, 0.99)
+    tokens_generated = random.randint(1, 20)
+    prefix_cache = random.uniform(0.0, 1.0)  # Added prefix cache score
+    
+    return {
+        "kv_cache_percentage": kv,
+        "input_token_length": input_tokens,
+        "num_request_waiting": waiting_requests,
+        "num_request_running": running_requests,
+        # Updated linear TTFT with noise - now includes prefix_cache_score
+        "actual_ttft_ms": (
+            input_tokens * 2.0
+            + waiting_requests * 3.0
+            + running_requests * 4.0
+            + kv * 50.0
+            + prefix_cache * 30.0  # New term for prefix cache
+            + 95 + random.uniform(-10, 10)
+        ),
+        # TPOT formula remains unchanged
+        "actual_tpot_ms": (
+            kv * 100.0
+            + input_tokens * 0.5
+            + tokens_generated * 1.0
+            + running_requests * 5.0
+            + 9 + random.uniform(-5, 5)
+        ),
+        "num_tokens_generated": tokens_generated,
+        "prefix_cache_score": prefix_cache,  # Added prefix cache score
+    }
+
+
+def generate_bulk_training_payload(size=1000):
+    """Generate a bulk training payload with specified number of entries."""
+    entries = []
+    for _ in range(size):
+        entries.append(generate_random_training_payload())
+    return {"entries": entries}
+
+
+async def async_post_request(session, url, payload, request_id):
+    """Make an async POST request and return result with metadata."""
+    start_time = time.time()
+    try:
+        async with session.post(url, json=payload, timeout=aiohttp.ClientTimeout(total=5)) as response:
+            end_time = time.time()
+            response_data = await response.json()
+            return {
+                'request_id': request_id,
+                'status_code': response.status,
+                'response_time': end_time - start_time,
+                'success': response.status in [200, 202],
+                'response_data': response_data,
+                'request_type': 'predict' if '/predict' in url else 'training',
+                'model_type': response_data.get('model_type') if response.status == 200 else None
+            }
+    except Exception as e:
+        end_time = time.time()
+        return {
+            'request_id': request_id,
+            'status_code': 0,
+            'response_time': end_time - start_time,
+            'success': False,
+            'error': str(e),
+            'request_type': 'predict' if '/predict' in url else 'training',
+            'model_type': None
+        }
+
+async def run_stress_test_async(duration_seconds=10, target_qps=300):
+    interval = 1.0/target_qps
+    start = time.time()
+    connector = aiohttp.TCPConnector(limit=10000, limit_per_host=10000, ttl_dns_cache=300, use_dns_cache=True)
+    async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=2)) as sess:
+        tasks = []
+        req_id = 0
+        next_time = start
+        while time.time() - start < duration_seconds:
+            now = time.time()
+            while next_time <= now:
+                req_id += 1
+                if random.random()<0.5:
+                    url = f"{BASE_URL}/predict"
+                    payload = generate_random_prediction_payload()
+                else:
+                    url = f"{BASE_URL}/add_training_data_bulk"
+                    payload = {"entries":[ generate_random_training_payload() ]}
+                tasks.append(asyncio.create_task(async_post_request(sess, url, payload, req_id)))
+                next_time += interval
+            await asyncio.sleep(0.0001)
+
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        
+        valid_results = [r for r in results if isinstance(r, dict)]
+    
+        # Calculate actual QPS achieved
+        if valid_results:
+            actual_duration = duration_seconds
+            actual_qps = len(valid_results) / actual_duration
+            print(f"Target QPS: {target_qps}, Actual QPS: {actual_qps:.0f}")
+    
+        return valid_results
+
+
+def fetch_and_parse_xgb_json(path_suffix):
+    """
+    Download the XGBoost JSON dump for `path_suffix` (ttft or tpot),
+    parse into a Python list of dicts, and return it.
+    """
+    url = f"{BASE_URL}/model/{path_suffix}/xgb/json"
+    r = requests.get(url, timeout=10)
+    assert r.status_code == 200, f"Failed to fetch JSON for {path_suffix}"
+    trees = r.json()
+    assert isinstance(trees, list), "Expected a JSON array of trees"
+    assert len(trees) > 0, "Tree list should not be empty"
+    assert isinstance(trees[0], dict), "Each tree must be a JSON object"
+    return trees
+
+
+async def async_fetch_and_parse_xgb_json(session, suffix, request_id):
+    """
+    Async GET /model/<suffix>/xgb/json and return timing + status.
+    """
+    url = f"{BASE_URL}/model/{suffix}/xgb/json"
+    start = time.time()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=10)) as resp:
+            data = await resp.json()
+            elapsed = time.time() - start
+            return {
+                'request_id': request_id,
+                'request_type': f'download_{suffix}',
+                'status_code': resp.status,
+                'response_time': elapsed,
+                'success': resp.status == 200,
+                'tree_count': len(data) if isinstance(data, list) else None
+            }
+    except Exception as e:
+        elapsed = time.time() - start
+        return {
+            'request_id': request_id,
+            'request_type': f'download_{suffix}',
+            'status_code': 0,
+            'response_time': elapsed,
+            'success': False,
+            'error': str(e)
+        }
+
+
+async def run_simplified_stress_test(duration_seconds=10, target_qps=2):
+    """
+    Simplified stress test: bulk training vs predictions and tree downloads (XGBoost only).
+    """
+    info_r = requests.get(f"{BASE_URL}/model/download/info", timeout=5.0)
+    model_type = info_r.json().get("model_type", "bayesian_ridge")
+
+    interval = 1.0 / target_qps
+    start = time.time()
+    connector = aiohttp.TCPConnector(limit=1000, limit_per_host=1000)
+    async with aiohttp.ClientSession(connector=connector) as sess:
+        tasks = []
+        req_id = 0
+        next_time = start
+
+        while time.time() - start < duration_seconds:
+            now = time.time()
+            while next_time <= now:
+                req_id += 1
+
+                if random.random() < 0.5:
+                    # Either predictions or tree downloads (XGBoost only)
+                    if random.random() < 0.7:  # 70% predictions
+                        url = f"{BASE_URL}/predict"
+                        payload = generate_random_prediction_payload()
+                        task = asyncio.create_task(
+                            async_post_request_with_timeout(
+                                sess, url, payload, req_id,
+                                aiohttp.ClientTimeout(total=5), "predict"
+                            )
+                        )
+                    else:  # 30% tree downloads (only for XGBoost)
+                        if model_type == "xgboost":
+                            suffix = random.choice(["ttft", "tpot"])
+                            task = asyncio.create_task(
+                                async_fetch_and_parse_xgb_json(sess, suffix, req_id)
+                            )
+                        else:
+                            # For Bayesian Ridge, just do another prediction
+                            url = f"{BASE_URL}/predict"
+                            payload = generate_random_prediction_payload()
+                            task = asyncio.create_task(
+                                async_post_request_with_timeout(
+                                    sess, url, payload, req_id,
+                                    aiohttp.ClientTimeout(total=5), "predict"
+                                )
+                            )
+                else:
+                    # bulk training
+                    url = f"{BASE_URL}/add_training_data_bulk"
+                    payload = generate_bulk_training_payload(1000)
+                    task = asyncio.create_task(
+                        async_post_request_with_timeout(
+                            sess, url, payload, req_id,
+                            aiohttp.ClientTimeout(total=30), "bulk_training"
+                        )
+                    )
+
+                tasks.append(task)
+                next_time += interval
+
+            await asyncio.sleep(0.001)
+
+        print(f"Waiting for {len(tasks)} requests to complete…")
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        valid = [r for r in results if isinstance(r, dict)]
+
+        if valid:
+            actual_qps = len(valid) / duration_seconds
+            print(f"Target QPS: {target_qps}, Actual QPS: {actual_qps:.2f}")
+
+        return valid
+
+
+async def async_post_request_with_timeout(session, url, payload, request_id, timeout, request_type):
+    """Make an async POST request with custom timeout and return result with metadata."""
+    start_time = time.time()
+    try:
+        async with session.post(url, json=payload, timeout=timeout) as response:
+            end_time = time.time()
+            response_data = await response.json()
+            
+            # Count training entries for bulk requests
+            training_entries = len(payload.get("entries", [])) if request_type == "bulk_training" else 1
+            
+            return {
+                'request_id': request_id,
+                'status_code': response.status,
+                'response_time': end_time - start_time,
+                'success': response.status in [200, 202],
+                'response_data': response_data,
+                'request_type': request_type,
+                'training_entries': training_entries if request_type == "bulk_training" else 0,
+                'model_type': response_data.get('model_type') if response.status == 200 and request_type == 'predict' else None
+            }
+    except Exception as e:
+        end_time = time.time()
+        training_entries = len(payload.get("entries", [])) if request_type == "bulk_training" else 1
+        return {
+            'request_id': request_id,
+            'status_code': 0,
+            'response_time': end_time - start_time,
+            'success': False,
+            'error': str(e),
+            'request_type': request_type,
+            'training_entries': training_entries if request_type == "bulk_training" else 0,
+            'model_type': None
+        }
+
+
+def analyze_stress_test_results(results):
+    """Analyze and print stress test results with model type information."""
+    if not results:
+        print("No results to analyze")
+        return
+    
+    total_requests = len(results)
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    failed_requests = total_requests - successful_requests
+    
+    response_times = [r['response_time'] for r in results if r.get('response_time')]
+    avg_response_time = sum(response_times) / len(response_times) if response_times else 0
+    
+    status_codes = defaultdict(int)
+    for r in results:
+        status_codes[r.get('status_code', 0)] += 1
+    
+    request_types = defaultdict(int)
+    for r in results:
+        request_types[r.get('request_type', 'unknown')] += 1
+    
+    # Analyze model types in prediction responses
+    model_types = defaultdict(int)
+    for r in results:
+        if r.get('model_type'):
+            model_types[r['model_type']] += 1
+    
+    test_duration = max(response_times) if response_times else 0
+    actual_qps = total_requests / test_duration if test_duration > 0 else 0
+    
+    print(f"\n{'='*50}")
+    print("STRESS TEST RESULTS")
+    print(f"{'='*50}")
+    print(f"Total Requests: {total_requests}")
+    print(f"Successful: {successful_requests} ({successful_requests/total_requests*100:.1f}%)")
+    print(f"Failed: {failed_requests} ({failed_requests/total_requests*100:.1f}%)")
+    print(f"Average Response Time: {avg_response_time*1000:.2f}ms")
+    print(f"Actual QPS: {actual_qps:.0f}")
+    print(f"\nRequest Types:")
+    for req_type, count in request_types.items():
+        print(f"  {req_type}: {count}")
+    print(f"\nStatus Code Distribution:")
+    for status, count in status_codes.items():
+        print(f"  {status}: {count}")
+    
+    if model_types:
+        print(f"\nModel Types in Predictions:")
+        for model_type, count in model_types.items():
+            print(f"  {model_type}: {count}")
+    
+    if response_times:
+        sorted_times = sorted(response_times)
+        p50 = sorted_times[int(len(sorted_times) * 0.5)] * 1000
+        p95 = sorted_times[int(len(sorted_times) * 0.95)] * 1000
+        p99 = sorted_times[int(len(sorted_times) * 0.99)] * 1000
+        print(f"\nResponse Time Percentiles:")
+        print(f"  P50: {p50:.2f}ms")
+        print(f"  P95: {p95:.2f}ms")
+        print(f"  P99: {p99:.2f}ms")
+
+
+def analyze_bulk_training_results(results):
+    """Analyze and print bulk training stress test results with additional metrics."""
+    if not results:
+        print("No results to analyze")
+        return
+    
+    total_requests = len(results)
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    failed_requests = total_requests - successful_requests
+    
+    # Separate analysis by request type
+    prediction_results = [r for r in results if r.get('request_type') == 'predict']
+    bulk_training_results = [r for r in results if r.get('request_type') == 'bulk_training']
+    download_results = [r for r in results if r.get('request_type', '').startswith('download_')]
+    
+    # Calculate total training entries processed
+    total_training_entries = sum(r.get('training_entries', 0) for r in bulk_training_results)
+    
+    # Analyze model types in prediction responses
+    model_types = defaultdict(int)
+    for r in prediction_results:
+        if r.get('model_type'):
+            model_types[r['model_type']] += 1
+    
+    response_times = [r['response_time'] for r in results if r.get('response_time')]
+    avg_response_time = sum(response_times) / len(response_times) if response_times else 0
+    
+    status_codes = defaultdict(int)
+    for r in results:
+        status_codes[r.get('status_code', 0)] += 1
+    
+    request_types = defaultdict(int)
+    for r in results:
+        request_types[r.get('request_type', 'unknown')] += 1
+    
+    print(f"\n{'='*60}")
+    print("BULK TRAINING STRESS TEST RESULTS")
+    print(f"{'='*60}")
+    print(f"Total Requests: {total_requests}")
+    print(f"Successful: {successful_requests} ({successful_requests/total_requests*100:.1f}%)")
+    print(f"Failed: {failed_requests} ({failed_requests/total_requests*100:.1f}%)")
+    print(f"Average Response Time: {avg_response_time*1000:.2f}ms")
+    
+    print(f"\nRequest Type Breakdown:")
+    print(f"  Prediction requests: {len(prediction_results)}")
+    print(f"  Bulk training requests: {len(bulk_training_results)}")
+    print(f"  Model download requests: {len(download_results)}")
+    print(f"  Total training entries processed: {total_training_entries}")
+    
+    if model_types:
+        print(f"\nModel Types in Predictions:")
+        for model_type, count in model_types.items():
+            print(f"  {model_type}: {count}")
+    
+    print(f"\nStatus Code Distribution:")
+    for status, count in status_codes.items():
+        print(f"  {status}: {count}")
+    
+    # Response time analysis by request type
+    if prediction_results:
+        pred_times = [r['response_time'] for r in prediction_results if r.get('response_time')]
+        if pred_times:
+            avg_pred_time = sum(pred_times) / len(pred_times)
+            print(f"\nPrediction Request Response Times:")
+            print(f"  Average: {avg_pred_time*1000:.2f}ms")
+            print(f"  Min: {min(pred_times)*1000:.2f}ms")
+            print(f"  Max: {max(pred_times)*1000:.2f}ms")
+    
+    if bulk_training_results:
+        bulk_times = [r['response_time'] for r in bulk_training_results if r.get('response_time')]
+        if bulk_times:
+            avg_bulk_time = sum(bulk_times) / len(bulk_times)
+            print(f"\nBulk Training Request Response Times:")
+            print(f"  Average: {avg_bulk_time*1000:.2f}ms")
+            print(f"  Min: {min(bulk_times)*1000:.2f}ms")
+            print(f"  Max: {max(bulk_times)*1000:.2f}ms")
+    
+    if download_results:
+        download_times = [r['response_time'] for r in download_results if r.get('response_time')]
+        if download_times:
+            avg_download_time = sum(download_times) / len(download_times)
+            print(f"\nModel Download Request Response Times:")
+            print(f"  Average: {avg_download_time*1000:.2f}ms")
+            print(f"  Min: {min(download_times)*1000:.2f}ms")
+            print(f"  Max: {max(download_times)*1000:.2f}ms")
+    
+    if response_times:
+        sorted_times = sorted(response_times)
+        p50 = sorted_times[int(len(sorted_times) * 0.5)] * 1000
+        p95 = sorted_times[int(len(sorted_times) * 0.95)] * 1000
+        p99 = sorted_times[int(len(sorted_times) * 0.99)] * 1000
+        print(f"\nOverall Response Time Percentiles:")
+        print(f"  P50: {p50:.2f}ms")
+        print(f"  P95: {p95:.2f}ms")
+        print(f"  P99: {p99:.2f}ms")
+
+
+def test_stress_test_high_qps():
+    """
+    Stress test with 300 QPS for 10 seconds.
+    Sends predictions and training data in parallel.
+    """
+    results = asyncio.run(run_stress_test_async(duration_seconds=10, target_qps=300))
+    
+    analyze_stress_test_results(results)
+    
+    assert len(results) > 0, "No requests were made"
+    
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    success_rate = successful_requests / len(results)
+    
+    assert success_rate > 0.8, f"Success rate too low: {success_rate*100:.1f}%"
+    
+    print(f"Stress test completed successfully with {success_rate*100:.1f}% success rate")
+
+
+def test_stress_test_mixed_load():
+    """
+    Alternative stress test with mixed load patterns.
+    Tests server stability under varying load conditions.
+    """
+    print("Running mixed load stress test...")
+    
+    print("Phase 1: Ramping up load...")
+    results_phase1 = asyncio.run(run_stress_test_async(duration_seconds=5, target_qps=100))
+    
+    print("Phase 2: High sustained load...")
+    results_phase2 = asyncio.run(run_stress_test_async(duration_seconds=10, target_qps=300))
+    
+    print("Phase 3: Cooling down...")
+    results_phase3 = asyncio.run(run_stress_test_async(duration_seconds=5, target_qps=50))
+    
+    all_results = results_phase1 + results_phase2 + results_phase3
+    
+    print("\nCOMBINED RESULTS FOR ALL PHASES:")
+    analyze_stress_test_results(all_results)
+    
+    assert len(all_results) > 0, "No requests were made"
+    
+    successful_requests = sum(1 for r in all_results if r.get('success', False))
+    success_rate = successful_requests / len(all_results)
+    
+    assert success_rate > 0.75, f"Overall success rate too low: {success_rate*100:.1f}%"
+    
+    print(f"Mixed load stress test completed with {success_rate*100:.1f}% success rate")
+
+
+def test_simplified_stress_test():
+    """Simplified stress test focusing on predictions, training, and tree downloads with prefix cache."""
+    print("Running simplified stress test with prefix cache score support...")
+    print("Configuration: 2 QPS, 50% bulk training, 35% predictions, 15% tree downloads (XGBoost only)")
+    
+    results = asyncio.run(run_simplified_stress_test(duration_seconds=60, target_qps=2))
+    
+    analyze_bulk_training_results(results)
+    
+    assert len(results) > 0, "No requests were made"
+    
+    successful_requests = sum(1 for r in results if r.get('success', False))
+    success_rate = successful_requests / len(results)
+    
+    # Count request types
+    prediction_count = sum(1 for r in results if r.get('request_type') == 'predict')
+    bulk_training_count = sum(1 for r in results if r.get('request_type') == 'bulk_training')
+    download_count = sum(1 for r in results if r.get('request_type', '').startswith('download_'))
+    
+    assert success_rate > 0.8, f"Success rate too low: {success_rate*100:.1f}%"
+    assert prediction_count > 0, "No prediction requests were made"
+    assert bulk_training_count > 0, "No bulk training requests were made"
+    
+    print(f"✓ Simplified stress test with prefix cache completed:")
+    print(f"  Success rate: {success_rate*100:.1f}%")
+    print(f"  Prediction requests: {prediction_count}")
+    print(f"  Tree download requests: {download_count}")
+    print(f"  Bulk training requests: {bulk_training_count}")
+
+
+def test_model_type_consistency():
+    """
+    Test that the model type is consistent across all API endpoints.
+    """
+    print("Testing model type consistency across endpoints...")
+    
+    # Get model type from different endpoints
+    root_response = requests.get(f"{BASE_URL}/")
+    model_info_response = requests.get(f"{BASE_URL}/model/download/info")
+    
+    # Make a prediction to get model type from prediction response
+    prediction_request = generate_random_prediction_payload()
+    prediction_response = requests.post(f"{BASE_URL}/predict", json=prediction_request)
+    
+    # Extract model types
+    root_model_type = root_response.json().get("model_type")
+    model_info_model_type = model_info_response.json().get("model_type")
+    prediction_model_type = prediction_response.json().get("model_type")
+    
+    # Check consistency
+    assert root_model_type == model_info_model_type == prediction_model_type, (
+        f"Model type inconsistency: root={root_model_type}, "
+        f"model_info={model_info_model_type}, prediction={prediction_model_type}"
+    )
+    
+    print(f"Model type consistent across all endpoints: {root_model_type}")
+
+
+def test_xgboost_vs_bayesian_ridge_performance():
+    """
+    Performance comparison test (if both models are available).
+    This test will check model performance differences.
+    """
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_info = model_info_r.json()
+    
+    print(f"Current model: {model_info['model_type']}")
+    
+    # Generate test predictions with prefix cache scores
+    test_cases = [generate_random_prediction_payload() for _ in range(10)]
+    
+    predictions = []
+    response_times = []
+    
+    for test_case in test_cases:
+        start_time = time.time()
+        response = requests.post(f"{BASE_URL}/predict", json=test_case)
+        end_time = time.time()
+        
+        assert response.status_code == 200
+        predictions.append(response.json())
+        response_times.append((end_time - start_time) * 1000)  # Convert to ms
+    
+    avg_response_time = sum(response_times) / len(response_times)
+    avg_prefix_cache = sum(tc['prefix_cache_score'] for tc in test_cases) / len(test_cases)
+    
+    print(f"Model: {predictions[0]['model_type']}")
+    print(f"Average response time: {avg_response_time:.2f}ms")
+    print(f"Average prefix cache score: {avg_prefix_cache:.2f}")
+    print(f"Average TTFT prediction: {sum(p['ttft_ms'] for p in predictions)/len(predictions):.2f}ms")
+    print(f"Average TPOT prediction: {sum(p['tpot_ms'] for p in predictions)/len(predictions):.2f}ms")
+    print(f"Average TTFT uncertainty: {sum(p['ttft_uncertainty'] for p in predictions)/len(predictions):.2f}")
+    print(f"Average TPOT uncertainty: {sum(p['tpot_uncertainty'] for p in predictions)/len(predictions):.2f}")
+    
+    # Basic sanity checks
+    assert avg_response_time < 1000, f"Response time too slow: {avg_response_time:.2f}ms"
+    assert all(p['ttft_ms'] > 0 for p in predictions), "All TTFT predictions should be positive"
+    assert all(p['tpot_ms'] > 0 for p in predictions), "All TPOT predictions should be positive"
+
+
+def test_uncertainty_estimation_quality():
+    """
+    Test the quality of uncertainty estimation for both model types.
+    """
+    model_info_r = requests.get(f"{BASE_URL}/model/download/info")
+    model_type = model_info_r.json().get("model_type")
+    
+    # Generate multiple predictions for the same input
+    test_payload = {
+        "kv_cache_percentage": 0.5,
+        "input_token_length": 100,
+        "num_request_waiting": 2,
+        "num_request_running": 1,
+        "num_tokens_generated": 5,
+        "prefix_cache_score": 0.8,  # Added prefix cache score
+    }
+    
+    predictions = []
+    for _ in range(5):  # Make multiple identical requests
+        response = requests.post(f"{BASE_URL}/predict", json=test_payload)
+        assert response.status_code == 200
+        predictions.append(response.json())
+    
+    # Check that predictions are consistent (should be identical for same input)
+    ttft_values = [p['ttft_ms'] for p in predictions]
+    tpot_values = [p['tpot_ms'] for p in predictions]
+    
+    ttft_std = sum((x - ttft_values[0])**2 for x in ttft_values)**0.5 / len(ttft_values)
+    tpot_std = sum((x - tpot_values[0])**2 for x in tpot_values)**0.5 / len(tpot_values)
+    
+    # For deterministic models, predictions should be identical
+    if model_type == "bayesian_ridge":
+        assert ttft_std < 0.01, f"TTFT predictions should be consistent, got std: {ttft_std}"
+        assert tpot_std < 0.01, f"TPOT predictions should be consistent, got std: {tpot_std}"
+    
+    # Check uncertainty values are reasonable
+    pred = predictions[0]
+    ttft_uncertainty_ratio = pred['ttft_uncertainty'] / pred['ttft_ms']
+    tpot_uncertainty_ratio = pred['tpot_uncertainty'] / pred['tpot_ms']
+    
+    print(f"Model: {model_type}")
+    print(f"Prefix cache score: {test_payload['prefix_cache_score']}")
+    print(f"TTFT: {pred['ttft_ms']:.2f} ± {pred['ttft_uncertainty']:.2f} ({ttft_uncertainty_ratio*100:.1f}%)")
+    print(f"TPOT: {pred['tpot_ms']:.2f} ± {pred['tpot_uncertainty']:.2f} ({tpot_uncertainty_ratio*100:.1f}%)")
+    
+    # Uncertainty should be reasonable (not too high or too low)
+    assert 0.01 < ttft_uncertainty_ratio < 0.5, f"TTFT uncertainty ratio should be reasonable: {ttft_uncertainty_ratio}"
+    assert 0.01 < tpot_uncertainty_ratio < 0.5, f"TPOT uncertainty ratio should be reasonable: {tpot_uncertainty_ratio}"
+    
+    # Check prediction bounds contain the prediction
+    ttft_bounds = pred['ttft_prediction_bounds']
+    tpot_bounds = pred['tpot_prediction_bounds']
+    
+    assert ttft_bounds[0] <= pred['ttft_ms'] <= ttft_bounds[1], "TTFT should be within prediction bounds"
+    assert tpot_bounds[0] <= pred['tpot_ms'] <= tpot_bounds[1], "TPOT should be within prediction bounds"
+
+
+def test_edge_cases():
+    """
+    Test edge cases and boundary conditions with prefix cache score.
+    """
+    # Test minimum values
+    min_payload = {
+        "kv_cache_percentage": 0.0,
+        "input_token_length": 1,
+        "num_request_waiting": 0,
+        "num_request_running": 0,
+        "num_tokens_generated": 1,
+        "prefix_cache_score": 0.0,  # Added prefix cache score
+    }
+    
+    response = requests.post(f"{BASE_URL}/predict", json=min_payload)
+    assert response.status_code == 200
+    data = response.json()
+    assert data['ttft_ms'] > 0
+    assert data['tpot_ms'] > 0
+    
+    # Test maximum reasonable values
+    max_payload = {
+        "kv_cache_percentage": 1.0,
+        "input_token_length": 10000,
+        "num_request_waiting": 100,
+        "num_request_running": 50,
+        "num_tokens_generated": 1000,
+        "prefix_cache_score": 1.0,  # Added prefix cache score
+    }
+    
+    response = requests.post(f"{BASE_URL}/predict", json=max_payload)
+    assert response.status_code == 200
+    data = response.json()
+    assert data['ttft_ms'] > 0
+    assert data['tpot_ms'] > 0
+    
+    # Test invalid values (should fail validation)
+    invalid_payloads = [
+        {"kv_cache_percentage": -0.1, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 1.1, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 0.5, "input_token_length": -1, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 0.5, "input_token_length": 100, "num_request_waiting": -1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 0.5, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": -1, "num_tokens_generated": 10, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 0.5, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": -1, "prefix_cache_score": 0.5},
+        {"kv_cache_percentage": 0.5, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": -0.1},  # Invalid prefix cache
+        {"kv_cache_percentage": 0.5, "input_token_length": 100, "num_request_waiting": 1, "num_request_running": 1, "num_tokens_generated": 10, "prefix_cache_score": 1.1},   # Invalid prefix cache
+    ]
+    
+    for invalid_payload in invalid_payloads:
+        response = requests.post(f"{BASE_URL}/predict", json=invalid_payload)
+        assert response.status_code == 422, f"Should reject invalid payload: {invalid_payload}"
+
+
+def test_concurrent_training_and_prediction():
+    """
+    Test that training and prediction can happen concurrently without issues.
+    """
+    print("Testing concurrent training and prediction with prefix cache...")
+    
+    def make_predictions():
+        results = []
+        for _ in range(20):
+            payload = generate_random_prediction_payload()
+            try:
+                response = requests.post(f"{BASE_URL}/predict", json=payload, timeout=5)
+                results.append(response.status_code == 200)
+            except:
+                results.append(False)
+            time.sleep(0.1)
+        return results
+    
+    def send_training_data():
+        results = []
+        for _ in range(5):
+            payload = generate_bulk_training_payload(100)  # Smaller batches for faster processing
+            try:
+                response = requests.post(f"{BASE_URL}/add_training_data_bulk", json=payload, timeout=10)
+                results.append(response.status_code == 202)
+            except:
+                results.append(False)
+            time.sleep(0.5)
+        return results
+    
+    # Run both functions concurrently
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        prediction_future = executor.submit(make_predictions)
+        training_future = executor.submit(send_training_data)
+        
+        prediction_results = prediction_future.result()
+        training_results = training_future.result()
+    
+    prediction_success_rate = sum(prediction_results) / len(prediction_results)
+    training_success_rate = sum(training_results) / len(training_results)
+    
+    print(f"Prediction success rate: {prediction_success_rate*100:.1f}%")
\ No newline at end of file
diff --git a/latencypredictor/training_server.py b/latencypredictor/training_server.py
new file mode 100644
index 0000000..94a656c
--- /dev/null
+++ b/latencypredictor/training_server.py
@@ -0,0 +1,1843 @@
+# Copyright 2025 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+import random
+import time
+import logging
+import threading
+from datetime import datetime, timezone
+from collections import deque
+from typing import Any, Dict, List, Optional, Tuple, Union
+from enum import Enum
+
+from fastapi.responses import Response 
+from fastapi.responses import JSONResponse, FileResponse
+
+import joblib
+import uvicorn
+import numpy as np
+import pandas as pd
+from fastapi import FastAPI, HTTPException, status
+from pydantic import BaseModel, Field
+from sklearn.linear_model import BayesianRidge
+from sklearn.preprocessing import StandardScaler
+
+import tempfile
+import shutil
+import os  # Added this import
+
+try:
+    import xgboost as xgb
+    XGBOOST_AVAILABLE = True
+except ImportError:
+    XGBOOST_AVAILABLE = False
+    logging.warning("XGBoost not available. Please install with: pip install xgboost")
+
+try:
+    import lightgbm as lgb
+    LIGHTGBM_AVAILABLE = True
+except ImportError:
+    LIGHTGBM_AVAILABLE = False
+    logging.warning("LightGBM not available. Please install with: pip install lightgbm")
+
+
+class ModelType(str, Enum):
+    BAYESIAN_RIDGE = "bayesian_ridge"
+    XGBOOST = "xgboost"
+    LIGHTGBM = "lightgbm"
+
+
+class RandomDropDeque(deque):
+    def __init__(self, maxlen):
+        super().__init__()
+        self._maxlen = maxlen
+
+    def append(self, item):
+        if len(self) >= self._maxlen:
+            # pick a random index to evict
+            idx = random.randrange(len(self))
+            # rotate so that element at idx moves to the left end
+            self.rotate(-idx)
+            # remove it
+            self.popleft()
+            # rotate back to original ordering
+            self.rotate(idx)
+        super().append(item)
+
+    def appendleft(self, item):
+        if len(self) >= self._maxlen:
+            idx = random.randrange(len(self))
+            # rotate so that element at idx moves to the right end
+            self.rotate(len(self) - idx - 1)
+            self.pop()
+            # rotate back
+            self.rotate(-(len(self) - idx - 1))
+        super().appendleft(item)
+
+
+# --- Configuration ---
+class Settings:
+    """
+    Configuration class for the latency predictor server.
+    Reads settings from environment variables with sensible defaults.
+    """
+    TTFT_MODEL_PATH: str = os.getenv("LATENCY_TTFT_MODEL_PATH", "/tmp/models/ttft.joblib")
+    TPOT_MODEL_PATH: str = os.getenv("LATENCY_TPOT_MODEL_PATH", "/tmp/models/tpot.joblib")
+    TTFT_SCALER_PATH: str = os.getenv("LATENCY_TTFT_SCALER_PATH", "/tmp/models/ttft_scaler.joblib")
+    TPOT_SCALER_PATH: str = os.getenv("LATENCY_TPOT_SCALER_PATH", "/tmp/models/tpot_scaler.joblib")
+    RETRAINING_INTERVAL_SEC: int = int(os.getenv("LATENCY_RETRAINING_INTERVAL_SEC", 1800))
+    MIN_SAMPLES_FOR_RETRAIN_FRESH: int = int(os.getenv("LATENCY_MIN_SAMPLES_FOR_RETRAIN_FRESH", 10))
+    MIN_SAMPLES_FOR_RETRAIN: int = int(os.getenv("LATENCY_MIN_SAMPLES_FOR_RETRAIN", 1000))
+    MAX_TRAINING_DATA_SIZE_PER_BUCKET: int = int(os.getenv("LATENCY_MAX_TRAINING_DATA_SIZE_PER_BUCKET", 10000))
+    TEST_TRAIN_RATIO: float = float(os.getenv("LATENCY_TEST_TRAIN_RATIO", "0.1"))  # Default 1:10 (10% test, 90% train)
+    MAX_TEST_DATA_SIZE: int = int(os.getenv("LATENCY_MAX_TEST_DATA_SIZE", "1000"))  # Max test samples to keep
+    MODEL_TYPE: str = os.getenv("LATENCY_MODEL_TYPE", "xgboost")  # Default to XGBoost
+    QUANTILE_ALPHA: float = float(os.getenv("LATENCY_QUANTILE_ALPHA", "0.9"))  # p90 quantile
+    SAMPLE_WEIGHTING_FOR_PREFIX_CACHE: bool = os.getenv("LATENCY_SAMPLE_WEIGHTING_FOR_PREFIX_CACHE", "false").lower() == "true"
+
+settings = Settings()
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Add this to your Pydantic models section
+class ModelInfoResponse(BaseModel):
+    model_type: str
+    xgboost_available: bool
+    lightgbm_available: bool = Field(default=False, description="Whether LightGBM is available")  # FIXED: Added this field
+
+    is_ready: bool
+    ttft_training_samples: int = Field(default=0, description="Number of TTFT training samples")
+    tpot_training_samples: int = Field(default=0, description="Number of TPOT training samples") 
+    ttft_test_samples: int = Field(default=0, description="Number of TTFT test samples")
+    tpot_test_samples: int = Field(default=0, description="Number of TPOT test samples")
+    last_retrain_time: Optional[datetime] = Field(default=None, description="Last retraining timestamp")
+    min_samples_for_retrain: int = Field(default=0, description="Minimum samples required for retraining")
+    retraining_interval_sec: int = Field(default=0, description="Retraining interval in seconds")
+    
+
+class FlushRequest(BaseModel):
+    flush_training_data: bool = Field(default=True, description="Flush training data buckets")
+    flush_test_data: bool = Field(default=True, description="Flush test data")
+    flush_metrics: bool = Field(default=True, description="Flush quantile metric scores")
+    reason: Optional[str] = Field(default=None, description="Optional reason for flushing")
+
+class FlushResponse(BaseModel):
+    success: bool
+    flushed_at: datetime
+    reason: Optional[str] = None
+    ttft_training_samples_flushed: int
+    tpot_training_samples_flushed: int
+    ttft_test_samples_flushed: int
+    tpot_test_samples_flushed: int
+    metrics_cleared: bool
+    message: str
+
+def quantile_loss(y_true, y_pred, quantile):
+    """
+    Calculate quantile loss (also known as pinball loss).
+    
+    For quantile τ (tau), the loss is:
+    - (τ - 1) * (y_true - y_pred) if y_true < y_pred (under-prediction)
+    - τ * (y_true - y_pred) if y_true >= y_pred (over-prediction)
+    
+    Args:
+        y_true: actual values
+        y_pred: predicted quantile values
+        quantile: the quantile being predicted (e.g., 0.9 for p90)
+    
+    Returns:
+        Mean quantile loss
+    """
+    errors = y_true - y_pred
+    loss = np.where(errors >= 0, quantile * errors, (quantile - 1) * errors)
+    return np.mean(loss)
+
+
+def quantile_coverage(y_true, y_pred, quantile):
+    """
+    Calculate quantile coverage - the proportion of actual values that fall below the predicted quantile.
+    
+    For a well-calibrated p90 model, this should be close to 0.9 (90%).
+    
+    Args:
+        y_true: actual values
+        y_pred: predicted quantile values
+        quantile: the quantile being predicted (e.g., 0.9 for p90)
+    
+    Returns:
+        Coverage percentage (0-100)
+    """
+    below_prediction = np.sum(y_true <= y_pred)
+    coverage = below_prediction / len(y_true)
+    return coverage * 100
+
+
+def quantile_violation_rate(y_true, y_pred, quantile):
+    """
+    Calculate quantile violation rate - the proportion of times actual values exceed the predicted quantile.
+    
+    For a well-calibrated p90 model, this should be close to 0.1 (10%).
+    
+    Args:
+        y_true: actual values
+        y_pred: predicted quantile values
+        quantile: the quantile being predicted (e.g., 0.9 for p90)
+    
+    Returns:
+        Violation rate percentage (0-100)
+    """
+    violations = np.sum(y_true > y_pred)
+    violation_rate = violations / len(y_true)
+    return violation_rate * 100
+
+
+class LatencyPredictor:
+    """
+    Manages model training, prediction, and data handling.
+    """
+    def __init__(self, model_type: str = None):
+        # Set model type with validation
+        if model_type is None:
+            model_type = settings.MODEL_TYPE
+    
+        if model_type not in [ModelType.BAYESIAN_RIDGE, ModelType.XGBOOST, ModelType.LIGHTGBM]:
+            raise ValueError(f"Invalid model_type: {model_type}. Must be one of {list(ModelType)}")
+    
+        if model_type == ModelType.XGBOOST and not XGBOOST_AVAILABLE:
+            logging.warning("XGBoost requested but not available. Falling back to Bayesian Ridge.")
+            model_type = ModelType.BAYESIAN_RIDGE
+
+        if model_type == ModelType.LIGHTGBM and not LIGHTGBM_AVAILABLE:
+            logging.warning("LightGBM requested but not available. Falling back to Bayesian Ridge.")
+            model_type = ModelType.BAYESIAN_RIDGE
+
+        self.model_type = ModelType(model_type)
+        self.quantile = settings.QUANTILE_ALPHA
+        logging.info(f"Initialized LatencyPredictor with model type: {self.model_type}, quantile: {self.quantile}")
+
+        # Data buckets for sampling
+        self.cache_buckets = int(1.0 / 0.05)  # 20 buckets for cache percentage (0-100% in 5% increments)
+        self.queue_buckets = 5  # 0, 1-2, 3-5, 6-10, 11+ waiting requests
+        self.prefix_buckets = 4  # NEW: 0-25%, 25-50%, 50-75%, 75-100% prefix cache score
+        self.bucket_size = settings.MAX_TRAINING_DATA_SIZE_PER_BUCKET 
+
+        # Data buckets with tuple keys: (queue_bucket, cache_bucket, prefix_bucket)
+        self.ttft_data_buckets = {
+            (q, c, p): deque(maxlen=self.bucket_size)
+            for q in range(self.queue_buckets)
+            for c in range(self.cache_buckets)
+            for p in range(self.prefix_buckets)  # NEW: Added prefix dimension
+        }
+        self.tpot_data_buckets = {
+            (q, c, p): deque(maxlen=self.bucket_size)
+            for q in range(self.queue_buckets)
+            for c in range(self.cache_buckets)
+            for p in range(self.prefix_buckets)  # NEW: Added prefix dimension
+        }
+
+        # Test data storage with configurable max size
+        self.ttft_test_data = deque(maxlen=settings.MAX_TEST_DATA_SIZE)
+        self.tpot_test_data = deque(maxlen=settings.MAX_TEST_DATA_SIZE)
+    
+        # Quantile-specific metric tracking (store last 5 scores)
+        self.ttft_quantile_loss_scores = deque(maxlen=5)
+        self.tpot_quantile_loss_scores = deque(maxlen=5)
+        self.ttft_coverage_scores = deque(maxlen=5)
+        self.tpot_coverage_scores = deque(maxlen=5)
+        self.ttft_violation_rates = deque(maxlen=5)
+        self.tpot_violation_rates = deque(maxlen=5)
+
+        self.ttft_model = None
+        self.tpot_model = None
+        self.ttft_scaler = None
+        self.tpot_scaler = None
+    
+        self.ttft_coefficients = None  # Will store descaled coefficients as dict
+        self.tpot_coefficients = None  # Will store descaled coefficients as dict
+
+        self.lock = threading.Lock()
+        self.last_retrain_time = None
+        self._shutdown_event = threading.Event()
+        self._training_thread: threading.Thread = None
+        
+    def _get_prefix_bucket(self, prefix_score: float) -> int:
+        """Map prefix cache score to bucket index."""
+        score = max(0.0, min(1.0, prefix_score))
+        return min(int(score * self.prefix_buckets), self.prefix_buckets - 1)
+        
+    def _get_queue_bucket(self, num_waiting: int) -> int:
+        """Map number of waiting requests to queue bucket index."""
+        if num_waiting == 0:
+            return 0
+        elif num_waiting <= 2:
+            return 1
+        elif num_waiting <= 5:
+            return 2
+        elif num_waiting <= 10:
+            return 3
+        else:
+            return 4  # 11+ requests
+
+    def _get_cache_bucket(self, cache_percentage: float) -> int:
+        """Map cache percentage to cache bucket index."""
+        pct = max(0.0, min(1.0, cache_percentage))
+        return min(int(pct * self.cache_buckets), self.cache_buckets - 1)
+
+    def _get_bucket_key(self, sample: dict) -> tuple:
+        """Get (queue_bucket, cache_bucket) tuple key for a sample."""
+        queue_bucket = self._get_queue_bucket(sample['num_request_waiting'])
+        cache_bucket = self._get_cache_bucket(sample['kv_cache_percentage'])
+        prefix_bucket = self._get_prefix_bucket(sample['prefix_cache_score'])  # NEW
+
+        return (queue_bucket, cache_bucket, prefix_bucket)
+
+    def _store_descaled_coefficients(self, model, scaler, feature_names, model_name):
+        """
+        Store descaled coefficients for Bayesian Ridge models.
+        Returns a dict with feature names as keys and coefficients as values.
+        """
+        if self.model_type != ModelType.BAYESIAN_RIDGE or model is None or scaler is None:
+            return None
+            
+        try:
+            # Get scaled coefficients and scaler parameters
+            coef_scaled = model.coef_
+            scale, mean = scaler.scale_, scaler.mean_
+            
+            # Descale coefficients: w_original = w_scaled / scale
+            w_orig = coef_scaled / scale
+            
+            # Calculate descaled intercept: b_orig = b_scaled - sum(w_scaled * mean / scale)
+            intercept = float(model.intercept_) - float(np.dot(coef_scaled, mean / scale))
+            
+            # Create coefficient dictionary
+            coefficients = {"intercept": intercept}
+            for feature, coef in zip(feature_names, w_orig):
+                coefficients[feature] = float(coef)
+                
+            logging.info(f"Stored descaled coefficients for {model_name}: {coefficients}")
+            return coefficients
+            
+        except Exception as e:
+            logging.error(f"Error storing descaled coefficients for {model_name}: {e}")
+            return None
+        
+    def _prepare_features_with_interaction(self, df: pd.DataFrame, model_type: str) -> pd.DataFrame:
+        """
+        Prepare features with interaction terms for better model learning.
+    
+        Args:
+            df: DataFrame with raw features
+            model_type: 'ttft' or 'tpot'
+    
+        Returns:
+            DataFrame with engineered features including interactions
+        """
+        if model_type == "ttft":
+            # Create interaction: prefix score * input length
+            # This captures that prefix caching benefit scales with input size
+            df['effective_input_tokens'] = (1-df['prefix_cache_score']) * (df['input_token_length'])
+            df['prefill_score_bucket'] = (
+            (df['prefix_cache_score'].clip(0, 1) * self.prefix_buckets)
+            .astype(int)
+            .clip(upper=self.prefix_buckets - 1)
+        )
+
+            # make it categorical for tree models (safe for LGB, XGB with enable_categorical)
+            df['prefill_score_bucket'] = pd.Categorical(df['prefill_score_bucket'], categories=[0,1,2,3], ordered=True)
+
+
+            # Return TTFT features with interaction
+            feature_cols = [
+            'kv_cache_percentage',
+            'input_token_length',
+            'num_request_waiting',
+            'num_request_running',
+            'prefix_cache_score',
+            'effective_input_tokens',
+            'prefill_score_bucket'
+            ]
+        
+            return df[feature_cols]
+        
+        else:  # tpot
+            # TPOT doesn't use prefix_cache_score, so no interaction needed
+            feature_cols = [
+                'kv_cache_percentage',
+                'input_token_length',
+                'num_request_waiting',
+                'num_request_running',
+                'num_tokens_generated'
+            ]
+        
+            return df[feature_cols]
+
+
+    def shutdown(self):
+        """Signal the training thread to exit and join it."""
+        self._shutdown_event.set()
+        if self._training_thread is not None:
+            self._training_thread.join()
+
+    @property
+    def is_ready(self) -> bool:
+        """Checks if all models and scalers are loaded/trained."""
+        if self.model_type == ModelType.BAYESIAN_RIDGE:
+            return all([self.ttft_model, self.tpot_model, self.ttft_scaler, self.tpot_scaler])
+        else:  # XGBoost or LightGBM
+            return all([self.ttft_model, self.tpot_model])
+
+    @is_ready.setter
+    def is_ready(self, value: bool):
+        if not isinstance(value, bool):
+            raise ValueError("is_ready must be a boolean value.")
+        self._is_ready_override = value
+
+    def _all_samples(self, buckets: dict) -> list:
+        samples = []
+        for bucket_deque in buckets.values():
+            samples.extend(bucket_deque)
+        return samples
+    
+
+
+    def _train_model_with_scaling(self, features: pd.DataFrame, target: pd.Series, model_name: str = None, sample_weight: np.ndarray = None) -> Union[Tuple[BayesianRidge, StandardScaler], xgb.XGBRegressor, lgb.LGBMRegressor]:
+
+        try:
+            if len(features) == 0 or len(target) == 0:
+                raise ValueError("Empty training data")
+            if features.isnull().any().any() or target.isnull().any():
+                raise ValueError("Training data contains NaN values")
+            if np.isinf(features.values).any() or np.isinf(target.values).any():
+                raise ValueError("Training data contains infinite values")
+
+            if self.model_type == ModelType.BAYESIAN_RIDGE:
+                features = features.drop(columns=['prefill_score_bucket'], errors='ignore')
+                scaler = StandardScaler()
+                features_scaled = scaler.fit_transform(features)
+                if np.isnan(features_scaled).any() or np.isinf(features_scaled).any():
+                    raise ValueError("Scaling produced invalid values")
+
+                # For Bayesian Ridge, we'll approximate quantile regression by training on the mean
+                # but adjusting predictions later. This is not ideal but Bayesian Ridge doesn't
+                # natively support quantile regression.
+                model = BayesianRidge(compute_score=True)
+                model.fit(features_scaled, target, sample_weight=sample_weight)
+                return model, scaler
+
+            elif self.model_type == ModelType.XGBOOST:  # XGBoost with quantile regression
+                if model_name == "ttft":
+                     # enforce your TTFT feature order
+                        ttft_order = [
+                            "kv_cache_percentage", "input_token_length", "num_request_waiting",
+                            "num_request_running", "prefix_cache_score", "effective_input_tokens", "prefill_score_bucket"
+                        ]
+                        if list(features.columns) != ttft_order:
+                            try:
+                                features = features[ttft_order]
+                            except Exception:
+                                raise ValueError(f"TTFT features must be exactly {ttft_order}; got {list(features.columns)}")
+
+                        # ---- (A) Build a warm-start stump that must split on prefill_score_bucket ----
+                        # Train on the SAME full feature set, but freeze all other features to constants,
+                        # so the only useful split is the prefix bucket.
+                        features_stump = features.copy()
+                        for col in features_stump.columns:
+                            if col != "prefill_score_bucket":
+                                # keep dtype, set to a constant scalar
+                                const_val = features_stump[col].iloc[0]
+                                features_stump[col] = const_val
+
+                        # ensure prefill bucket is int codes if it's categorical
+                        if str(features_stump["prefill_score_bucket"].dtype) == "category":
+                            features_stump["prefill_score_bucket"] = (
+                                features_stump["prefill_score_bucket"].cat.codes.astype("int32")
+                            )
+                        else:
+                            features_stump["prefill_score_bucket"] = features_stump["prefill_score_bucket"].astype("int32")
+
+                        
+                        model = xgb.XGBRegressor(
+                            n_estimators=200,
+                            max_depth=6,
+                            learning_rate=0.05,
+                            subsample=0.8,
+                            colsample_bytree=0.8,
+                            min_child_weight=5,
+                            gamma=0.2,
+                            reg_alpha=0.01,
+                            reg_lambda=0.1,
+                            objective="reg:quantileerror",
+                            quantile_alpha=self.quantile,
+                            tree_method="hist",
+                            n_jobs=-1,
+                            random_state=42,
+                            verbosity=1,
+                            enable_categorical=True,
+                            )
+                        model.fit(features, target, sample_weight=sample_weight)    
+                        return model
+
+
+                elif model_name == "tpot":
+                    tpot_order = ["kv_cache_percentage","input_token_length","num_request_waiting","num_request_running","num_tokens_generated"]
+                    if list(features.columns) != tpot_order:
+                        try:
+                            features = features[tpot_order]
+                        except Exception as _:
+                            raise ValueError(f"TPOT features must be exactly {tpot_order}; got {list(features.columns)}")
+                    mono_str = "(1,1,1,1,1)" 
+                else:
+                    mono_str = "(0,0,0,0,0)"  # default
+                model = xgb.XGBRegressor(
+                n_estimators=200,            # Number of trees to build (moderate value for balanced accuracy and speed)
+                max_depth=6,                 # Depth of trees; 6 is typically a sweet spot balancing bias/variance
+                learning_rate=0.05,          # Smaller learning rate to achieve stable convergence
+                subsample=0.8,               # Use 80% of data per tree (adds regularization & reduces overfitting)
+                colsample_bytree=0.8,        # Use 80% of features per tree (improves generalization)
+        
+                # Key parameters for accurate quantile regression:
+                min_child_weight=5,          # Low value allows fine-grained splits near p90 boundary (prevents overprediction)
+                gamma=0.2,                  # Low gamma allows splits with small loss reduction (critical for quantile accuracy)
+                #monotone_constraints=mono_str,  # Enforce monotonicity based on feature impact on latency   
+                # Regularization to prevent overfitting:
+                reg_alpha=0.01,              # L1 regularization (Lasso) - encourages sparsity
+                reg_lambda=0.1,              # L2 regularization (Ridge) - prevents large coefficients
+        
+                # Quantile regression configuration:
+                objective="reg:quantileerror",    # Quantile loss (pinball loss) for quantile regression
+                quantile_alpha=self.quantile,     # Target quantile (e.g., 0.9 for p90)
+        
+                # Performance optimization:
+                tree_method='hist',          # Efficient histogram algorithm; optimal for large datasets
+                n_jobs=-1,                   # Utilize all CPU cores for parallel training
+                random_state=42,             # Ensures reproducible results
+                verbosity=1,
+                enable_categorical=True       # Enable categorical feature support   
+    )
+                model.fit(features, target, sample_weight=sample_weight)
+                return model
+            elif self.model_type == ModelType.LIGHTGBM:  # LightGBM with quantile regression
+                model = lgb.LGBMRegressor(
+                n_estimators=200,           # Number of trees
+                max_depth=6,                # Maximum tree depth
+                learning_rate=0.05,         # Learning rate
+                subsample=0.8,              # Row sampling ratio
+                colsample_bytree=0.8,       # Column sampling ratio
+                min_child_samples=20,       # Minimum samples in leaf
+                reg_alpha=0.1,              # L1 regularization
+                reg_lambda=0.1,             # L2 regularization
+                objective="quantile",       # Quantile regression objective
+                alpha=self.quantile,        # Quantile level (e.g., 0.9 for p90)
+                n_jobs=-1,                  # Use all cores
+                random_state=42,            # Reproducibility
+                verbosity=-1,               # Suppress warnings
+                force_col_wise=True         # Better for small datasets
+            )
+                model.fit(features, target, sample_weight=sample_weight, categorical_feature=['prefill_score_bucket'] if model_name == "ttft" else None)
+                return model
+                
+        except Exception as e:
+            logging.error(f"Error in _train_model_with_scaling: {e}", exc_info=True)
+            raise
+        
+    def _calculate_quantile_metrics_on_test(self, model, scaler, test_data, feature_cols, target_col):
+        """Calculate quantile-specific metrics on test data"""
+        try:
+            df = pd.DataFrame(test_data).dropna()
+            df = df[df[target_col] > 0]
+            
+            if len(df) < 2:
+                return None, None, None
+            
+            X = df[feature_cols]
+            if self.model_type == ModelType.BAYESIAN_RIDGE and scaler is not None:
+                X = scaler.transform(X)
+            
+            y_true = df[target_col].values
+            y_pred = model.predict(X)
+            
+            # For Bayesian Ridge (which doesn't do true quantile regression), 
+            # we'll estimate the quantile by adding a factor to the mean prediction
+            if self.model_type == ModelType.BAYESIAN_RIDGE:
+                # Rough approximation: add some multiple of std to get to desired quantile
+                # This is a simplification - in practice you'd want proper quantile regression
+                std_factor = 1.28 if self.quantile == 0.9 else (2.0 if self.quantile == 0.95 else 0.674)
+                _, y_std = model.predict(X, return_std=True)
+                y_pred = y_pred + std_factor * y_std
+            
+            # Calculate quantile-specific metrics
+            ql = quantile_loss(y_true, y_pred, self.quantile)
+            coverage = quantile_coverage(y_true, y_pred, self.quantile)
+            violation_rate = quantile_violation_rate(y_true, y_pred, self.quantile)
+            
+            return ql, coverage, violation_rate
+            
+        except Exception as e:
+            logging.error(f"Error calculating quantile metrics: {e}", exc_info=True)
+            return None, None, None
+
+    def _create_default_model(self, model_type: str) -> Union[Tuple[BayesianRidge, StandardScaler], xgb.XGBRegressor, lgb.LGBMRegressor]:
+
+        """Creates and trains a simple default model with initial priors."""
+        try:
+            logging.info(f"Creating default '{model_type}' model with priors.")
+            if model_type == "ttft":
+                features = pd.DataFrame({
+                    'kv_cache_percentage': [0.0, ],
+                    'input_token_length': [1, ],
+                    'num_request_waiting': [0, ],
+                    'num_request_running': [0, ],
+                    'prefix_cache_score': [0.0, ]  # Added prefix_cache_score
+                })
+                features = self._prepare_features_with_interaction(features, "ttft")
+                target = pd.Series([10.0])
+              
+            else:
+                features = pd.DataFrame({
+                    'kv_cache_percentage': [0.0],
+                    'input_token_length': [1],  # Added input_token_length
+                    'num_request_waiting': [0, ],
+                    'num_request_running': [0, ],
+                    'num_tokens_generated': [1,]
+                })
+                target = pd.Series([10.0])
+            return self._train_model_with_scaling(features, target, model_name=model_type)
+        except Exception as e:
+            logging.error(f"Error creating default model for {model_type}: {e}", exc_info=True)
+            raise
+
+    def train(self):
+        try:
+            with self.lock:
+                ttft_snap = list(self._all_samples(self.ttft_data_buckets))
+                tpot_snap = list(self._all_samples(self.tpot_data_buckets))
+                total = len(ttft_snap) + len(tpot_snap)
+                if total < settings.MIN_SAMPLES_FOR_RETRAIN:
+                    logging.info(f"Skipping training: only {total} samples (< {settings.MIN_SAMPLES_FOR_RETRAIN}).")
+                    return
+                logging.info(f"Initiating training with {total} samples using {self.model_type} for quantile {self.quantile}.")
+
+            new_ttft_model = new_ttft_scaler = None
+            new_tpot_model = new_tpot_scaler = None
+
+            # Train TTFT
+            if ttft_snap:
+                raw_ttft = pd.DataFrame(ttft_snap).dropna()
+                raw_ttft = raw_ttft[raw_ttft['actual_ttft_ms'] > 0]
+                df_ttft = self._prepare_features_with_interaction(raw_ttft.copy(), model_type="ttft")
+                print(f"TTFT training data size: {len(df_ttft)} with sample data: {df_ttft.columns.tolist()}")
+                if len(df_ttft) >= settings.MIN_SAMPLES_FOR_RETRAIN:
+                    # Updated TTFT features to include prefix_cache_score
+                    ttft_feature_cols_tree = [
+                    'kv_cache_percentage','input_token_length','num_request_waiting',
+                    'num_request_running','prefix_cache_score','effective_input_tokens','prefill_score_bucket'
+                ]
+                ttft_feature_cols_br = [
+                    'kv_cache_percentage','input_token_length','num_request_waiting',
+                    'num_request_running','prefix_cache_score','effective_input_tokens'
+                ]
+
+                # Build X_ttft for all model types, then trim for BR
+                X_ttft = df_ttft[ttft_feature_cols_tree]
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    X_ttft = X_ttft[ttft_feature_cols_br]
+
+                y_ttft = raw_ttft['actual_ttft_ms']
+
+                try:
+                        # raw_ttft still has the original columns including 'prefix_cache_score'
+                        raw_ttft['_prefix_bucket'] = raw_ttft['prefix_cache_score'].clip(0, 1).apply(
+                            lambda s: min(int(s * self.prefix_buckets), self.prefix_buckets - 1)
+                        )
+
+                        bucket_counts = raw_ttft['_prefix_bucket'].value_counts().to_dict()
+                        total_ttft = len(raw_ttft)
+                        num_buckets = max(1, len(bucket_counts))
+                        epsilon = 1.0
+                        bucket_weights = {p: total_ttft / (num_buckets * (cnt + epsilon)) for p, cnt in bucket_counts.items()}
+                        sample_weight_ttft = None
+                        if settings.SAMPLE_WEIGHTING_FOR_PREFIX_CACHE:
+                            sample_weight_ttft = raw_ttft['_prefix_bucket'].map(bucket_weights).astype(float).to_numpy()
+                            sample_weight_ttft *= (len(sample_weight_ttft) / sample_weight_ttft.sum())
+
+                        result = self._train_model_with_scaling(X_ttft, y_ttft, model_name="ttft", sample_weight=sample_weight_ttft)
+                        if self.model_type == ModelType.BAYESIAN_RIDGE:
+                            new_ttft_model, new_ttft_scaler = result
+                        else:
+                            new_ttft_model = result
+                            new_ttft_scaler = None
+
+                        # Quantile metrics on test set
+                        ql = coverage = violation_rate = None
+                        ttft_test_raw = pd.DataFrame(list(self.ttft_test_data)).dropna()
+                        if not ttft_test_raw.empty:
+                            ttft_test_fe = self._prepare_features_with_interaction(ttft_test_raw.copy(), "ttft")
+                            cols = ttft_feature_cols_br if self.model_type == ModelType.BAYESIAN_RIDGE else ttft_feature_cols_tree
+                            test_records = ttft_test_fe[cols].assign(actual_ttft_ms=ttft_test_raw['actual_ttft_ms']).to_dict("records")
+                            ql, coverage, violation_rate = self._calculate_quantile_metrics_on_test(
+                                new_ttft_model, new_ttft_scaler, test_records, cols, 'actual_ttft_ms'
+                            )
+
+
+                        
+                        if ql is not None:
+                            self.ttft_quantile_loss_scores.append(ql)
+                            self.ttft_coverage_scores.append(coverage)
+                            self.ttft_violation_rates.append(violation_rate)
+                            logging.info(f"TTFT model trained on {len(df_ttft)} samples. "
+                                       f"Quantile Loss = {ql:.4f}, "
+                                       f"Coverage = {coverage:.2f}% (target: {self.quantile*100:.0f}%), "
+                                       f"Violation Rate = {violation_rate:.2f}% (target: {(1-self.quantile)*100:.0f}%)")
+                        else:
+                            logging.info(f"TTFT model trained on {len(df_ttft)} samples. Quantile metrics = N/A (insufficient test data)")
+
+                except Exception:
+                        logging.error("Error training TTFT model", exc_info=True)
+
+
+            # Train TPOT
+            if tpot_snap:
+                df_tpot = pd.DataFrame(tpot_snap).dropna()
+                df_tpot = df_tpot[df_tpot['actual_tpot_ms'] > 0]
+                if len(df_tpot) >= settings.MIN_SAMPLES_FOR_RETRAIN:
+                    # TPOT features remain unchanged
+                    X_tpot = df_tpot[['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 'num_request_running', 'num_tokens_generated']]
+                    y_tpot = df_tpot['actual_tpot_ms']
+                    try:
+                        result = self._train_model_with_scaling(X_tpot, y_tpot, model_name="tpot")
+                        if self.model_type == ModelType.BAYESIAN_RIDGE:
+                            new_tpot_model, new_tpot_scaler = result
+                        else:
+                            new_tpot_model = result
+                            new_tpot_scaler = None
+                        
+                        # Calculate quantile metrics on test data
+                        tpot_feature_cols = ['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 'num_request_running', 'num_tokens_generated']
+                        ql, coverage, violation_rate = self._calculate_quantile_metrics_on_test(
+                            new_tpot_model, new_tpot_scaler, 
+                            list(self.tpot_test_data), tpot_feature_cols, 'actual_tpot_ms'
+                        )
+                        
+                        if ql is not None:
+                            self.tpot_quantile_loss_scores.append(ql)
+                            self.tpot_coverage_scores.append(coverage)
+                            self.tpot_violation_rates.append(violation_rate)
+                            logging.info(f"TPOT model trained on {len(df_tpot)} samples. "
+                                       f"Quantile Loss = {ql:.4f}, "
+                                       f"Coverage = {coverage:.2f}% (target: {self.quantile*100:.0f}%), "
+                                       f"Violation Rate = {violation_rate:.2f}% (target: {(1-self.quantile)*100:.0f}%)")
+                        else:
+                            logging.info(f"TPOT model trained on {len(df_tpot)} samples. Quantile metrics = N/A (insufficient test data)")
+                            
+                    except Exception:
+                        logging.error("Error training TPOT model", exc_info=True)
+                else:
+                    logging.warning("Not enough TPOT samples, skipping TPOT training.")
+
+            with self.lock:
+                if new_ttft_model:
+                    self.ttft_model = new_ttft_model
+                    if new_ttft_scaler is not None:
+                        self.ttft_scaler = new_ttft_scaler
+                    
+                    # Store descaled coefficients for Bayesian Ridge
+                    if self.model_type == ModelType.BAYESIAN_RIDGE:
+                        ttft_features = ttft_feature_cols_br  # no 'prefill_score_bucket' for BR
+                        self.ttft_coefficients = self._store_descaled_coefficients(
+                        new_ttft_model, new_ttft_scaler, ttft_features, "TTFT"
+                )
+                        
+                if new_tpot_model:
+                    self.tpot_model = new_tpot_model
+                    if new_tpot_scaler is not None:
+                        self.tpot_scaler = new_tpot_scaler
+                    
+                    # Store descaled coefficients for Bayesian Ridge
+                    if self.model_type == ModelType.BAYESIAN_RIDGE:
+                        tpot_features = ['kv_cache_percentage', 'input_token_length', 
+                                       'num_request_waiting', 'num_request_running', 'num_tokens_generated']
+                        self.tpot_coefficients = self._store_descaled_coefficients(
+                            new_tpot_model, new_tpot_scaler, tpot_features, "TPOT"
+                        )
+                
+                if self.is_ready:
+                    self.last_retrain_time = datetime.now(timezone.utc)
+                    try:
+                        self._save_models_unlocked()
+                    except Exception:
+                        logging.error("Error saving models after training.", exc_info=True)
+        except Exception as e:
+            logging.error(f"Critical error in train(): {e}", exc_info=True)
+
+    def predict(self, features: dict) -> Tuple[float, float, float, float]:
+        try:
+            with self.lock:
+                if not self.is_ready:
+                    raise HTTPException(status_code=503, detail="Models not ready")
+                required = ['kv_cache_percentage', 'input_token_length', 'num_request_waiting', 'num_request_running', 'num_tokens_generated', 'prefix_cache_score']
+                for f in required:
+                    if f not in features:
+                        raise ValueError(f"Missing required feature: {f}")
+                    if not isinstance(features[f], (int, float)):
+                        raise ValueError(f"Invalid type for feature {f}: expected number")
+
+                # Updated TTFT features to include prefix_cache_score
+                ttft_cols = ['kv_cache_percentage','input_token_length','num_request_waiting','num_request_running','prefix_cache_score']
+                tpot_cols = ['kv_cache_percentage','input_token_length','num_request_waiting','num_request_running','num_tokens_generated']
+                
+                # Create DataFrames for predictions
+                df_ttft = pd.DataFrame([{col: features[col] for col in ttft_cols}])
+                # Add interaction term for TTFT
+                df_ttft = self._prepare_features_with_interaction(df_ttft, model_type="ttft")
+                df_tpot = pd.DataFrame([{col: features[col] for col in tpot_cols}])
+
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    # Use scaling for Bayesian Ridge
+                    df_ttft = df_ttft.drop(columns=['prefill_score_bucket'], errors='ignore')
+                    ttft_scaled = self.ttft_scaler.transform(df_ttft)
+                    tpot_scaled = self.tpot_scaler.transform(df_tpot)
+
+                    ttft_pred_mean, ttft_std = self.ttft_model.predict(ttft_scaled, return_std=True)
+                    tpot_pred_mean, tpot_std = self.tpot_model.predict(tpot_scaled, return_std=True)
+                    
+                    # Approximate quantile prediction by adding factor to mean
+                    std_factor = 1.28 if self.quantile == 0.9 else (2.0 if self.quantile == 0.95 else 0.674)
+                    ttft_pred = ttft_pred_mean[0] + std_factor * ttft_std[0]
+                    tpot_pred = tpot_pred_mean[0] + std_factor * tpot_std[0]
+                    
+                    return ttft_pred, tpot_pred, ttft_std[0], tpot_std[0]
+                
+                elif self.model_type == ModelType.XGBOOST:
+                    # XGBoost quantile regression directly predicts the quantile
+                    ttft_pred = self.ttft_model.predict(df_ttft)
+                    tpot_pred = self.tpot_model.predict(df_tpot)
+                
+                    # For XGBoost quantile regression, uncertainty estimation is more complex
+                    ttft_std = ttft_pred[0] * 0.1  # 10% of prediction as uncertainty estimate
+                    tpot_std = tpot_pred[0] * 0.1
+                
+                    return ttft_pred[0], tpot_pred[0], ttft_std, tpot_std
+                
+                else:  # LightGBM with quantile regression
+                    # LightGBM quantile regression directly predicts the quantile
+                    ttft_pred = self.ttft_model.predict(df_ttft)
+                    tpot_pred = self.tpot_model.predict(df_tpot)
+                
+                    # For LightGBM quantile regression, use a similar uncertainty estimate as XGBoost
+                    ttft_std = ttft_pred[0] * 0.1  # 10% of prediction as uncertainty estimate
+                    tpot_std = tpot_pred[0] * 0.1
+                
+                    return ttft_pred[0], tpot_pred[0], ttft_std, tpot_std
+                    
+        except ValueError as ve:
+            logging.warning(f"Client error in predict(): {ve}")
+            raise HTTPException(status_code=400, detail=str(ve))
+        except HTTPException:
+            raise
+        except Exception as e:
+            logging.error("Error in predict():", exc_info=True)
+            raise HTTPException(status_code=500, detail="Internal error during prediction")
+
+    def add_training_sample(self, sample: dict):
+        try:
+            required = ['kv_cache_percentage', 'actual_ttft_ms', 'actual_tpot_ms', 'num_tokens_generated', 'input_token_length', 'num_request_waiting', 'num_request_running', 'prefix_cache_score']
+            for field in required:
+                if field not in sample or not isinstance(sample[field], (int, float)):
+                    logging.warning(f"Invalid sample field: {field}")
+                    return
+        
+            # Use hash-based deterministic split to ensure consistent train/test assignment
+            # This ensures the same sample always goes to the same split
+            sample_hash = hash(str(sorted(sample.items())))
+            is_test = (sample_hash % 100) < (settings.TEST_TRAIN_RATIO * 100)
+        
+            # Create subsets based on conditions
+            ttft_valid = sample['actual_ttft_ms'] > 0
+            tpot_valid = sample['actual_tpot_ms'] > 0
+        
+            if is_test:
+                # Add to test data only if the respective metric is valid
+                if ttft_valid:
+                    self.ttft_test_data.append(sample.copy())
+                if tpot_valid:
+                    self.tpot_test_data.append(sample.copy())
+            else:
+                # Add to training buckets only if the respective metric is valid
+                bucket_key = self._get_bucket_key(sample)
+            
+                if ttft_valid:
+                    self.ttft_data_buckets[bucket_key].append(sample)
+                if tpot_valid:
+                    self.tpot_data_buckets[bucket_key].append(sample)
+                
+        except Exception as e:
+            logging.error(f"Error adding training sample: {e}", exc_info=True)
+            
+    
+    def add_training_samples(self, samples: list):
+        """Bulk-add multiple training samples in one go."""
+        with self.lock:
+            for sample in samples:
+                try:
+                    # reuse the single-sample logic
+                    self.add_training_sample(sample)
+                except Exception:
+                    # log & continue on individual failures
+                    logging.exception("Failed to add one sample in bulk ingestion")
+
+
+    # Update the _save_models_unlocked method to handle LightGBM model exports
+    def _save_models_unlocked(self):
+        try:
+            if self.ttft_model:
+                os.makedirs(os.path.dirname(settings.TTFT_MODEL_PATH), exist_ok=True)
+                joblib.dump(self.ttft_model, settings.TTFT_MODEL_PATH)
+                logging.info("TTFT model saved.")
+        
+                # Save model-specific exports
+                if self.model_type == ModelType.XGBOOST:
+                    try:
+                        booster = self.ttft_model.get_booster()
+                        raw_trees = booster.get_dump(dump_format="json")
+                        trees = [json.loads(t) for t in raw_trees]
+                
+                        ttft_json_path = settings.TTFT_MODEL_PATH.replace('.joblib', '_trees.json')
+                        with open(ttft_json_path, 'w') as f:
+                            json.dump(trees, f, indent=2)
+                        logging.info(f"TTFT XGBoost trees saved to {ttft_json_path}")
+                    except Exception as e:
+                        logging.error(f"Error saving TTFT XGBoost trees: {e}", exc_info=True)
+            
+                elif self.model_type == ModelType.LIGHTGBM:
+                    try:
+                        # Save LightGBM model as text format
+                        ttft_txt_path = settings.TTFT_MODEL_PATH.replace('.joblib', '_lgb.txt')
+                        self.ttft_model.booster_.save_model(ttft_txt_path)
+                    
+                        # Save feature importances as JSON
+                        feature_names = ['kv_cache_percentage', 'input_token_length', 
+                                       'num_request_waiting', 'num_request_running', 'prefix_cache_score', 'effective_input_tokens', 'prefill_score_bucket']
+                        importances = dict(zip(feature_names, self.ttft_model.feature_importances_))
+                    
+                        ttft_imp_path = settings.TTFT_MODEL_PATH.replace('.joblib', '_importances.json')
+                        with open(ttft_imp_path, 'w') as f:
+                            json.dump(importances, f, indent=2)
+                    
+                        logging.info(f"TTFT LightGBM model saved to {ttft_txt_path}")
+                        logging.info(f"TTFT LightGBM importances saved to {ttft_imp_path}")
+                    except Exception as e:
+                        logging.error(f"Error saving TTFT LightGBM exports: {e}", exc_info=True)
+        
+            if self.ttft_scaler and self.model_type == ModelType.BAYESIAN_RIDGE:
+                os.makedirs(os.path.dirname(settings.TTFT_SCALER_PATH), exist_ok=True)
+                joblib.dump(self.ttft_scaler, settings.TTFT_SCALER_PATH)
+                logging.info("TTFT scaler saved.")
+        
+            if self.tpot_model:
+                os.makedirs(os.path.dirname(settings.TPOT_MODEL_PATH), exist_ok=True)
+                joblib.dump(self.tpot_model, settings.TPOT_MODEL_PATH)
+                logging.info("TPOT model saved.")
+        
+                # Save model-specific exports
+                if self.model_type == ModelType.XGBOOST:
+                    try:
+                        booster = self.tpot_model.get_booster()
+                        raw_trees = booster.get_dump(dump_format="json")
+                        trees = [json.loads(t) for t in raw_trees]
+                
+                        tpot_json_path = settings.TPOT_MODEL_PATH.replace('.joblib', '_trees.json')
+                        with open(tpot_json_path, 'w') as f:
+                            json.dump(trees, f, indent=2)
+                        logging.info(f"TPOT XGBoost trees saved to {tpot_json_path}")
+                    except Exception as e:
+                        logging.error(f"Error saving TPOT XGBoost trees: {e}", exc_info=True)
+            
+                elif self.model_type == ModelType.LIGHTGBM:
+                    try:
+                        # Save LightGBM model as text format
+                        tpot_txt_path = settings.TPOT_MODEL_PATH.replace('.joblib', '_lgb.txt')
+                        self.tpot_model.booster_.save_model(tpot_txt_path)
+                    
+                        # Save feature importances as JSON
+                        feature_names = ['kv_cache_percentage', 'input_token_length', 
+                                       'num_request_waiting', 'num_request_running', 'num_tokens_generated']
+                        importances = dict(zip(feature_names, self.tpot_model.feature_importances_))
+                    
+                        tpot_imp_path = settings.TPOT_MODEL_PATH.replace('.joblib', '_importances.json')
+                        with open(tpot_imp_path, 'w') as f:
+                            json.dump(importances, f, indent=2)
+                    
+                        logging.info(f"TPOT LightGBM model saved to {tpot_txt_path}")
+                        logging.info(f"TPOT LightGBM importances saved to {tpot_imp_path}")
+                    except Exception as e:
+                        logging.error(f"Error saving TPOT LightGBM exports: {e}", exc_info=True)
+        
+            if self.tpot_scaler and self.model_type == ModelType.BAYESIAN_RIDGE:
+                os.makedirs(os.path.dirname(settings.TPOT_SCALER_PATH), exist_ok=True)
+                joblib.dump(self.tpot_scaler, settings.TPOT_SCALER_PATH)
+                logging.info("TPOT scaler saved.")
+        
+        except Exception as e:
+            logging.error(f"Error saving models: {e}", exc_info=True)
+            
+    def flush_training_data(self, flush_training: bool = True, flush_test: bool = True, 
+                   flush_metrics: bool = True, reason: str = None) -> dict:
+        """
+        Manually flush training data, test data, and/or metrics.
+        Returns statistics about what was flushed.
+
+        Args:
+            flush_training: Whether to flush training data buckets
+            flush_test: Whether to flush test data
+            flush_metrics: Whether to flush quantile metric scores
+            reason: Optional reason for flushing (for logging)
+
+        Returns:
+            Dictionary with flush statistics
+        """
+        try:
+            with self.lock:
+                # Count samples before flushing (handles 3D buckets)
+                ttft_training_count = sum(len(bucket) for bucket in self.ttft_data_buckets.values())
+                tpot_training_count = sum(len(bucket) for bucket in self.tpot_data_buckets.values())
+                ttft_test_count = len(self.ttft_test_data)
+                tpot_test_count = len(self.tpot_test_data)
+        
+                reason_str = f" Reason: {reason}" if reason else ""
+                logging.info(
+                    f"Manual flush requested.{reason_str} "
+                    f"Training: {flush_training}, Test: {flush_test}, Metrics: {flush_metrics}"
+                )
+        
+                # Flush training data (now handles 3D buckets automatically)
+                if flush_training:
+                    for bucket_key in self.ttft_data_buckets:
+                        self.ttft_data_buckets[bucket_key].clear()
+                    for bucket_key in self.tpot_data_buckets:
+                        self.tpot_data_buckets[bucket_key].clear()
+                    logging.info(
+                        f"Flushed {ttft_training_count} TTFT and {tpot_training_count} TPOT training samples"
+                    )
+        
+                # Flush test data
+                if flush_test:
+                    self.ttft_test_data.clear()
+                    self.tpot_test_data.clear()
+                    logging.info(
+                        f"Flushed {ttft_test_count} TTFT and {tpot_test_count} TPOT test samples"
+                    )
+        
+                # Clear metrics
+                metrics_cleared = False
+                if flush_metrics:
+                    self.ttft_quantile_loss_scores.clear()
+                    self.tpot_quantile_loss_scores.clear()
+                    self.ttft_coverage_scores.clear()
+                    self.tpot_coverage_scores.clear()
+                    self.ttft_violation_rates.clear()
+                    self.tpot_violation_rates.clear()
+                    metrics_cleared = True
+                    logging.info("Cleared all quantile metric scores")
+        
+                return {
+                    "success": True,
+                    "ttft_training_samples_flushed": ttft_training_count if flush_training else 0,
+                    "tpot_training_samples_flushed": tpot_training_count if flush_training else 0,
+                    "ttft_test_samples_flushed": ttft_test_count if flush_test else 0,
+                    "tpot_test_samples_flushed": tpot_test_count if flush_test else 0,
+                    "metrics_cleared": metrics_cleared
+                }
+        
+        except Exception as e:
+            logging.error(f"Error flushing data: {e}", exc_info=True)
+            raise
+    
+    def load_models(self):
+        try:
+            with self.lock:
+                if os.path.exists(settings.TTFT_MODEL_PATH):
+                    self.ttft_model = joblib.load(settings.TTFT_MODEL_PATH)
+                    if self.model_type == ModelType.BAYESIAN_RIDGE and os.path.exists(settings.TTFT_SCALER_PATH):
+                        self.ttft_scaler = joblib.load(settings.TTFT_SCALER_PATH)
+                else:
+                    result = self._create_default_model("ttft")
+                    if self.model_type == ModelType.BAYESIAN_RIDGE:
+                        self.ttft_model, self.ttft_scaler = result
+                    else:
+                        self.ttft_model = result
+                    settings.MIN_SAMPLES_FOR_RETRAIN = settings.MIN_SAMPLES_FOR_RETRAIN_FRESH
+                    self._save_models_unlocked()
+
+                if os.path.exists(settings.TPOT_MODEL_PATH):
+                    self.tpot_model = joblib.load(settings.TPOT_MODEL_PATH)
+                    if self.model_type == ModelType.BAYESIAN_RIDGE and os.path.exists(settings.TPOT_SCALER_PATH):
+                        self.tpot_scaler = joblib.load(settings.TPOT_SCALER_PATH)
+                else:
+                    result = self._create_default_model("tpot")
+                    if self.model_type == ModelType.BAYESIAN_RIDGE:
+                        self.tpot_model, self.tpot_scaler = result
+                    else:
+                        self.tpot_model = result
+                    settings.MIN_SAMPLES_FOR_RETRAIN = settings.MIN_SAMPLES_FOR_RETRAIN_FRESH
+                    self._save_models_unlocked()
+
+                if not self.is_ready:
+                    raise RuntimeError("Failed to initialize models/scalers")
+        except Exception as e:
+            logging.error(f"Critical error in load_models: {e}", exc_info=True)
+            raise
+        
+    def get_metrics(self) -> str:
+        """Render Prometheus-style metrics: model, coefficients/importances, bucket counts, and quantile-specific scores."""
+        try:
+        # Snapshot models & scalers
+            ttft_model, tpot_model = self.ttft_model, self.tpot_model
+            ttft_scaler, tpot_scaler = self.ttft_scaler, self.tpot_scaler
+
+            lines: List[str] = []
+            # 1) Model type and quantile info
+            lines.append(f'model_type{{type="{self.model_type.value}"}} 1')
+            lines.append(f'model_quantile{{}} {self.quantile}')
+
+            # Helper: emit linear‐model coefs or tree importances
+            def emit_metrics(model, coefficients, feats, prefix):
+                if model is None:
+                    # placeholders
+                    lines.append(f'{prefix}_intercept{{}} 0.0')
+                    kind = "coef" if self.model_type == ModelType.BAYESIAN_RIDGE else "importance"
+                    for f in feats:
+                        lines.append(f'{prefix}_{kind}{{feature="{f}"}} 0.0')
+                    return
+
+                if self.model_type == ModelType.BAYESIAN_RIDGE:
+                    # Use stored descaled coefficients
+                    if coefficients:
+                        lines.append(f'{prefix}_intercept{{}} {coefficients.get("intercept", 0.0):.6f}')
+                        for f in feats:
+                            coef_value = coefficients.get(f, 0.0)
+                            lines.append(f'{prefix}_coef{{feature="{f}"}} {coef_value:.6f}')
+                    else:
+                        # Fallback to zeros if coefficients not available
+                        lines.append(f'{prefix}_intercept{{}} 0.0')
+                        for f in feats:
+                            lines.append(f'{prefix}_coef{{feature="{f}"}} 0.0')
+                else:
+                    # XGBoost/LightGBM importances
+                    try:
+                        imps = model.feature_importances_
+                    except Exception:
+                        imps = [0.0]*len(feats)
+                    lines.append(f'{prefix}_intercept{{}} 0.0')
+                    for f, imp in zip(feats, imps):
+                        lines.append(f'{prefix}_importance{{feature="{f}"}} {imp:.6f}')
+
+            if self.model_type == ModelType.BAYESIAN_RIDGE:
+                ttft_feats = ["kv_cache_percentage","input_token_length","num_request_waiting",
+                  "num_request_running","prefix_cache_score","effective_input_tokens"]
+            else:
+                ttft_feats = ["kv_cache_percentage","input_token_length","num_request_waiting",
+                  "num_request_running","prefix_cache_score","effective_input_tokens","prefill_score_bucket"]
+
+            tpot_feats = ["kv_cache_percentage","input_token_length","num_request_waiting",
+              "num_request_running","num_tokens_generated"]
+            emit_metrics(ttft_model, self.ttft_coefficients, ttft_feats, "ttft")
+            emit_metrics(tpot_model, self.tpot_coefficients, tpot_feats, "tpot")
+
+            # 3) Multi-dimensional bucket counts with 3D keys
+            for (queue_bucket, cache_bucket, prefix_bucket), bucket_deque in self.ttft_data_buckets.items():
+                count = len(bucket_deque)
+                lines.append(f'training_samples_count{{model="ttft",queue_bucket="{queue_bucket}",cache_bucket="{cache_bucket}",prefix_bucket="{prefix_bucket}"}} {count}')
+    
+            for (queue_bucket, cache_bucket, prefix_bucket), bucket_deque in self.tpot_data_buckets.items():
+                count = len(bucket_deque)
+                lines.append(f'training_samples_count{{model="tpot",queue_bucket="{queue_bucket}",cache_bucket="{cache_bucket}",prefix_bucket="{prefix_bucket}"}} {count}')
+    
+            # Summary metrics by queue state
+            for q in range(self.queue_buckets):
+                ttft_total = sum(len(self.ttft_data_buckets[(q, c, p)]) 
+                               for c in range(self.cache_buckets) 
+                               for p in range(self.prefix_buckets))
+                tpot_total = sum(len(self.tpot_data_buckets[(q, c, p)]) 
+                               for c in range(self.cache_buckets) 
+                               for p in range(self.prefix_buckets))
+                lines.append(f'training_samples_queue_total{{model="ttft",queue_bucket="{q}"}} {ttft_total}')
+                lines.append(f'training_samples_queue_total{{model="tpot",queue_bucket="{q}"}} {tpot_total}')
+
+            # Summary metrics by cache state
+            for c in range(self.cache_buckets):
+                ttft_total = sum(len(self.ttft_data_buckets[(q, c, p)]) 
+                               for q in range(self.queue_buckets) 
+                               for p in range(self.prefix_buckets))
+                tpot_total = sum(len(self.tpot_data_buckets[(q, c, p)]) 
+                               for q in range(self.queue_buckets) 
+                               for p in range(self.prefix_buckets))
+                lines.append(f'training_samples_cache_total{{model="ttft",cache_bucket="{c}"}} {ttft_total}')
+                lines.append(f'training_samples_cache_total{{model="tpot",cache_bucket="{c}"}} {tpot_total}')
+
+            # Summary metrics by prefix state
+            for p in range(self.prefix_buckets):
+                ttft_total = sum(len(self.ttft_data_buckets[(q, c, p)]) 
+                               for q in range(self.queue_buckets) 
+                               for c in range(self.cache_buckets))
+                tpot_total = sum(len(self.tpot_data_buckets[(q, c, p)]) 
+                               for q in range(self.queue_buckets) 
+                               for c in range(self.cache_buckets))
+            
+                # Calculate prefix range for this bucket
+                prefix_low = p / self.prefix_buckets
+                prefix_high = (p + 1) / self.prefix_buckets
+            
+                lines.append(f'training_samples_prefix_total{{model="ttft",prefix_bucket="{p}",range="{prefix_low:.2f}-{prefix_high:.2f}"}} {ttft_total}')
+                lines.append(f'training_samples_prefix_total{{model="tpot",prefix_bucket="{p}",range="{prefix_low:.2f}-{prefix_high:.2f}"}} {tpot_total}')
+
+            # Add prefix score distribution statistics
+            all_ttft_samples = self._all_samples(self.ttft_data_buckets)
+            if all_ttft_samples:
+                prefix_scores = [s['prefix_cache_score'] for s in all_ttft_samples]
+                ttfts = [s['actual_ttft_ms'] for s in all_ttft_samples]
+            
+                lines.append(f'prefix_score_mean{{}} {np.mean(prefix_scores):.4f}')
+                lines.append(f'prefix_score_std{{}} {np.std(prefix_scores):.4f}')
+                lines.append(f'prefix_score_min{{}} {np.min(prefix_scores):.4f}')
+                lines.append(f'prefix_score_max{{}} {np.max(prefix_scores):.4f}')
+            
+                # Average TTFT by prefix bucket
+                for p in range(self.prefix_buckets):
+                    prefix_low = p / self.prefix_buckets
+                    prefix_high = (p + 1) / self.prefix_buckets
+                
+                    if p == self.prefix_buckets - 1:
+                        mask = [(prefix_low <= score <= prefix_high) for score in prefix_scores]  # include 1.0
+                    else:
+                        mask = [(prefix_low <= score <  prefix_high) for score in prefix_scores]
+                    bucket_ttfts = [t for t, m in zip(ttfts, mask) if m]
+                
+                    if bucket_ttfts:
+                        avg_ttft = np.mean(bucket_ttfts)
+                        median_ttft = np.median(bucket_ttfts)
+                        lines.append(f'avg_ttft_by_prefix{{prefix_bucket="{p}",range="{prefix_low:.2f}-{prefix_high:.2f}"}} {avg_ttft:.2f}')
+                        lines.append(f'median_ttft_by_prefix{{prefix_bucket="{p}",range="{prefix_low:.2f}-{prefix_high:.2f}"}} {median_ttft:.2f}')
+
+            # 4) Quantile Loss scores (last up to 5)
+            for idx, score in enumerate(self.ttft_quantile_loss_scores):
+                lines.append(f'ttft_quantile_loss{{idx="{idx}"}} {score:.6f}')
+            for idx, score in enumerate(self.tpot_quantile_loss_scores):
+                lines.append(f'tpot_quantile_loss{{idx="{idx}"}} {score:.6f}')
+
+            # 5) Coverage scores (should be close to quantile * 100)
+            for idx, coverage in enumerate(self.ttft_coverage_scores):
+                lines.append(f'ttft_coverage_percent{{idx="{idx}"}} {coverage:.6f}')
+            for idx, coverage in enumerate(self.tpot_coverage_scores):
+                lines.append(f'tpot_coverage_percent{{idx="{idx}"}} {coverage:.6f}')
+
+            # 6) Violation rates (should be close to (1-quantile) * 100)
+            for idx, violation_rate in enumerate(self.ttft_violation_rates):
+                lines.append(f'ttft_violation_rate_percent{{idx="{idx}"}} {violation_rate:.6f}')
+            for idx, violation_rate in enumerate(self.tpot_violation_rates):
+                lines.append(f'tpot_violation_rate_percent{{idx="{idx}"}} {violation_rate:.6f}')
+
+            # 7) Target metrics for reference
+            target_coverage = self.quantile * 100
+            target_violation_rate = (1 - self.quantile) * 100
+            lines.append(f'target_coverage_percent{{}} {target_coverage:.1f}')
+            lines.append(f'target_violation_rate_percent{{}} {target_violation_rate:.1f}')
+
+            return "\n".join(lines) + "\n"
+
+        except Exception as e:
+            logging.error(f"Error generating metrics: {e}", exc_info=True)
+            return "# error_generating_metrics 1\n"
+                
+
+# --- FastAPI Application ---
+app = FastAPI(
+    title="Latency Predictor Service",
+    description="A service to predict TTFT and TPOT using quantile regression with continuous training and feature scaling.",
+)
+
+predictor = LatencyPredictor()
+
+# --- Pydantic Models for API ---
+class TrainingEntry(BaseModel):
+    kv_cache_percentage: float = Field(..., ge=0.0, le=1.0)
+    input_token_length: int = Field(..., ge=0)
+    num_request_waiting: int = Field(..., ge=0)
+    num_request_running: int = Field(..., ge=0)
+    actual_ttft_ms: float = Field(..., ge=0.0)
+    actual_tpot_ms: float = Field(..., ge=0.0)
+    num_tokens_generated: int = Field(..., ge=0)
+    prefix_cache_score: float = Field(..., ge=0.0, le=1.0, description="Prefix cache hit ratio score (0.0 to 1.0)")
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+class PredictionRequest(BaseModel):
+    kv_cache_percentage: float = Field(..., ge=0.0, le=1.0)
+    input_token_length: int = Field(..., ge=0)
+    num_request_waiting: int = Field(..., ge=0)
+    num_request_running: int = Field(..., ge=0)
+    num_tokens_generated: int = Field(..., ge=0)
+    prefix_cache_score: float = Field(..., ge=0.0, le=1.0, description="Prefix cache hit ratio score (0.0 to 1.0)")
+
+class PredictionResponse(BaseModel):
+    ttft_ms: float = Field(..., description=f"Predicted {settings.QUANTILE_ALPHA:.0%} quantile TTFT in milliseconds")
+    tpot_ms: float = Field(..., description=f"Predicted {settings.QUANTILE_ALPHA:.0%} quantile TPOT in milliseconds")
+    ttft_uncertainty: float = Field(..., description="Uncertainty estimate for TTFT prediction")
+    tpot_uncertainty: float = Field(..., description="Uncertainty estimate for TPOT prediction")
+    ttft_prediction_bounds: Tuple[float, float] = Field(..., description="Approximate prediction bounds for TTFT")
+    tpot_prediction_bounds: Tuple[float, float] = Field(..., description="Approximate prediction bounds for TPOT")
+    predicted_at: datetime
+    model_type: ModelType = Field(default=predictor.model_type.value, description="Type of model used for prediction")
+    quantile: float = Field(default=settings.QUANTILE_ALPHA, description="Quantile being predicted")
+    
+class BulkTrainingRequest(BaseModel):
+    entries: List[TrainingEntry]
+
+
+
+# --- Background Training Loop ---
+def continuous_training_loop():
+    time.sleep(10)
+    while not predictor._shutdown_event.is_set():
+        try:
+            logging.debug("Checking if training should run...")
+            predictor.train()
+        except Exception:
+            logging.error("Error in periodic retraining", exc_info=True)
+        if predictor._shutdown_event.wait(timeout=settings.RETRAINING_INTERVAL_SEC):
+            break
+    logging.info("Training loop exiting.")
+
+# --- FastAPI Events ---
+@app.on_event("startup")
+async def startup_event():
+    logging.info("Server starting up...")
+    predictor.load_models()
+    t = threading.Thread(target=continuous_training_loop, daemon=True)
+    predictor._training_thread = t
+    t.start()
+    logging.info("Background training started.")
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    logging.info("Server shutting down...")
+    predictor.shutdown()
+    
+
+@app.post("/add_training_data_bulk", status_code=status.HTTP_202_ACCEPTED)
+async def add_training_data_bulk(batch: BulkTrainingRequest):
+     """
+     Accepts a JSON body like:
+       { "entries": [ { …TrainingEntry… }, { … }, … ] }
+     """
+     try:
+        predictor.add_training_samples([e.dict() for e in batch.entries])
+        return {"message": f"Accepted {len(batch.entries)} training samples."}
+     except Exception:
+         logging.error("Failed to add bulk training data", exc_info=True)
+         raise HTTPException(status_code=500, detail="Failed to add training data in bulk")
+
+@app.post("/predict", response_model=PredictionResponse)
+async def predict_endpoint(request: PredictionRequest):
+    try:
+        ttft_pred, tpot_pred, ttft_std, tpot_std = predictor.predict(request.dict())
+        ttft_pred = max(0, ttft_pred)
+        tpot_pred = max(0, tpot_pred)
+        ttft_bounds = (max(0, ttft_pred - 2*ttft_std), ttft_pred + 2*ttft_std)
+        tpot_bounds = (max(0, tpot_pred - 2*tpot_std), tpot_pred + 2*tpot_std)
+        return PredictionResponse(
+            ttft_ms=ttft_pred,
+            tpot_ms=tpot_pred,
+            ttft_uncertainty=ttft_std,
+            tpot_uncertainty=tpot_std,
+            ttft_prediction_bounds=ttft_bounds,
+            tpot_prediction_bounds=tpot_bounds,
+            predicted_at=datetime.now(timezone.utc),
+            model_type=predictor.model_type.value,
+            quantile=predictor.quantile
+        )
+    except HTTPException:
+        raise
+    except Exception:
+        logging.error("Prediction failed", exc_info=True)
+        raise HTTPException(status_code=500, detail="An internal error occurred during prediction.")
+
+
+
+@app.get("/healthz", status_code=status.HTTP_200_OK)
+async def health_check():
+    return {"status": "ok"}
+
+@app.get("/readyz", status_code=status.HTTP_200_OK)
+async def readiness_check():
+    if not predictor.is_ready:
+        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Models are not ready.")
+    return {"status": "ready"}
+
+
+@app.get("/metrics", status_code=status.HTTP_200_OK)
+async def metrics():
+    """Prometheus metrics including coefficients/importances, bucket counts, and quantile-specific metrics."""
+    try:
+        content = predictor.get_metrics()
+        return Response(content, media_type="text/plain; version=0.0.4")
+    except Exception as e:
+        logging.error(f"Error in metrics endpoint: {e}", exc_info=True)
+        return Response("# Error generating metrics\n", media_type="text/plain; version=0.0.4")
+    
+@app.get("/", include_in_schema=False)
+async def root():
+    return {
+        "message": "Latency Predictor is running.",
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile,
+        "description": f"Predicting {predictor.quantile:.0%} quantile for TTFT and TPOT latencies"
+    }
+
+@app.post("/flush", response_model=FlushResponse, status_code=status.HTTP_200_OK)
+async def flush_data(request: FlushRequest = FlushRequest()):
+    """
+    Manually flush training data, test data, and/or metrics.
+    
+    Useful when:
+    - Server workload has changed significantly
+    - You want to start fresh with new data
+    - Testing or debugging model behavior
+    - Forcing a clean state after deployment
+    
+    Example requests:
+    - Flush everything: POST /flush with empty body
+    - Flush only training: POST /flush with {"flush_test_data": false, "flush_metrics": false}
+    - Flush with reason: POST /flush with {"reason": "New deployment"}
+    """
+    try:
+        result = predictor.flush_training_data(
+            flush_training=request.flush_training_data,
+            flush_test=request.flush_test_data,
+            flush_metrics=request.flush_metrics,
+            reason=request.reason
+        )
+        
+        total_flushed = (
+            result["ttft_training_samples_flushed"] + 
+            result["tpot_training_samples_flushed"] +
+            result["ttft_test_samples_flushed"] + 
+            result["tpot_test_samples_flushed"]
+        )
+        
+        message_parts = []
+        if request.flush_training_data:
+            message_parts.append(
+                f"{result['ttft_training_samples_flushed']} TTFT and "
+                f"{result['tpot_training_samples_flushed']} TPOT training samples"
+            )
+        if request.flush_test_data:
+            message_parts.append(
+                f"{result['ttft_test_samples_flushed']} TTFT and "
+                f"{result['tpot_test_samples_flushed']} TPOT test samples"
+            )
+        if request.flush_metrics:
+            message_parts.append("all metric scores")
+        
+        message = f"Successfully flushed: {', '.join(message_parts)}" if message_parts else "No data flushed"
+        
+        return FlushResponse(
+            success=True,
+            flushed_at=datetime.now(timezone.utc),
+            reason=request.reason,
+            ttft_training_samples_flushed=result["ttft_training_samples_flushed"],
+            tpot_training_samples_flushed=result["tpot_training_samples_flushed"],
+            ttft_test_samples_flushed=result["ttft_test_samples_flushed"],
+            tpot_test_samples_flushed=result["tpot_test_samples_flushed"],
+            metrics_cleared=result["metrics_cleared"],
+            message=message
+        )
+        
+    except Exception as e:
+        logging.error(f"Error in flush endpoint: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500, 
+            detail=f"Failed to flush data: {str(e)}"
+        )
+
+
+@app.get("/data/status", status_code=status.HTTP_200_OK)
+async def get_data_status():
+    """
+    Get current status of training data.
+    Useful for monitoring and deciding whether to flush.
+    """
+    ttft_training_count = sum(len(bucket) for bucket in predictor.ttft_data_buckets.values())
+    tpot_training_count = sum(len(bucket) for bucket in predictor.tpot_data_buckets.values())
+    
+    bucket_distribution = {}
+    for (q, c, p), bucket in predictor.ttft_data_buckets.items():
+        if len(bucket) > 0:
+            key = f"queue_{q}_cache_{c}_prefix_{p}"
+            bucket_distribution[key] = len(bucket)
+    
+    return {
+        "training_data": {
+            "ttft_samples": ttft_training_count,
+            "tpot_samples": tpot_training_count,
+            "total_samples": ttft_training_count + tpot_training_count
+        },
+        "test_data": {
+            "ttft_samples": len(predictor.ttft_test_data),
+            "tpot_samples": len(predictor.tpot_test_data),
+            "total_samples": len(predictor.ttft_test_data) + len(predictor.tpot_test_data)
+        },
+        "metrics": {
+            "ttft_scores_count": len(predictor.ttft_quantile_loss_scores),
+            "tpot_scores_count": len(predictor.tpot_quantile_loss_scores)
+        },
+        "bucket_distribution": bucket_distribution,
+        "model_ready": predictor.is_ready,
+        "last_retrain": predictor.last_retrain_time.isoformat() if predictor.last_retrain_time else None
+    }
+    
+@app.get("/model/download/info")
+async def model_download_info():
+    """
+    Get information about available model downloads and coefficients.
+    """
+    info = {
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile,
+        "available_endpoints": {}
+    }
+    
+    if predictor.model_type == ModelType.BAYESIAN_RIDGE:
+        info["available_endpoints"]["coefficients"] = "/metrics"
+        info["coefficients_info"] = {
+            "ttft_coefficients_available": predictor.ttft_coefficients is not None,
+            "tpot_coefficients_available": predictor.tpot_coefficients is not None,
+            "description": "Descaled coefficients available in Prometheus metrics endpoint"
+        }
+    elif predictor.model_type == ModelType.XGBOOST:
+        info["available_endpoints"]["trees"] = {
+            "ttft_trees": "/model/ttft/xgb/json",
+            "tpot_trees": "/model/tpot/xgb/json"
+        }
+    else: 
+        info["available_endpoints"]["lightgbm"] = {
+            "ttft_model_txt": "/model/ttft/lgb/txt",
+            "tpot_model_txt": "/model/tpot/lgb/txt",
+            "ttft_importances": "/model/ttft/lgb/importances",
+            "tpot_importances": "/model/tpot/lgb/importances"
+        }
+    
+    info["model_status"] = {
+        "ttft_model_ready": predictor.ttft_model is not None,
+        "tpot_model_ready": predictor.tpot_model is not None,
+    }
+    
+    if predictor.model_type == ModelType.BAYESIAN_RIDGE:
+        info["model_status"]["ttft_coefficients_ready"] = predictor.ttft_coefficients is not None
+        info["model_status"]["tpot_coefficients_ready"] = predictor.tpot_coefficients is not None
+    
+    # Add quantile-specific evaluation info
+    info["evaluation_info"] = {
+        "quantile_loss": "Pinball loss for quantile regression evaluation",
+        "coverage_percent": f"Percentage of actual values below predicted {predictor.quantile:.0%} quantile (target: {predictor.quantile*100:.1f}%)",
+        "violation_rate_percent": f"Percentage of actual values above predicted {predictor.quantile:.0%} quantile (target: {(1-predictor.quantile)*100:.1f}%)"
+    }
+    
+    return info
+
+@app.get("/model/ttft/xgb/json")
+async def ttft_xgb_json():
+    """
+    Dump the TTFT XGBoost model as JSON trees.
+    """
+    if predictor.model_type != ModelType.XGBOOST:
+        raise HTTPException(status_code=404, detail="TTFT model is not XGBoost")
+    
+    if not predictor.ttft_model:
+        raise HTTPException(status_code=404, detail="TTFT model not available")
+        
+    try:
+        booster = predictor.ttft_model.get_booster()
+        # get_dump with dump_format="json" gives one JSON string per tree
+        raw_trees = booster.get_dump(dump_format="json")
+        # parse each string into a dict so the response is a JSON array of objects
+        trees = [json.loads(t) for t in raw_trees]
+        return JSONResponse(content=trees)
+    except Exception as e:
+        logging.error(f"Error dumping TTFT XGBoost trees: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error dumping TTFT XGBoost trees")
+
+
+@app.get("/model/tpot/xgb/json")
+async def tpot_xgb_json():
+    """
+    Dump the TPOT XGBoost model as JSON trees.
+    """
+    if predictor.model_type != ModelType.XGBOOST:
+        raise HTTPException(status_code=404, detail="TPOT model is not XGBoost")
+    
+    if not predictor.tpot_model:
+        raise HTTPException(status_code=404, detail="TPOT model not available")
+        
+    try:
+        booster = predictor.tpot_model.get_booster()
+        raw_trees = booster.get_dump(dump_format="json")
+        trees = [json.loads(t) for t in raw_trees]
+        return JSONResponse(content=trees)
+    except Exception as e:
+        logging.error(f"Error dumping TPOT XGBoost trees: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Error dumping TPOT XGBoost trees")
+
+
+
+@app.get("/model/{model_name}/info")
+async def model_info(model_name: str):
+    """Get model file information including last modified time."""
+    model_paths = {
+        "ttft": settings.TTFT_MODEL_PATH,
+        "tpot": settings.TPOT_MODEL_PATH,
+        "ttft_scaler": settings.TTFT_SCALER_PATH,
+        "tpot_scaler": settings.TPOT_SCALER_PATH
+    }
+    
+    if model_name not in model_paths:
+        raise HTTPException(status_code=404, detail=f"Unknown model: {model_name}")
+    
+    model_path = model_paths[model_name]
+    
+    if not os.path.exists(model_path):
+        raise HTTPException(status_code=404, detail=f"Model {model_name} not found")
+    
+    # Get file stats
+    stat = os.stat(model_path)
+    last_modified = datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc)
+    
+    return {
+        "model_name": model_name,
+        "path": model_path,
+        "size_bytes": stat.st_size,
+        "last_modified": last_modified.isoformat(),
+        "exists": True,
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile if model_name in ["ttft", "tpot"] else None
+    }
+
+
+@app.get("/model/{model_name}/download")
+async def download_model(model_name: str):
+    """Download a model file."""
+    model_paths = {
+        "ttft": settings.TTFT_MODEL_PATH,
+        "tpot": settings.TPOT_MODEL_PATH,
+        "ttft_scaler": settings.TTFT_SCALER_PATH,
+        "tpot_scaler": settings.TPOT_SCALER_PATH
+    }
+    
+    if model_name not in model_paths:
+        raise HTTPException(status_code=404, detail=f"Unknown model: {model_name}")
+    
+    model_path = model_paths[model_name]
+    
+    if not os.path.exists(model_path):
+        raise HTTPException(status_code=404, detail=f"Model {model_name} not found")
+    
+    # Return the file
+    filename = f"{model_name}.joblib"
+    return FileResponse(
+        model_path,
+        media_type='application/octet-stream',
+        filename=filename
+    )
+
+
+@app.get("/models/list")
+async def list_models():
+    """List all available models with their status."""
+    models = {}
+    model_paths = {
+        "ttft": settings.TTFT_MODEL_PATH,
+        "tpot": settings.TPOT_MODEL_PATH,
+        "ttft_scaler": settings.TTFT_SCALER_PATH,
+        "tpot_scaler": settings.TPOT_SCALER_PATH
+    }
+    
+    for model_name, model_path in model_paths.items():
+        if os.path.exists(model_path):
+            stat = os.stat(model_path)
+            models[model_name] = {
+                "exists": True,
+                "size_bytes": stat.st_size,
+                "last_modified": datetime.fromtimestamp(stat.st_mtime, tz=timezone.utc).isoformat()
+            }
+        else:
+            models[model_name] = {
+                "exists": False,
+                "size_bytes": 0,
+                "last_modified": None
+            }
+    
+    return {
+        "models": models,
+        "model_type": predictor.model_type.value,
+        "quantile": predictor.quantile,
+        "server_time": datetime.now(timezone.utc).isoformat(),
+        "evaluation_metrics": {
+            "quantile_loss": "Lower is better",
+            "coverage_percent": f"Target: {predictor.quantile*100:.1f}%",
+            "violation_rate_percent": f"Target: {(1-predictor.quantile)*100:.1f}%"
+        }
+    }
+
+# Add new API endpoints for LightGBM model exports
+@app.get("/model/ttft/lgb/txt")
+async def ttft_lgb_txt():
+    """
+    Download the TTFT LightGBM model as text format.
+    """
+    if predictor.model_type != ModelType.LIGHTGBM:
+        raise HTTPException(status_code=404, detail="TTFT model is not LightGBM")
+    
+    if not predictor.ttft_model:
+        raise HTTPException(status_code=404, detail="TTFT model not available")
+        
+    txt_path = settings.TTFT_MODEL_PATH.replace('.joblib', '_lgb.txt')
+    if not os.path.exists(txt_path):
+        raise HTTPException(status_code=404, detail="TTFT LightGBM text model not found")
+    
+    return FileResponse(
+        txt_path,
+        media_type='text/plain',
+        filename='ttft_lgb_model.txt'
+    )
+
+@app.get("/model/tpot/lgb/txt")
+async def tpot_lgb_txt():
+    """
+    Download the TPOT LightGBM model as text format.
+    """
+    if predictor.model_type != ModelType.LIGHTGBM:
+        raise HTTPException(status_code=404, detail="TPOT model is not LightGBM")
+    
+    if not predictor.tpot_model:
+        raise HTTPException(status_code=404, detail="TPOT model not available")
+        
+    txt_path = settings.TPOT_MODEL_PATH.replace('.joblib', '_lgb.txt')
+    if not os.path.exists(txt_path):
+        raise HTTPException(status_code=404, detail="TPOT LightGBM text model not found")
+    
+    return FileResponse(
+        txt_path,
+        media_type='text/plain',
+        filename='tpot_lgb_model.txt'
+    )
+
+@app.get("/model/ttft/lgb/importances")
+async def ttft_lgb_importances():
+    """
+    Get TTFT LightGBM feature importances as JSON.
+    """
+    if predictor.model_type != ModelType.LIGHTGBM:
+        raise HTTPException(status_code=404, detail="TTFT model is not LightGBM")
+    
+    if not predictor.ttft_model:
+        raise HTTPException(status_code=404, detail="TTFT model not available")
+        
+    imp_path = settings.TTFT_MODEL_PATH.replace('.joblib', '_importances.json')
+    if not os.path.exists(imp_path):
+        raise HTTPException(status_code=404, detail="TTFT LightGBM importances not found")
+    
+    with open(imp_path, 'r') as f:
+        importances = json.load(f)
+    
+    return JSONResponse(content=importances)
+
+@app.get("/model/tpot/lgb/importances")
+async def tpot_lgb_importances():
+    """
+    Get TPOT LightGBM feature importances as JSON.
+    """
+    if predictor.model_type != ModelType.LIGHTGBM:
+        raise HTTPException(status_code=404, detail="TPOT model is not LightGBM")
+    
+    if not predictor.tpot_model:
+        raise HTTPException(status_code=404, detail="TPOT model not available")
+        
+    imp_path = settings.TPOT_MODEL_PATH.replace('.joblib', '_importances.json')
+    if not os.path.exists(imp_path):
+        raise HTTPException(status_code=404, detail="TPOT LightGBM importances not found")
+    
+    with open(imp_path, 'r') as f:
+        importances = json.load(f)
+    
+    return JSONResponse(content=importances)
+
+@app.get("/debug/prefix_distribution")
+async def prefix_distribution():
+    """
+    Debug endpoint to analyze the relationship between prefix_cache_score and TTFT.
+    This helps verify that the model is seeing the data it needs to learn.
+    """
+    all_samples = predictor._all_samples(predictor.ttft_data_buckets)
+    if not all_samples:
+        return {"error": "No training samples available"}
+    
+    prefix_scores = [s['prefix_cache_score'] for s in all_samples]
+    ttfts = [s['actual_ttft_ms'] for s in all_samples]
+    
+    # Group by prefix score ranges
+    ranges = [(0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
+    distribution = {}
+
+    for low, high in ranges:
+        # include the right edge only for the final bin so 1.0 is counted
+        if high == 1.0:
+            mask = [(low <= p <= high) for p in prefix_scores]
+        else:
+            mask = [(low <= p <  high) for p in prefix_scores]
+
+        range_ttfts  = [t for t, m in zip(ttfts, mask) if m]
+        range_prefix = [p for p, m in zip(prefix_scores, mask) if m]
+        distribution[f"{low}-{high}"] = {
+            "count": len(range_ttfts),
+            "mean_ttft_ms": float(np.mean(range_ttfts)) if range_ttfts else 0,
+            "median_ttft_ms": float(np.median(range_ttfts)) if range_ttfts else 0,
+            "std_ttft_ms": float(np.std(range_ttfts)) if range_ttfts else 0,
+            "mean_prefix_score": float(np.mean(range_prefix)) if range_prefix else 0,
+            "min_ttft_ms": float(np.min(range_ttfts)) if range_ttfts else 0,
+            "max_ttft_ms": float(np.max(range_ttfts)) if range_ttfts else 0,
+        }
+    
+    # Overall statistics
+    overall = {
+        "total_samples": len(all_samples),
+        "prefix_score_mean": float(np.mean(prefix_scores)),
+        "prefix_score_std": float(np.std(prefix_scores)),
+        "prefix_score_min": float(np.min(prefix_scores)),
+        "prefix_score_max": float(np.max(prefix_scores)),
+        "ttft_mean": float(np.mean(ttfts)),
+        "ttft_std": float(np.std(ttfts)),
+        "correlation": float(np.corrcoef(prefix_scores, ttfts)[0, 1])
+    }
+    
+    return {
+        "overall_stats": overall,
+        "distribution_by_prefix_range": distribution,
+        "interpretation": {
+            "correlation": "Negative correlation means higher prefix score → lower TTFT (good!)",
+            "check_distribution": "All ranges should have samples. Empty ranges mean missing data.",
+            "expected_pattern": "TTFT should decrease significantly as prefix score increases"
+        }
+    }
+
+if __name__ == "__main__":
+    uvicorn.run("__main__:app", host="0.0.0.0", port=8000, reload=True)
\ No newline at end of file
diff --git a/sidecars/latencypredictorasync/latencypredictor_async.go b/sidecars/latencypredictorasync/latencypredictor_async.go
new file mode 100644
index 0000000..32a41f3
--- /dev/null
+++ b/sidecars/latencypredictorasync/latencypredictor_async.go
@@ -0,0 +1,232 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math/rand"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/go-logr/logr"
+)
+
+// --- Predictor Client ---
+
+type Predictor struct {
+	config     *Config
+	httpClient *http.Client
+	logger     logr.Logger
+	rng        *rand.Rand
+
+	metricsMu     sync.RWMutex
+	cachedMetrics *MetricsResponse
+	modelInfo     *ModelInfo
+	serverStatus  *ServerStatusResponse
+
+	bufferMu sync.Mutex
+	pending  []TrainingEntry
+
+	wg   sync.WaitGroup
+	done chan struct{}
+}
+
+func New(config *Config, logger logr.Logger) *Predictor {
+	if config == nil {
+		config = ConfigFromEnv()
+	}
+	p := &Predictor{
+		config:     config,
+		httpClient: &http.Client{Timeout: config.HTTPTimeout},
+		logger:     logger.WithName("latency-predictor-client"),
+		rng:        rand.New(rand.NewSource(time.Now().UnixNano())),
+		done:       make(chan struct{}),
+	}
+	p.wg.Add(1)
+	go p.backgroundLoop()
+	return p
+}
+
+// Start initializes the predictor by fetching server status and model info.
+func (p *Predictor) Start(ctx context.Context) error {
+	// Get initial server status
+	if err := p.refreshServerStatus(ctx); err != nil {
+		return fmt.Errorf("failed to get initial server status: %v", err)
+	}
+
+	// Get initial model info if training server is available
+	if err := p.refreshModelInfo(ctx); err != nil {
+		return fmt.Errorf("failed to get initial model info: %v", err)
+	}
+
+	p.logger.Info("Latency predictor async client started.",
+		"training_url", p.config.TrainingURL,
+		"prediction_urls", p.config.PredictionURLs,
+		"max_sample_size", p.config.MaxSampleSize,
+		"flush_interval", p.config.FlushInterval,
+		"use_native_xgboost", p.config.UseNativeXGBoost,
+		"max_bulk_size", p.config.MaxBulkSize)
+	return nil
+}
+
+// Stop stops background work, then does a final flush/refresh.
+func (p *Predictor) Stop() {
+	close(p.done)
+	p.wg.Wait() // Wait for the background loop to finish
+	// final flush & refresh
+	p.flushTraining()
+	p.refreshMetrics()
+	p.logger.Info("Latency predictor async client stopped.")
+}
+
+// backgroundLoop runs flush & refresh at configured intervals.
+func (p *Predictor) backgroundLoop() {
+	defer p.wg.Done()
+	flushTicker := time.NewTicker(p.config.FlushInterval)
+	metricsTicker := time.NewTicker(p.config.MetricsRefreshInterval)
+	defer flushTicker.Stop()
+	defer metricsTicker.Stop()
+
+	for {
+		select {
+		case <-flushTicker.C:
+			p.flushTraining()
+		case <-metricsTicker.C:
+			p.refreshMetrics()
+			// Also refresh server status periodically
+			ctx, cancel := context.WithTimeout(context.Background(), p.config.HTTPTimeout)
+			if err := p.refreshServerStatus(ctx); err != nil {
+				p.logger.Error(err, "failed to refresh server status during background refresh")
+			}
+			cancel()
+		case <-p.done:
+			return
+		}
+	}
+}
+
+// GetXGBoostTrees returns the cached XGBoost tree data. It does not fetch new data.
+func (p *Predictor) GetXGBoostTrees(ctx context.Context) (*XGBoostTrees, error) {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+	if p.cachedMetrics == nil || p.cachedMetrics.XGBoostTrees == nil {
+		return nil, errors.New("no cached XGBoost trees available")
+	}
+	return p.cachedMetrics.XGBoostTrees, nil
+}
+
+// GetModelInfo fetches the latest model info from the training server.
+func (p *Predictor) GetModelInfo(ctx context.Context) (*ModelInfo, error) {
+	if err := p.refreshModelInfo(ctx); err != nil {
+		return nil, err
+	}
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+
+	return p.modelInfo, nil
+}
+
+// GetCachedMetrics returns the last metrics fetched. The bool indicates if a value is cached.
+func (p *Predictor) GetCachedMetrics() (*MetricsResponse, bool) {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+	if p.cachedMetrics == nil {
+		return nil, false
+	}
+	return p.cachedMetrics, true
+}
+
+// IsXGBoostReady returns true if native XGBoost models are loaded and ready.
+func (p *Predictor) IsXGBoostReady() bool {
+	return p.modelInfo != nil && p.modelInfo.ModelType == xgBoostModelType
+}
+
+// IsLightGBMReady returns true if LightGBM models are available via HTTP.
+func (p *Predictor) IsLightGBMReady() bool {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+	return p.modelInfo != nil && p.modelInfo.ModelType == gbmModelType && len(p.config.PredictionURLs) > 0
+}
+
+// IsBayesianRidgeReady returns true if Bayesian Ridge coefficients are cached.
+func (p *Predictor) IsBayesianRidgeReady() bool {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+	return p.cachedMetrics != nil && p.cachedMetrics.Coefficients != nil
+}
+
+// GetCurrentModelType returns the current model type from cached server status or model info.
+func (p *Predictor) GetCurrentModelType() string {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+
+	// Prefer server status if available
+	if p.serverStatus != nil {
+		return p.serverStatus.ModelType
+	}
+
+	if p.modelInfo == nil {
+		return ""
+	}
+	return p.modelInfo.ModelType
+}
+
+// GetCurrentQuantile returns the current quantile from server status or defaults to 0.9
+func (p *Predictor) GetCurrentQuantile() float64 {
+	p.metricsMu.RLock()
+	defer p.metricsMu.RUnlock()
+
+	// Prefer server status if available
+	if p.serverStatus != nil && p.serverStatus.Quantile > 0 {
+		return p.serverStatus.Quantile
+	}
+
+	if p.modelInfo != nil && p.modelInfo.Quantile > 0 {
+		return p.modelInfo.Quantile
+	}
+
+	return 0.9 // Default quantile
+}
+
+// IsReady returns true if a prediction method is ready based on the current model type.
+func (p *Predictor) IsReady() bool {
+	switch p.GetCurrentModelType() {
+	case bayesianRidgeModelType:
+		return p.IsBayesianRidgeReady()
+	case xgBoostModelType:
+		// Ready if we have prediction URLs for HTTP calls
+		return len(p.config.PredictionURLs) > 0
+	case gbmModelType:
+		// Ready if we have prediction URLs for HTTP calls
+		return p.IsLightGBMReady()
+	default:
+		return false
+	}
+}
+
+// GetPredictionURLs returns the list of configured prediction URLs for debugging/monitoring.
+func (p *Predictor) GetPredictionURLs() []string {
+	return p.config.PredictionURLs
+}
+
+// GetTrainingURL returns the configured training URL for debugging/monitoring.
+func (p *Predictor) GetTrainingURL() string {
+	return p.config.TrainingURL
+}
diff --git a/sidecars/latencypredictorasync/metrics.go b/sidecars/latencypredictorasync/metrics.go
new file mode 100644
index 0000000..7c0f2bb
--- /dev/null
+++ b/sidecars/latencypredictorasync/metrics.go
@@ -0,0 +1,193 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strconv"
+	"strings"
+
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+// GetMetrics fetches & parses metrics from the training server (for Bayesian Ridge).
+func (p *Predictor) GetMetrics(ctx context.Context) (*MetricsResponse, error) {
+	url := p.config.TrainingURL + "/metrics"
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create metrics request: %w", err)
+	}
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call training server /metrics endpoint: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("training server returned non-200 status: %d %s, body: %s", resp.StatusCode, resp.Status, string(body))
+	}
+
+	rawMetricsBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read metrics response body: %w", err)
+	}
+	rawMetrics := string(rawMetricsBytes)
+
+	metricsResponse := &MetricsResponse{
+		RawMetrics: rawMetrics,
+		ModelType:  bayesianRidgeModelType, // Assume Bayesian Ridge when calling /metrics
+	}
+
+	coeffs, buckets, err := p.parsePrometheusMetrics(rawMetrics)
+	if err != nil {
+		p.logger.Error(err, "Failed to parse Prometheus metrics, caching raw only")
+	} else {
+		metricsResponse.Coefficients = coeffs
+		metricsResponse.BucketCounts = buckets
+	}
+
+	p.metricsMu.Lock()
+	p.cachedMetrics = metricsResponse
+	p.metricsMu.Unlock()
+
+	p.logger.V(logutil.DEBUG).Info("Successfully retrieved and cached Bayesian Ridge metrics.")
+	return metricsResponse, nil
+}
+
+// parsePrometheusMetrics parses the Prometheus-format metrics into structured data.
+func (p *Predictor) parsePrometheusMetrics(rawMetrics string) (*ModelCoefficients, *BucketCounts, error) {
+	lines := strings.Split(rawMetrics, "\n")
+
+	coefficients := &ModelCoefficients{
+		TTFTCoeffs: make(map[string]float64),
+		TPOTCoeffs: make(map[string]float64),
+	}
+	bucketCounts := &BucketCounts{
+		TTFTBuckets: make(map[int]int),
+		TPOTBuckets: make(map[int]int),
+	}
+	var firstErr error
+
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+
+		if err := p.parseMetricLine(line, coefficients, bucketCounts); err != nil {
+			if firstErr == nil {
+				firstErr = err // Save first error to return
+			}
+			p.logger.V(logutil.TRACE).Info("Skipping unparseable metric line", "line", line, "error", err)
+		}
+	}
+	return coefficients, bucketCounts, firstErr
+}
+
+// parseMetricLine parses a single line of Prometheus-formatted text.
+func (p *Predictor) parseMetricLine(line string, coefficients *ModelCoefficients, bucketCounts *BucketCounts) error {
+	lastSpaceIdx := strings.LastIndexAny(line, " \t")
+	if lastSpaceIdx == -1 {
+		return errors.New("invalid metric format: no space found")
+	}
+
+	metricPart := strings.TrimSpace(line[:lastSpaceIdx])
+	valueStr := strings.TrimSpace(line[lastSpaceIdx+1:])
+
+	value, err := strconv.ParseFloat(valueStr, 64)
+	if err != nil {
+		return fmt.Errorf("could not parse value '%s': %w", valueStr, err)
+	}
+
+	metricName := metricPart
+	if openBrace := strings.Index(metricPart, "{"); openBrace != -1 {
+		metricName = metricPart[:openBrace]
+	}
+
+	switch metricName {
+	case "ttft_intercept":
+		coefficients.TTFTIntercept = value
+	case "tpot_intercept":
+		coefficients.TPOTIntercept = value
+	case "ttft_coef":
+		if feature := p.extractLabel(metricPart, "feature"); feature != "" {
+			coefficients.TTFTCoeffs[feature] = value
+		}
+	case "tpot_coef":
+		if feature := p.extractLabel(metricPart, "feature"); feature != "" {
+			coefficients.TPOTCoeffs[feature] = value
+		}
+	case "training_samples_count":
+		model := p.extractLabel(metricPart, "model")
+		bucketStr := p.extractLabel(metricPart, "bucket")
+		if bucket, err := strconv.Atoi(bucketStr); err == nil {
+			switch model {
+			case "ttft":
+				bucketCounts.TTFTBuckets[bucket] = int(value)
+			case "tpot":
+				bucketCounts.TPOTBuckets[bucket] = int(value)
+			}
+		}
+	}
+	return nil
+}
+
+// extractLabel extracts a label value from a Prometheus metric string.
+// Example: `metric{key="value"}`, `key` -> `"value"`
+func (p *Predictor) extractLabel(metricPart, labelName string) string {
+	searchStr := labelName + `="`
+	start := strings.Index(metricPart, searchStr)
+	if start == -1 {
+		return ""
+	}
+	start += len(searchStr)
+	end := strings.Index(metricPart[start:], `"`)
+	if end == -1 {
+		return ""
+	}
+	return metricPart[start : start+end]
+}
+
+// GetModelCoefficients fetches the latest metrics and returns the parsed coefficients.
+func (p *Predictor) GetModelCoefficients(ctx context.Context) (*ModelCoefficients, error) {
+	metrics, err := p.GetMetrics(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if metrics.Coefficients == nil {
+		return nil, errors.New("coefficients not available in fetched metrics")
+	}
+	return metrics.Coefficients, nil
+}
+
+// GetBucketCounts fetches the latest metrics and returns the parsed bucket counts.
+func (p *Predictor) GetBucketCounts(ctx context.Context) (*BucketCounts, error) {
+	metrics, err := p.GetMetrics(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if metrics.BucketCounts == nil {
+		return nil, errors.New("bucket counts not available in fetched metrics")
+	}
+	return metrics.BucketCounts, nil
+}
diff --git a/sidecars/latencypredictorasync/prediction.go b/sidecars/latencypredictorasync/prediction.go
new file mode 100644
index 0000000..cf251ec
--- /dev/null
+++ b/sidecars/latencypredictorasync/prediction.go
@@ -0,0 +1,351 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+// Predict uses cached coefficients (Bayesian Ridge) or HTTP calls (XGBoost/LightGBM) for prediction.
+func (p *Predictor) Predict(ctx context.Context, req PredictionRequest) (*PredictionResponse, error) {
+	// Get current model type from server status first, fall back to model info
+	p.metricsMu.RLock()
+	modelType := ""
+	quantile := 0.9 // default
+
+	if p.serverStatus != nil {
+		modelType = p.serverStatus.ModelType
+		quantile = p.serverStatus.Quantile
+	} else if p.modelInfo != nil {
+		modelType = p.modelInfo.ModelType
+		if p.modelInfo.Quantile > 0 {
+			quantile = p.modelInfo.Quantile
+		}
+	}
+
+	mr := p.cachedMetrics
+	p.metricsMu.RUnlock()
+
+	if modelType == "" {
+		return nil, errors.New("model type not yet available from server")
+	}
+
+	switch modelType {
+	case bayesianRidgeModelType:
+		return p.predictBayesianRidge(req, mr, quantile)
+	case xgBoostModelType, gbmModelType:
+		return p.predictHTTP(ctx, req)
+	default:
+		return nil, fmt.Errorf("unsupported or unknown model type: %s", modelType)
+	}
+}
+
+// PredictBulk makes bulk predictions with error handling (allows partial failures)
+func (p *Predictor) PredictBulk(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error) {
+	if len(requests) == 0 {
+		return nil, errors.New("no prediction requests provided")
+	}
+
+	if len(requests) > p.config.MaxBulkSize {
+		return nil, fmt.Errorf("too many requests: %d (max: %d)", len(requests), p.config.MaxBulkSize)
+	}
+
+	// Validate all requests first
+	for i, req := range requests {
+		if err := p.ValidatePredictionRequest(req); err != nil {
+			return nil, fmt.Errorf("validation failed for request %d: %w", i, err)
+		}
+	}
+
+	payload := BulkPredictionRequest{Requests: requests}
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal bulk prediction request: %w", err)
+	}
+
+	predictionURL := p.getRandomPredictionURL()
+	url := predictionURL + "/predict/bulk"
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create bulk prediction request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call bulk prediction endpoint %s: %w", url, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("bulk prediction server returned non-200 status: %d %s, body: %s", resp.StatusCode, resp.Status, string(body))
+	}
+
+	var bulkResp BulkPredictionResponseWithErrors
+	if err := json.NewDecoder(resp.Body).Decode(&bulkResp); err != nil {
+		return nil, fmt.Errorf("failed to decode bulk prediction response: %w", err)
+	}
+
+	// Convert to standard bulk response format
+	var predictions []PredictionResponse
+	for _, pred := range bulkResp.Predictions {
+		if pred != nil {
+			predictions = append(predictions, *pred)
+		}
+	}
+
+	return &BulkPredictionResponse{
+		Predictions:           predictions,
+		TotalRequests:         bulkResp.TotalRequests,
+		SuccessfulPredictions: bulkResp.SuccessfulPredictions,
+		FailedPredictions:     bulkResp.FailedPredictions,
+		ProcessingTimeMs:      bulkResp.ProcessingTimeMs,
+	}, nil
+}
+
+// PredictBulkStrict makes bulk predictions that fail if any single prediction fails
+func (p *Predictor) PredictBulkStrict(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error) {
+	if len(requests) == 0 {
+		return nil, errors.New("no prediction requests provided")
+	}
+
+	if len(requests) > p.config.MaxBulkSize {
+		return nil, fmt.Errorf("too many requests: %d (max: %d)", len(requests), p.config.MaxBulkSize)
+	}
+
+	// Validate all requests first
+	for i, req := range requests {
+		if err := p.ValidatePredictionRequest(req); err != nil {
+			return nil, fmt.Errorf("validation failed for request %d: %w", i, err)
+		}
+	}
+
+	payload := BulkPredictionRequest{Requests: requests}
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal bulk prediction request: %w", err)
+	}
+
+	predictionURL := p.getRandomPredictionURL()
+	url := predictionURL + "/predict/bulk/strict"
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create bulk prediction request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call bulk prediction endpoint %s: %w", url, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("bulk prediction server returned non-200 status: %d %s, body: %s", resp.StatusCode, resp.Status, string(body))
+	}
+
+	var bulkResp BulkPredictionResponse
+	if err := json.NewDecoder(resp.Body).Decode(&bulkResp); err != nil {
+		return nil, fmt.Errorf("failed to decode bulk prediction response: %w", err)
+	}
+
+	return &bulkResp, nil
+}
+
+// predictBayesianRidge uses cached coefficients for linear prediction
+func (p *Predictor) predictBayesianRidge(req PredictionRequest, mr *MetricsResponse, quantile float64) (*PredictionResponse, error) {
+	if mr == nil || mr.Coefficients == nil {
+		return nil, errors.New("no cached Bayesian Ridge coefficients available for prediction")
+	}
+	c := mr.Coefficients
+
+	// Updated linear combination for TTFT to include prefix_cache_score
+	ttft := c.TTFTIntercept +
+		c.TTFTCoeffs["kv_cache_percentage"]*req.KVCachePercentage +
+		c.TTFTCoeffs["input_token_length"]*float64(req.InputTokenLength) +
+		c.TTFTCoeffs["num_request_waiting"]*float64(req.NumRequestWaiting) +
+		c.TTFTCoeffs["num_request_running"]*float64(req.NumRequestRunning) +
+		c.TTFTCoeffs["prefix_cache_score"]*req.PrefixCacheScore
+
+	// Linear combination for TPOT (remains unchanged - no prefix cache effect)
+	tpot := c.TPOTIntercept +
+		c.TPOTCoeffs["kv_cache_percentage"]*req.KVCachePercentage +
+		c.TPOTCoeffs["input_token_length"]*float64(req.InputTokenLength) +
+		c.TPOTCoeffs["num_request_waiting"]*float64(req.NumRequestWaiting) +
+		c.TPOTCoeffs["num_request_running"]*float64(req.NumRequestRunning) +
+		c.TPOTCoeffs["num_tokens_generated"]*float64(req.NumTokensGenerated)
+
+	return &PredictionResponse{
+		TTFT:        ttft,
+		TPOT:        tpot,
+		PredictedAt: time.Now(),
+		ModelType:   "bayesian_ridge",
+		Quantile:    quantile,
+	}, nil
+}
+
+// predictHTTP makes an HTTP call to a randomly selected prediction server for XGBoost/LightGBM predictions
+func (p *Predictor) predictHTTP(ctx context.Context, req PredictionRequest) (*PredictionResponse, error) {
+	data, err := json.Marshal(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal prediction request: %w", err)
+	}
+
+	// Get random prediction URL for load balancing
+	predictionURL := p.getRandomPredictionURL()
+	url := predictionURL + "/predict"
+
+	p.logger.V(logutil.TRACE).Info("Making prediction request", "url", url)
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(data))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+
+	resp, err := p.httpClient.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call prediction endpoint %s: %w", url, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("prediction server returned non-200 status: %d %s, body: %s", resp.StatusCode, resp.Status, string(body))
+	}
+
+	var predResp PredictionResponse
+	if err := json.NewDecoder(resp.Body).Decode(&predResp); err != nil {
+		return nil, fmt.Errorf("failed to decode prediction response: %w", err)
+	}
+
+	return &predResp, nil
+}
+
+// ValidatePredictionRequest validates that a prediction request has all required fields
+// with valid values, including the new prefix_cache_score field.
+func (p *Predictor) ValidatePredictionRequest(req PredictionRequest) error {
+	if req.KVCachePercentage < 0.0 || req.KVCachePercentage > 1.0 {
+		return fmt.Errorf("kv_cache_percentage must be between 0.0 and 1.0, got %f", req.KVCachePercentage)
+	}
+	if req.InputTokenLength < 0 {
+		return fmt.Errorf("input_token_length must be non-negative, got %d", req.InputTokenLength)
+	}
+	if req.NumRequestWaiting < 0 {
+		return fmt.Errorf("num_request_waiting must be non-negative, got %d", req.NumRequestWaiting)
+	}
+	if req.NumRequestRunning < 0 {
+		return fmt.Errorf("num_request_running must be non-negative, got %d", req.NumRequestRunning)
+	}
+	if req.NumTokensGenerated < 0 {
+		return fmt.Errorf("num_tokens_generated must be non-negative, got %d", req.NumTokensGenerated)
+	}
+	if req.PrefixCacheScore < 0.0 || req.PrefixCacheScore > 1.0 {
+		return fmt.Errorf("prefix_cache_score must be between 0.0 and 1.0, got %f", req.PrefixCacheScore)
+	}
+	return nil
+}
+
+// NewPredictionRequest is a helper function to create a new PredictionRequest with proper validation.
+func NewPredictionRequest(
+	kvCachePercentage float64,
+	inputTokenLength int,
+	numRequestWaiting int,
+	numRequestRunning int,
+	numTokensGenerated int,
+	prefixCacheScore float64,
+) (PredictionRequest, error) {
+	req := PredictionRequest{
+		KVCachePercentage:  kvCachePercentage,
+		InputTokenLength:   inputTokenLength,
+		NumRequestWaiting:  numRequestWaiting,
+		NumRequestRunning:  numRequestRunning,
+		NumTokensGenerated: numTokensGenerated,
+		PrefixCacheScore:   prefixCacheScore,
+	}
+
+	// Create a temporary predictor for validation (could be optimized)
+	p := &Predictor{}
+	if err := p.ValidatePredictionRequest(req); err != nil {
+		return PredictionRequest{}, err
+	}
+
+	return req, nil
+}
+
+// refreshServerStatus gets current server status from a prediction server
+func (p *Predictor) refreshServerStatus(ctx context.Context) error {
+	predictionURL := p.getRandomPredictionURL()
+	url := predictionURL + "/status"
+
+	p.logger.V(logutil.DEBUG).Info("Fetching server status", "url", url)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create server status request: %w", err)
+	}
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("failed to call /status endpoint: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("server %s returned non-200 status: %d %s, body: %s", url, resp.StatusCode, resp.Status, string(body))
+	}
+
+	var status ServerStatusResponse
+	if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
+		return fmt.Errorf("failed to decode server status response: %w", err)
+	}
+
+	p.metricsMu.Lock()
+	p.serverStatus = &status
+	p.metricsMu.Unlock()
+
+	p.logger.V(logutil.DEBUG).Info("Retrieved server status",
+		"model_type", status.ModelType,
+		"quantile", status.Quantile,
+		"is_ready", status.IsReady)
+	return nil
+}
+
+// getRandomPredictionURL returns a randomly selected prediction URL for load balancing
+func (p *Predictor) getRandomPredictionURL() string {
+	if len(p.config.PredictionURLs) == 0 {
+		return p.config.TrainingURL // Fallback to training URL
+	}
+	if len(p.config.PredictionURLs) == 1 {
+		return p.config.PredictionURLs[0]
+	}
+	index := p.rng.Intn(len(p.config.PredictionURLs))
+	return p.config.PredictionURLs[index]
+}
diff --git a/sidecars/latencypredictorasync/training.go b/sidecars/latencypredictorasync/training.go
new file mode 100644
index 0000000..4bbf240
--- /dev/null
+++ b/sidecars/latencypredictorasync/training.go
@@ -0,0 +1,420 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
+)
+
+// AddTrainingDataBulk buffers entries for periodic flush.
+func (p *Predictor) AddTrainingDataBulk(entries []TrainingEntry) error {
+	p.bufferMu.Lock()
+	p.pending = append(p.pending, entries...)
+	p.bufferMu.Unlock()
+	return nil
+}
+
+// randomSample returns up to maxSize entries via stratified sampling to preserve
+// the ratio of TTFT entries (ActualTTFT > 0) and TPOT entries (ActualTPOT > 0).
+func (p *Predictor) randomSample(entries []TrainingEntry, maxSize int) []TrainingEntry {
+	if len(entries) <= maxSize {
+		return entries
+	}
+
+	// Separate entries into three groups
+	var ttftEntries []TrainingEntry
+	var tpotEntries []TrainingEntry
+	var otherEntries []TrainingEntry
+
+	for _, entry := range entries {
+		hasTTFT := entry.ActualTTFT > 0
+		hasTPOT := entry.ActualTPOT > 0
+
+		switch {
+		case hasTTFT:
+			ttftEntries = append(ttftEntries, entry)
+		case hasTPOT:
+			tpotEntries = append(tpotEntries, entry)
+		default:
+			otherEntries = append(otherEntries, entry)
+		}
+	}
+
+	totalEntries := len(entries)
+	if totalEntries == 0 {
+		return entries
+	}
+
+	// Calculate proportional sample sizes
+	ttftSampleSize := int(float64(len(ttftEntries)) / float64(totalEntries) * float64(maxSize))
+	tpotSampleSize := int(float64(len(tpotEntries)) / float64(totalEntries) * float64(maxSize))
+	otherSampleSize := int(float64(len(otherEntries)) / float64(totalEntries) * float64(maxSize))
+
+	// Adjust for rounding errors to ensure we reach exactly maxSize
+	totalSampled := ttftSampleSize + tpotSampleSize + otherSampleSize
+	if totalSampled < maxSize {
+		remaining := maxSize - totalSampled
+
+		// Distribute remaining samples to the largest *original* group
+		switch {
+		case len(ttftEntries) >= len(tpotEntries) && len(ttftEntries) >= len(otherEntries):
+			ttftSampleSize += remaining
+		case len(tpotEntries) >= len(otherEntries):
+			tpotSampleSize += remaining
+		default:
+			otherSampleSize += remaining
+		}
+
+	} else if totalSampled > maxSize {
+		excess := totalSampled - maxSize
+
+		// Reduce from the largest *sampled* group
+		switch {
+		case ttftSampleSize >= tpotSampleSize && ttftSampleSize >= otherSampleSize:
+			ttftSampleSize -= excess
+		case tpotSampleSize >= otherSampleSize:
+			tpotSampleSize -= excess
+		default:
+			otherSampleSize -= excess
+		}
+	}
+
+	var result []TrainingEntry
+
+	// Sample from each group
+	if ttftSampleSize > 0 && len(ttftEntries) > 0 {
+		ttftSample := p.sampleFromSlice(ttftEntries, min(ttftSampleSize, len(ttftEntries)))
+		result = append(result, ttftSample...)
+	}
+
+	if tpotSampleSize > 0 && len(tpotEntries) > 0 {
+		tpotSample := p.sampleFromSlice(tpotEntries, min(tpotSampleSize, len(tpotEntries)))
+		result = append(result, tpotSample...)
+	}
+
+	if otherSampleSize > 0 && len(otherEntries) > 0 {
+		otherSample := p.sampleFromSlice(otherEntries, min(otherSampleSize, len(otherEntries)))
+		result = append(result, otherSample...)
+	}
+
+	return result
+}
+
+// Helper function to sample from a slice
+func (p *Predictor) sampleFromSlice(entries []TrainingEntry, sampleSize int) []TrainingEntry {
+	if len(entries) <= sampleSize {
+		return entries
+	}
+
+	// Create a copy and shuffle
+	sample := make([]TrainingEntry, len(entries))
+	copy(sample, entries)
+	p.rng.Shuffle(len(sample), func(i, j int) {
+		sample[i], sample[j] = sample[j], sample[i]
+	})
+
+	return sample[:sampleSize]
+}
+
+// Helper function to get minimum of two integers
+func min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// flushTraining sends buffered entries to training server in one bulk POST, with error handling.
+func (p *Predictor) flushTraining() {
+	p.bufferMu.Lock()
+	if len(p.pending) == 0 {
+		p.bufferMu.Unlock()
+		return
+	}
+	batch := p.pending
+	p.pending = nil
+	p.bufferMu.Unlock()
+
+	originalSize := len(batch)
+	if originalSize > p.config.MaxSampleSize {
+		batch = p.randomSample(batch, p.config.MaxSampleSize)
+		p.logger.V(logutil.DEBUG).Info("Sampled training entries for flush",
+			"original_size", originalSize,
+			"sampled_size", len(batch))
+	}
+
+	payload := BulkTrainingRequest{Entries: batch}
+	data, err := json.Marshal(payload)
+	if err != nil {
+		p.logger.Error(err, "Failed to marshal bulk payload")
+		return // Cannot send if marshalling fails
+	}
+
+	// Send training data to training server
+	url := p.config.TrainingURL + "/add_training_data_bulk"
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, url, bytes.NewBuffer(data))
+	if err != nil {
+		p.logger.Error(err, "Failed to create bulk POST request", "url", url)
+		return
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		p.logger.Error(err, "Bulk POST failed", "url", url)
+		return
+	}
+	defer resp.Body.Close()
+
+	// Ensure body is read and closed
+	if _, err := io.Copy(io.Discard, resp.Body); err != nil {
+		p.logger.Error(err, "failed to read response body", "url", url)
+	}
+
+	if resp.StatusCode != http.StatusAccepted {
+		p.logger.Error(fmt.Errorf("status %d", resp.StatusCode),
+			"Bulk POST returned non-202 status", "url", url)
+	} else {
+		p.logger.V(logutil.DEBUG).Info("Flushed training batch", "sent_count", len(batch), "original_count", originalSize)
+	}
+}
+
+// ValidateTrainingEntry validates that a training entry has all required fields
+// with valid values, including the new prefix_cache_score field.
+func (p *Predictor) ValidateTrainingEntry(entry TrainingEntry) error {
+	if entry.KVCachePercentage < 0.0 || entry.KVCachePercentage > 1.0 {
+		return fmt.Errorf("kv_cache_percentage must be between 0.0 and 1.0, got %f", entry.KVCachePercentage)
+	}
+	if entry.InputTokenLength < 0 {
+		return fmt.Errorf("input_token_length must be non-negative, got %d", entry.InputTokenLength)
+	}
+	if entry.NumRequestWaiting < 0 {
+		return fmt.Errorf("num_request_waiting must be non-negative, got %d", entry.NumRequestWaiting)
+	}
+	if entry.NumRequestRunning < 0 {
+		return fmt.Errorf("num_request_running must be non-negative, got %d", entry.NumRequestRunning)
+	}
+	if entry.NumTokensGenerated < 0 {
+		return fmt.Errorf("num_tokens_generated must be non-negative, got %d", entry.NumTokensGenerated)
+	}
+	if entry.ActualTTFT < 0.0 {
+		return fmt.Errorf("actual_ttft_ms must be non-negative, got %f", entry.ActualTTFT)
+	}
+	if entry.ActualTPOT < 0.0 {
+		return fmt.Errorf("actual_tpot_ms must be non-negative, got %f", entry.ActualTPOT)
+	}
+	if entry.PrefixCacheScore < 0.0 || entry.PrefixCacheScore > 1.0 {
+		return fmt.Errorf("prefix_cache_score must be between 0.0 and 1.0, got %f", entry.PrefixCacheScore)
+	}
+	return nil
+}
+
+// NewTrainingEntry is a helper function to create a new TrainingEntry with proper validation.
+func NewTrainingEntry(
+	kvCachePercentage float64,
+	inputTokenLength int,
+	numRequestWaiting int,
+	numRequestRunning int,
+	numTokensGenerated int,
+	actualTTFT float64,
+	actualTPOT float64,
+	prefixCacheScore float64,
+) (TrainingEntry, error) {
+	entry := TrainingEntry{
+		KVCachePercentage:  kvCachePercentage,
+		InputTokenLength:   inputTokenLength,
+		NumRequestWaiting:  numRequestWaiting,
+		NumRequestRunning:  numRequestRunning,
+		NumTokensGenerated: numTokensGenerated,
+		ActualTTFT:         actualTTFT,
+		ActualTPOT:         actualTPOT,
+		PrefixCacheScore:   prefixCacheScore,
+		Timestamp:          time.Now(),
+	}
+
+	// Create a temporary predictor for validation (could be optimized)
+	p := &Predictor{}
+	if err := p.ValidateTrainingEntry(entry); err != nil {
+		return TrainingEntry{}, err
+	}
+
+	return entry, nil
+}
+
+// refreshModelInfo gets current model type and readiness info from training server
+func (p *Predictor) refreshModelInfo(ctx context.Context) error {
+	url := p.config.TrainingURL + "/model/download/info"
+	p.logger.V(logutil.DEBUG).Info("Fetching model info", "url", url)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create model info request: %w", err)
+	}
+
+	resp, err := p.httpClient.Do(req)
+	if err != nil {
+		return fmt.Errorf("failed to call /model/download/info endpoint: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("server %s returned non-200 status: %d %s, body: %s", url, resp.StatusCode, resp.Status, string(body))
+	}
+
+	var modelInfo ModelInfo
+	if err := json.NewDecoder(resp.Body).Decode(&modelInfo); err != nil {
+		return fmt.Errorf("failed to decode model info response: %w", err)
+	}
+
+	p.metricsMu.Lock()
+	p.modelInfo = &modelInfo
+	p.metricsMu.Unlock()
+
+	p.logger.V(logutil.DEBUG).Info("Retrieved model info", "model_type", modelInfo.ModelType, "model_status", modelInfo.ModelStatus)
+	return nil
+}
+
+// refreshMetrics GETs /metrics from training server and caches parsed coefficients or fetches XGBoost trees.
+func (p *Predictor) refreshMetrics() {
+	ctx, cancel := context.WithTimeout(context.Background(), p.config.HTTPTimeout)
+	defer cancel()
+
+	// Refresh model info first
+	if err := p.refreshModelInfo(ctx); err != nil {
+		p.logger.Error(err, "Failed to refresh model info during periodic refresh")
+		return
+	}
+
+	p.metricsMu.RLock()
+	modelType := ""
+	if p.modelInfo != nil {
+		modelType = p.modelInfo.ModelType
+	}
+	p.metricsMu.RUnlock()
+
+	if modelType == "" {
+		p.logger.V(logutil.DEBUG).Info("Cannot refresh metrics: model type is unknown")
+		return
+	}
+
+	switch modelType {
+	case bayesianRidgeModelType:
+		if _, err := p.GetMetrics(ctx); err != nil {
+			p.logger.Error(err, "Failed to refresh Bayesian Ridge metrics")
+		}
+	case xgBoostModelType:
+		if p.config.UseNativeXGBoost {
+			// Fetch XGBoost trees for native predictions
+			trees, err := p.getXGBoostTrees(ctx)
+			if err != nil {
+				p.logger.Error(err, "Failed to fetch XGBoost trees")
+				return
+			}
+
+			p.metricsMu.Lock()
+			if p.cachedMetrics == nil {
+				p.cachedMetrics = &MetricsResponse{}
+			}
+			p.cachedMetrics.ModelType = modelType
+			p.cachedMetrics.XGBoostTrees = trees
+			p.metricsMu.Unlock()
+
+			p.logger.V(logutil.DEBUG).Info("Updated XGBoost trees for native predictions")
+		} else {
+			// Just update model type for HTTP-based predictions
+			p.metricsMu.Lock()
+			if p.cachedMetrics == nil {
+				p.cachedMetrics = &MetricsResponse{}
+			}
+			p.cachedMetrics.ModelType = modelType
+			p.metricsMu.Unlock()
+
+			p.logger.V(logutil.DEBUG).Info("Updated model type for HTTP-based predictions", "model_type", modelType)
+		}
+	case gbmModelType:
+		// LightGBM only supports HTTP calls, no native tree caching needed
+		p.metricsMu.Lock()
+		if p.cachedMetrics == nil {
+			p.cachedMetrics = &MetricsResponse{}
+		}
+		p.cachedMetrics.ModelType = modelType
+		p.metricsMu.Unlock()
+
+		p.logger.V(logutil.DEBUG).Info("Updated model type for HTTP-based predictions", "model_type", modelType)
+	default:
+		p.logger.Info("Unknown model type, cannot refresh metrics", "model_type", modelType)
+	}
+}
+
+// getXGBoostTrees fetches tree JSON from the training server
+func (p *Predictor) getXGBoostTrees(ctx context.Context) (*XGBoostTrees, error) {
+	trees := &XGBoostTrees{}
+
+	// Fetch TTFT trees from training server
+	ttftURL := p.config.TrainingURL + "/model/ttft/xgb/json"
+	ttftReq, err := http.NewRequestWithContext(ctx, http.MethodGet, ttftURL, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create TTFT trees request: %w", err)
+	}
+
+	ttftResp, err := p.httpClient.Do(ttftReq)
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch TTFT trees: %w", err)
+	}
+	defer ttftResp.Body.Close()
+
+	if ttftResp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(ttftResp.Body)
+		return nil, fmt.Errorf("TTFT trees request failed: %d %s, body: %s", ttftResp.StatusCode, ttftResp.Status, string(body))
+	}
+
+	if err := json.NewDecoder(ttftResp.Body).Decode(&trees.TTFTTrees); err != nil {
+		return nil, fmt.Errorf("failed to decode TTFT trees: %w", err)
+	}
+
+	// Fetch TPOT trees from training server
+	tpotURL := p.config.TrainingURL + "/model/tpot/xgb/json"
+	tpotReq, err := http.NewRequestWithContext(ctx, http.MethodGet, tpotURL, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create TPOT trees request: %w", err)
+	}
+
+	tpotResp, err := p.httpClient.Do(tpotReq)
+	if err != nil {
+		return nil, fmt.Errorf("failed to fetch TPOT trees: %w", err)
+	}
+	defer tpotResp.Body.Close()
+
+	if tpotResp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(tpotResp.Body)
+		return nil, fmt.Errorf("TPOT trees request failed: %d %s, body: %s", tpotResp.StatusCode, tpotResp.Status, string(body))
+	}
+
+	if err := json.NewDecoder(tpotResp.Body).Decode(&trees.TPOTTrees); err != nil {
+		return nil, fmt.Errorf("failed to decode TPOT trees: %w", err)
+	}
+
+	return trees, nil
+}
diff --git a/sidecars/latencypredictorasync/types.go b/sidecars/latencypredictorasync/types.go
new file mode 100644
index 0000000..4b4a1ca
--- /dev/null
+++ b/sidecars/latencypredictorasync/types.go
@@ -0,0 +1,233 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"context"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const (
+	gbmModelType           = "lightgbm"
+	bayesianRidgeModelType = "bayesian_ridge"
+	xgBoostModelType       = "xgboost"
+)
+
+// --- Configuration ---
+
+type Config struct {
+	// TrainingURL is the base URL of the Python training server.
+	TrainingURL string
+	// PredictionURLs is a list of prediction server URLs for load balancing.
+	PredictionURLs []string
+	// MaxSampleSize is the maximum number of training entries to send in each flush.
+	// If the buffer contains more entries, they will be randomly sampled.
+	MaxSampleSize int
+	// FlushInterval determines how often to flush training & refresh metrics.
+	FlushInterval time.Duration
+	// UseNativeXGBoost when true, attempts to use local XGBoost models for prediction.
+	// When false, falls back to HTTP calls to the Python server for XGBoost predictions.
+	UseNativeXGBoost bool
+	// HTTPTimeout is the timeout for HTTP requests to the Python server.
+	HTTPTimeout time.Duration
+	// MetricsRefreshInterval determines how often to refresh cached metrics.
+	MetricsRefreshInterval time.Duration
+	// MaxBulkSize is the maximum number of predictions to send in a single bulk request.
+	MaxBulkSize int
+}
+
+func DefaultConfig() *Config {
+	return &Config{
+		TrainingURL:            "http://localhost:8000",
+		PredictionURLs:         []string{"http://localhost:8001"},
+		MaxSampleSize:          1000,
+		FlushInterval:          1 * time.Second,
+		MetricsRefreshInterval: 60 * time.Second,
+		UseNativeXGBoost:       true,
+		HTTPTimeout:            10 * time.Second,
+		MaxBulkSize:            100,
+	}
+}
+
+func ConfigFromEnv() *Config {
+	cfg := DefaultConfig()
+
+	// Training URL (single URL for training data submission)
+	if url := os.Getenv("TRAINING_SERVER_URL"); url != "" {
+		cfg.TrainingURL = url
+	}
+
+	// Prediction URLs (comma-separated list for load balancing)
+	if urls := os.Getenv("PREDICTION_SERVER_URL"); urls != "" {
+		predictionURLs := strings.Split(urls, ",")
+		for i, url := range predictionURLs {
+			predictionURLs[i] = strings.TrimSpace(url)
+		}
+		cfg.PredictionURLs = predictionURLs
+	}
+
+	if sizeStr := os.Getenv("LATENCY_MAX_SAMPLE_SIZE"); sizeStr != "" {
+		if size, err := strconv.Atoi(sizeStr); err == nil && size > 0 {
+			cfg.MaxSampleSize = size
+		}
+	}
+	if intervalStr := os.Getenv("LATENCY_FLUSH_INTERVAL_SEC"); intervalStr != "" {
+		if sec, err := strconv.Atoi(intervalStr); err == nil && sec > 0 {
+			cfg.FlushInterval = time.Duration(sec) * time.Second
+		}
+	}
+	if nativeStr := os.Getenv("LATENCY_USE_NATIVE_XGBOOST"); nativeStr != "" {
+		cfg.UseNativeXGBoost = strings.ToLower(nativeStr) == "true"
+	}
+	if timeoutStr := os.Getenv("LATENCY_HTTP_TIMEOUT_SEC"); timeoutStr != "" {
+		if sec, err := strconv.Atoi(timeoutStr); err == nil && sec > 0 {
+			cfg.HTTPTimeout = time.Duration(sec) * time.Second
+		}
+	}
+	if s := os.Getenv("LATENCY_METRICS_INTERVAL_SEC"); s != "" {
+		if sec, err := strconv.Atoi(s); err == nil && sec > 0 {
+			cfg.MetricsRefreshInterval = time.Duration(sec) * time.Second
+		}
+	}
+	if bulkStr := os.Getenv("LATENCY_MAX_BULK_SIZE"); bulkStr != "" {
+		if size, err := strconv.Atoi(bulkStr); err == nil && size > 0 && size <= 100 {
+			cfg.MaxBulkSize = size
+		}
+	}
+	return cfg
+}
+
+// Predictor defines the interface for latency prediction and training.
+type PredictorInterface interface {
+	Predict(ctx context.Context, req PredictionRequest) (*PredictionResponse, error)
+	PredictBulk(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error)
+	PredictBulkStrict(ctx context.Context, requests []PredictionRequest) (*BulkPredictionResponse, error)
+	AddTrainingDataBulk(entry []TrainingEntry) error
+	GetServerStatus(ctx context.Context) (*ServerStatusResponse, error)
+}
+
+// --- Data Models ---
+
+type TrainingEntry struct {
+	KVCachePercentage  float64   `json:"kv_cache_percentage"`
+	InputTokenLength   int       `json:"input_token_length"`
+	NumRequestWaiting  int       `json:"num_request_waiting"`
+	NumRequestRunning  int       `json:"num_request_running"`
+	NumTokensGenerated int       `json:"num_tokens_generated"`
+	ActualTTFT         float64   `json:"actual_ttft_ms"`
+	ActualTPOT         float64   `json:"actual_tpot_ms"`
+	PrefixCacheScore   float64   `json:"prefix_cache_score"`
+	Timestamp          time.Time `json:"timestamp"`
+}
+
+type BulkTrainingRequest struct {
+	Entries []TrainingEntry `json:"entries"`
+}
+
+type PredictionRequest struct {
+	KVCachePercentage  float64 `json:"kv_cache_percentage"`
+	InputTokenLength   int     `json:"input_token_length"`
+	NumRequestWaiting  int     `json:"num_request_waiting"`
+	NumRequestRunning  int     `json:"num_request_running"`
+	NumTokensGenerated int     `json:"num_tokens_generated"`
+	PrefixCacheScore   float64 `json:"prefix_cache_score"`
+}
+
+type PredictionResponse struct {
+	TTFT                 float64    `json:"ttft_ms"`
+	TPOT                 float64    `json:"tpot_ms"`
+	TTFTUncertainty      float64    `json:"ttft_uncertainty,omitempty"`
+	TPOTUncertainty      float64    `json:"tpot_uncertainty,omitempty"`
+	TTFTPredictionBounds [2]float64 `json:"ttft_prediction_bounds,omitempty"`
+	TPOTPredictionBounds [2]float64 `json:"tpot_prediction_bounds,omitempty"`
+	PredictedAt          time.Time  `json:"predicted_at"`
+	ModelType            string     `json:"model_type"`
+	Quantile             float64    `json:"quantile"`
+	LastModelLoad        *time.Time `json:"last_model_load"`
+}
+
+// New data models for bulk predictions
+type BulkPredictionRequest struct {
+	Requests []PredictionRequest `json:"requests"`
+}
+
+type BulkPredictionResponse struct {
+	Predictions           []PredictionResponse `json:"predictions"`
+	TotalRequests         int                  `json:"total_requests"`
+	SuccessfulPredictions int                  `json:"successful_predictions"`
+	FailedPredictions     int                  `json:"failed_predictions"`
+	ProcessingTimeMs      float64              `json:"processing_time_ms"`
+}
+
+type BulkPredictionError struct {
+	Index   int               `json:"index"`
+	Error   string            `json:"error"`
+	Request PredictionRequest `json:"request"`
+}
+
+type BulkPredictionResponseWithErrors struct {
+	Predictions           []*PredictionResponse `json:"predictions"`
+	Errors                []BulkPredictionError `json:"errors"`
+	TotalRequests         int                   `json:"total_requests"`
+	SuccessfulPredictions int                   `json:"successful_predictions"`
+	FailedPredictions     int                   `json:"failed_predictions"`
+	ProcessingTimeMs      float64               `json:"processing_time_ms"`
+}
+
+// Server status response
+type ServerStatusResponse struct {
+	IsReady           bool            `json:"is_ready"`
+	ModelType         string          `json:"model_type"`
+	Quantile          float64         `json:"quantile"`
+	LastModelLoad     *time.Time      `json:"last_model_load"`
+	TrainingServerURL string          `json:"training_server_url"`
+	ModelsExist       map[string]bool `json:"models_exist"`
+}
+
+type ModelCoefficients struct {
+	TTFTIntercept float64            `json:"ttft_intercept"`
+	TTFTCoeffs    map[string]float64 `json:"ttft_coefficients"`
+	TPOTIntercept float64            `json:"tpot_intercept"`
+	TPOTCoeffs    map[string]float64 `json:"tpot_coefficients"`
+}
+
+type XGBoostTrees struct {
+	TTFTTrees []interface{} `json:"ttft_trees"`
+	TPOTTrees []interface{} `json:"tpot_trees"`
+}
+
+type BucketCounts struct {
+	TTFTBuckets map[int]int `json:"ttft_buckets"`
+	TPOTBuckets map[int]int `json:"tpot_buckets"`
+}
+
+type ModelInfo struct {
+	ModelType   string          `json:"model_type"`
+	ModelStatus map[string]bool `json:"model_status"`
+	Quantile    float64         `json:"quantile"`
+}
+
+type MetricsResponse struct {
+	ModelType    string             `json:"model_type"`
+	Coefficients *ModelCoefficients `json:"coefficients"`
+	XGBoostTrees *XGBoostTrees      `json:"xgboost_trees"`
+	BucketCounts *BucketCounts      `json:"bucket_counts"`
+	RawMetrics   string             `json:"raw_metrics"`
+}
diff --git a/site-src/guides/getting-started-latest.md b/site-src/guides/getting-started-latest.md
index bf7413b..7629f4a 100644
--- a/site-src/guides/getting-started-latest.md
+++ b/site-src/guides/getting-started-latest.md
@@ -56,7 +56,7 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
       1. Requirements
          - Gateway API [CRDs](https://gateway-api.sigs.k8s.io/guides/#installing-gateway-api) installed.
 
-      2. Install Istio
+      1. Install Istio:
 
          ```
          TAG=$(curl https://storage.googleapis.com/istio-build/dev/1.28-dev)
@@ -79,14 +79,14 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
 
          - Gateway API [CRDs](https://gateway-api.sigs.k8s.io/guides/#installing-gateway-api) installed.
 
-      2. Set the Kgateway version and install the Kgateway CRDs.
+      1. Set the Kgateway version and install the Kgateway CRDs:
 
          ```bash
          KGTW_VERSION=v2.1.0
          helm upgrade -i --create-namespace --namespace kgateway-system --version $KGTW_VERSION kgateway-crds oci://cr.kgateway.dev/kgateway-dev/charts/kgateway-crds
          ```
 
-      3. Install Kgateway
+      1. Install Kgateway:
 
          ```bash
          helm upgrade -i --namespace kgateway-system --version $KGTW_VERSION kgateway oci://cr.kgateway.dev/kgateway-dev/charts/kgateway --set inferenceExtension.enabled=true
@@ -114,7 +114,7 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
          See [Deploy Inference Gateways](https://cloud.google.com/kubernetes-engine/docs/how-to/deploy-gke-inference-gateway)
          for detailed instructions.
 
-      2. Deploy the Inference Gateway:
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/gke/gateway.yaml
@@ -127,13 +127,13 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
          NAME                CLASS               ADDRESS         PROGRAMMED   AGE
          inference-gateway   inference-gateway   <MY_ADDRESS>    True         22s
          ```
-      3. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/gke/httproute.yaml
          ```
 
-      4. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -144,7 +144,7 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
       Please note that this feature is currently in an experimental phase and is not intended for production use.
       The implementation and user experience are subject to changes as we continue to iterate on this project.
  
-      1. Deploy the Inference Gateway
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/istio/gateway.yaml
@@ -157,13 +157,13 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
          inference-gateway   inference-gateway   <MY_ADDRESS>    True         22s
          ```
 
-      2. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/istio/httproute.yaml
          ```
 
-      3. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -176,7 +176,7 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
       implementation. Kgateway supports Inference Gateway with the [agentgateway](https://agentgateway.dev/) data plane. Follow these steps
       to run Kgateway as an Inference Gateway:
 
-      1. Deploy the Inference Gateway
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/kgateway/gateway.yaml
@@ -187,13 +187,13 @@ kubectl apply -k https://github.com/kubernetes-sigs/gateway-api-inference-extens
          kubectl get gateway inference-gateway
          ```
 
-      2. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/kgateway/httproute.yaml
          ```
 
-      3. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -216,7 +216,7 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
    The following instructions assume you would like to cleanup ALL resources that were created in this quickstart guide.
    Please be careful not to delete resources you'd like to keep.
 
-   1. Uninstall the InferencePool, InferenceObjective and model server resources
+   1. Uninstall the InferencePool, InferenceObjective and model server resources:
 
       ```bash
       helm uninstall vllm-llama3-8b-instruct
@@ -227,7 +227,7 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
       kubectl delete secret hf-token --ignore-not-found
       ```
 
-   1. Uninstall the Gateway API Inference Extension CRDs
+   1. Uninstall the Gateway API Inference Extension CRDs:
 
       ```bash
       kubectl delete -k https://github.com/kubernetes-sigs/gateway-api-inference-extension/config/crd --ignore-not-found
@@ -252,13 +252,13 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
 
       The following steps assume you would like to clean up ALL Istio resources that were created in this quickstart guide.
 
-      1. Uninstall All Istio resources
+      1. Uninstall All Istio resources:
 
          ```bash
          istioctl uninstall -y --purge
          ```
 
-      2. Remove the Istio namespace
+      1. Remove the Istio namespace:
 
          ```bash
          kubectl delete ns istio-system
@@ -273,19 +273,19 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
 
       The following steps assume you would like to cleanup ALL Kgateway resources that were created in this quickstart guide.
 
-      1. Uninstall Kgateway
+      1. Uninstall Kgateway:
 
          ```bash
          helm uninstall kgateway -n kgateway-system
          ```
 
-      2. Uninstall the Kgateway CRDs.
+      1. Uninstall the Kgateway CRDs:
 
          ```bash
          helm uninstall kgateway-crds -n kgateway-system
          ```
 
-      3. Remove the Kgateway namespace.
+      1. Remove the Kgateway namespace:
 
          ```bash
          kubectl delete ns kgateway-system
diff --git a/site-src/guides/index.md b/site-src/guides/index.md
index 205e678..db47db0 100644
--- a/site-src/guides/index.md
+++ b/site-src/guides/index.md
@@ -52,7 +52,7 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
       1. Requirements
          - Gateway API [CRDs](https://gateway-api.sigs.k8s.io/guides/#installing-gateway-api) installed.
 
-      2. Install Istio
+      1. Install Istio:
 
          ```
          TAG=$(curl https://storage.googleapis.com/istio-build/dev/1.28-dev)
@@ -75,14 +75,14 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
 
          - Gateway API [CRDs](https://gateway-api.sigs.k8s.io/guides/#installing-gateway-api) installed.
 
-      2. Set the Kgateway version and install the Kgateway CRDs.
+      1. Set the Kgateway version and install the Kgateway CRDs:
 
          ```bash
          KGTW_VERSION=v2.2.0-main
          helm upgrade -i --create-namespace --namespace kgateway-system --version $KGTW_VERSION kgateway-crds oci://cr.kgateway.dev/kgateway-dev/charts/kgateway-crds
          ```
 
-      3. Install Kgateway
+      1. Install Kgateway:
 
          ```bash
          helm upgrade -i --namespace kgateway-system --version $KGTW_VERSION kgateway oci://cr.kgateway.dev/kgateway-dev/charts/kgateway --set inferenceExtension.enabled=true
@@ -110,7 +110,7 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
          See [Deploy Inference Gateways](https://cloud.google.com/kubernetes-engine/docs/how-to/deploy-gke-inference-gateway)
          for detailed instructions.
 
-      2. Deploy the Inference Gateway:
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://raw.githubusercontent.com/kubernetes-sigs/gateway-api-inference-extension/refs/tags/v1.1.0/config/manifests/gateway/gke/gateway.yaml
@@ -123,13 +123,13 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
          NAME                CLASS               ADDRESS         PROGRAMMED   AGE
          inference-gateway   inference-gateway   <MY_ADDRESS>    True         22s
          ```
-      3. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://raw.githubusercontent.com/kubernetes-sigs/gateway-api-inference-extension/refs/tags/v1.1.0/config/manifests/gateway/gke/httproute.yaml
          ```
 
-      4. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -140,7 +140,7 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
       Please note that this feature is currently in an experimental phase and is not intended for production use.
       The implementation and user experience are subject to changes as we continue to iterate on this project.
 
-      1. Deploy the Inference Gateway
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/istio/gateway.yaml
@@ -153,13 +153,13 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
          inference-gateway   inference-gateway   <MY_ADDRESS>    True         22s
          ```
 
-      2. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/raw/main/config/manifests/gateway/istio/httproute.yaml
          ```
 
-      3. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -172,7 +172,7 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
       implementation. Kgateway supports Inference Gateway with the [agentgateway](https://agentgateway.dev/) data plane. Follow these steps
       to run Kgateway as an Inference Gateway:
 
-      1. Deploy the Inference Gateway
+      1. Deploy the Inference Gateway:
 
          ```bash
          kubectl apply -f https://raw.githubusercontent.com/kubernetes-sigs/gateway-api-inference-extension/refs/tags/v1.1.0/config/manifests/gateway/agentgateway/gateway.yaml
@@ -183,13 +183,13 @@ kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extens
          kubectl get gateway inference-gateway
          ```
 
-      2. Deploy the HTTPRoute
+      1. Deploy the HTTPRoute:
 
          ```bash
          kubectl apply -f https://raw.githubusercontent.com/kubernetes-sigs/gateway-api-inference-extension/refs/tags/v1.1.0/config/manifests/gateway/agentgateway/httproute.yaml
          ```
 
-      3. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
+      1. Confirm that the HTTPRoute status conditions include `Accepted=True` and `ResolvedRefs=True`:
 
          ```bash
          kubectl get httproute llm-route -o yaml
@@ -212,7 +212,7 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
    The following instructions assume you would like to cleanup ALL resources that were created in this quickstart guide.
    Please be careful not to delete resources you'd like to keep.
 
-   1. Uninstall the InferencePool, InferenceObjective and model server resources
+   1. Uninstall the InferencePool, InferenceObjective and model server resources:
 
       ```bash
       helm uninstall vllm-llama3-8b-instruct
@@ -223,7 +223,7 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
       kubectl delete secret hf-token --ignore-not-found
       ```
 
-   1. Uninstall the Gateway API Inference Extension CRDs
+   1. Uninstall the Gateway API Inference Extension CRDs:
 
       ```bash
       kubectl delete -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.1.0/manifests.yaml --ignore-not-found
@@ -249,13 +249,13 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
 
       The following steps assume you would like to clean up ALL Istio resources that were created in this quickstart guide.
 
-      1. Uninstall All Istio resources
+      1. Uninstall All Istio resources:
 
          ```bash
          istioctl uninstall -y --purge
          ```
 
-      2. Remove the Istio namespace
+      1. Remove the Istio namespace:
 
          ```bash
          kubectl delete ns istio-system
@@ -270,19 +270,19 @@ Deploy the sample InferenceObjective which allows you to specify priority of req
 
       The following steps assume you would like to cleanup ALL Kgateway resources that were created in this quickstart guide.
 
-      1. Uninstall Kgateway
+      1. Uninstall Kgateway:
 
          ```bash
          helm uninstall kgateway -n kgateway-system
          ```
 
-      2. Uninstall the Kgateway CRDs.
+      1. Uninstall the Kgateway CRDs:
 
          ```bash
          helm uninstall kgateway-crds -n kgateway-system
          ```
 
-      3. Remove the Kgateway namespace.
+      1. Remove the Kgateway namespace:
 
          ```bash
          kubectl delete ns kgateway-system



Test patch:
diff --git a/sidecars/latencypredictorasync/latencypredictor_async_test.go b/sidecars/latencypredictorasync/latencypredictor_async_test.go
new file mode 100644
index 0000000..baabe3e
--- /dev/null
+++ b/sidecars/latencypredictorasync/latencypredictor_async_test.go
@@ -0,0 +1,2245 @@
+/*
+Copyright 2025 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package latencypredictorasync
+
+import (
+	"context"
+	"math/rand"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/go-logr/logr"
+	"github.com/go-logr/zapr"
+	"go.uber.org/zap"
+)
+
+func TestLatencyPredictorIntegration(t *testing.T) {
+	// Setup logger
+	zapLog, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatalf("Failed to create logger: %v", err)
+	}
+	logger := zapr.NewLogger(zapLog)
+
+	// Check if server URLs are set
+	predictionURLs := os.Getenv("PREDICTION_SERVER_URL")
+	trainingURL := os.Getenv("TRAINING_SERVER_URL")
+
+	if predictionURLs == "" {
+		t.Skip("PREDICTION_SERVER_URL not set, skipping integration test")
+	}
+	if trainingURL == "" {
+		// Fallback to first prediction URL for training if not set
+		urls := strings.Split(predictionURLs, ",")
+		if len(urls) > 0 {
+			trainingURL = strings.TrimSpace(urls[0])
+		} else {
+			t.Skip("No valid URLs available for testing")
+		}
+	}
+
+	// Parse prediction URLs
+	urls := strings.Split(predictionURLs, ",")
+	parsedPredictionURLs := make([]string, 0, len(urls))
+	for _, url := range urls {
+		parsedPredictionURLs = append(parsedPredictionURLs, strings.TrimSpace(url))
+	}
+
+	// Create config with the actual server URLs
+	config := &Config{
+		TrainingURL:            trainingURL,
+		PredictionURLs:         parsedPredictionURLs,
+		MaxSampleSize:          1000,
+		FlushInterval:          500 * time.Millisecond, // Shorter for testing
+		MetricsRefreshInterval: 1 * time.Second,        // Longer for metrics
+		UseNativeXGBoost:       true,
+		HTTPTimeout:            30 * time.Second, // Longer timeout for tests
+		MaxBulkSize:            50,               // Test bulk size
+	}
+
+	// Create predictor
+	predictor := New(config, logger)
+	defer predictor.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	// Start the predictor
+	err = predictor.Start(ctx)
+	if err != nil {
+		t.Fatalf("Failed to start predictor: %v", err)
+	}
+
+	t.Run("TestModelInfo", func(t *testing.T) {
+		testModelInfo(t, ctx, predictor)
+	})
+
+	t.Run("TestBulkTrainingData", func(t *testing.T) {
+		testBulkTrainingData(t, predictor)
+	})
+
+	t.Run("TestPrediction", func(t *testing.T) {
+		testPrediction(t, ctx, predictor)
+	})
+
+	t.Run("TestBulkPredictions", func(t *testing.T) {
+		testBulkPredictions(t, ctx, predictor)
+	})
+
+	t.Run("TestBulkPredictionsStrict", func(t *testing.T) {
+		testBulkPredictionsStrict(t, ctx, predictor)
+	})
+
+	t.Run("TestPredictionWithPrefixCache", func(t *testing.T) {
+		testPredictionWithPrefixCache(t, ctx, predictor)
+	})
+
+	t.Run("TestHTTPFallbackPrediction", func(t *testing.T) {
+		testHTTPFallbackPrediction(t, ctx, predictor)
+	})
+
+	t.Run("TestLightGBMSupport", func(t *testing.T) {
+		testLightGBMSupport(t, ctx, predictor)
+	})
+
+	t.Run("TestPredictionPerformance", func(t *testing.T) {
+		testPredictionPerformance(t, ctx, predictor)
+	})
+
+	t.Run("TestBulkPredictionPerformance", func(t *testing.T) {
+		testBulkPredictionPerformance(t, ctx, predictor)
+	})
+
+	t.Run("TestHTTPOnlyPerformance", func(t *testing.T) {
+		testHTTPOnlyPerformance(t, ctx)
+	})
+
+	t.Run("TestXGBoostJSONStructure", func(t *testing.T) {
+		testXGBoostJSONStructure(t, ctx, predictor)
+	})
+
+	t.Run("TestHTTPOnlyPrediction", func(t *testing.T) {
+		testHTTPOnlyPrediction(t, ctx)
+	})
+
+	t.Run("TestMetricsRetrieval", func(t *testing.T) {
+		testMetricsRetrieval(t, ctx, predictor)
+	})
+
+	t.Run("TestLoadBalancing", func(t *testing.T) {
+		testLoadBalancing(t, ctx, predictor)
+	})
+
+	t.Run("TestPrefixCacheValidation", func(t *testing.T) {
+		testPrefixCacheValidation(t, predictor)
+	})
+
+	t.Run("TestPredictionConstructors", func(t *testing.T) {
+		testPredictionConstructors(t)
+	})
+}
+
+func testModelInfo(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing model info retrieval...")
+
+	modelInfo, err := predictor.GetModelInfo(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get model info: %v", err)
+	}
+
+	t.Logf("Model Info - Type: %s, Model Status: %v, Quantile: %.2f",
+		modelInfo.ModelType, modelInfo.ModelStatus, modelInfo.Quantile)
+
+	if modelInfo.ModelType == "" {
+		t.Error("Model type should not be empty")
+	}
+
+	// Store model type for other tests
+	currentModelType := predictor.GetCurrentModelType()
+	t.Logf("Current model type from predictor: %s", currentModelType)
+
+	// Log URLs being used
+	t.Logf("Training URL: %s", predictor.GetTrainingURL())
+	t.Logf("Prediction URLs: %v", predictor.GetPredictionURLs())
+}
+
+func testBulkTrainingData(t *testing.T, predictor *Predictor) {
+	t.Log("Testing bulk training data submission with prefix cache score...")
+
+	// Generate 1000 random training entries including prefix cache scores
+	entries := generateTrainingEntries(1000)
+
+	err := predictor.AddTrainingDataBulk(entries)
+	if err != nil {
+		t.Fatalf("Failed to add bulk training data: %v", err)
+	}
+
+	t.Logf("Successfully added %d training entries to buffer (with prefix cache scores)", len(entries))
+
+	// Wait a bit for the background flush to occur
+	time.Sleep(2 * time.Second)
+
+	t.Log("Training data should have been flushed to training server")
+}
+
+func testPrediction(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing prediction functionality...")
+
+	// Log current predictor state
+	t.Logf("Predictor state:")
+	t.Logf("  Current model type: %s", predictor.GetCurrentModelType())
+	t.Logf("  Current quantile: %.2f", predictor.GetCurrentQuantile())
+	t.Logf("  Overall ready: %t", predictor.IsReady())
+	t.Logf("  XGBoost ready: %t", predictor.IsXGBoostReady())
+	t.Logf("  LightGBM ready: %t", predictor.IsLightGBMReady())
+	t.Logf("  Bayesian Ridge ready: %t", predictor.IsBayesianRidgeReady())
+
+	// Wait for models to be ready
+	maxWait := 30 * time.Second
+	waitTime := 100 * time.Millisecond
+	elapsed := time.Duration(0)
+
+	for elapsed < maxWait {
+		if predictor.IsReady() {
+			break
+		}
+		time.Sleep(waitTime)
+		elapsed += waitTime
+	}
+
+	if !predictor.IsReady() {
+		t.Log("Warning: Predictor not ready after waiting, attempting prediction anyway")
+	}
+
+	// Create a sample prediction request with prefix cache score
+	req := PredictionRequest{
+		KVCachePercentage:  0.755, // 75.5% as a fraction
+		InputTokenLength:   512,
+		NumRequestWaiting:  3,
+		NumRequestRunning:  2,
+		NumTokensGenerated: 100,
+		PrefixCacheScore:   0.8, // 80% prefix cache hit rate
+	}
+
+	t.Logf("Making prediction request: %+v", req)
+
+	response, err := predictor.Predict(ctx, req)
+	if err != nil {
+		t.Fatalf("Failed to make prediction: %v", err)
+	}
+
+	t.Logf("Prediction Response:")
+	t.Logf("  TTFT: %.2f ms (uncertainty: %.2f)", response.TTFT, response.TTFTUncertainty)
+	t.Logf("  TPOT: %.2f ms (uncertainty: %.2f)", response.TPOT, response.TPOTUncertainty)
+	t.Logf("  TTFT Bounds: [%.2f, %.2f]", response.TTFTPredictionBounds[0], response.TTFTPredictionBounds[1])
+	t.Logf("  TPOT Bounds: [%.2f, %.2f]", response.TPOTPredictionBounds[0], response.TPOTPredictionBounds[1])
+	t.Logf("  Model Type: %s", response.ModelType)
+	t.Logf("  Quantile: %.2f", response.Quantile)
+	t.Logf("  Predicted At: %s", response.PredictedAt.Format(time.RFC3339))
+
+	// Validate response
+	if response.TTFT <= 0 {
+		t.Error("TTFT should be positive")
+	}
+	if response.TPOT <= 0 {
+		t.Error("TPOT should be positive")
+	}
+	if response.ModelType == "" {
+		t.Error("Model type should not be empty")
+	}
+	if response.Quantile <= 0 || response.Quantile >= 1 {
+		t.Errorf("Quantile should be between 0 and 1, got: %.2f", response.Quantile)
+	}
+
+	// Test multiple predictions to ensure consistency
+	t.Log("Testing multiple predictions with varying prefix cache scores...")
+	for i := 0; i < 5; i++ {
+		testReq := PredictionRequest{
+			KVCachePercentage:  float64(50+i*10) / 100.0, // Convert percentage to fraction
+			InputTokenLength:   256 + i*128,
+			NumRequestWaiting:  i,
+			NumRequestRunning:  1 + i,
+			NumTokensGenerated: 50 + i*25,
+			PrefixCacheScore:   float64(i*20) / 100.0, // Vary prefix cache from 0% to 80%
+		}
+
+		resp, err := predictor.Predict(ctx, testReq)
+		if err != nil {
+			t.Errorf("Prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		t.Logf("Prediction %d: TTFT=%.2f, TPOT=%.2f (prefix_cache=%.1f%%, quantile=%.2f)",
+			i+1, resp.TTFT, resp.TPOT, testReq.PrefixCacheScore*100, resp.Quantile)
+	}
+}
+
+func testBulkPredictions(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing bulk predictions with error tolerance...")
+
+	if !predictor.IsReady() {
+		t.Skip("Predictor not ready for bulk prediction testing")
+	}
+
+	// Create multiple prediction requests
+	requests := make([]PredictionRequest, 10)
+	for i := 0; i < 10; i++ {
+		requests[i] = PredictionRequest{
+			KVCachePercentage:  float64(40+i*5) / 100.0, // 40% to 85%
+			InputTokenLength:   200 + i*50,              // 200 to 650
+			NumRequestWaiting:  i % 5,                   // 0 to 4
+			NumRequestRunning:  (i % 3) + 1,             // 1 to 3
+			NumTokensGenerated: 25 + i*10,               // 25 to 115
+			PrefixCacheScore:   float64(i) / 9.0,        // 0.0 to 1.0
+		}
+	}
+
+	t.Logf("Making bulk prediction request with %d requests", len(requests))
+
+	bulkResponse, err := predictor.PredictBulk(ctx, requests)
+	if err != nil {
+		t.Fatalf("Bulk prediction failed: %v", err)
+	}
+
+	t.Logf("Bulk Prediction Response:")
+	t.Logf("  Total Requests: %d", bulkResponse.TotalRequests)
+	t.Logf("  Successful: %d", bulkResponse.SuccessfulPredictions)
+	t.Logf("  Failed: %d", bulkResponse.FailedPredictions)
+	t.Logf("  Processing Time: %.2f ms", bulkResponse.ProcessingTimeMs)
+
+	if bulkResponse.TotalRequests != len(requests) {
+		t.Errorf("Expected %d total requests, got %d", len(requests), bulkResponse.TotalRequests)
+	}
+
+	if bulkResponse.SuccessfulPredictions != len(bulkResponse.Predictions) {
+		t.Errorf("Successful count (%d) doesn't match predictions length (%d)",
+			bulkResponse.SuccessfulPredictions, len(bulkResponse.Predictions))
+	}
+
+	// Validate each prediction in the response
+	for i, prediction := range bulkResponse.Predictions {
+		if prediction.TTFT <= 0 {
+			t.Errorf("Prediction %d: TTFT should be positive, got %.2f", i, prediction.TTFT)
+		}
+		if prediction.TPOT <= 0 {
+			t.Errorf("Prediction %d: TPOT should be positive, got %.2f", i, prediction.TPOT)
+		}
+		if prediction.ModelType == "" {
+			t.Errorf("Prediction %d: Model type should not be empty", i)
+		}
+
+		t.Logf("  Prediction %d: TTFT=%.2f, TPOT=%.2f, quantile=%.2f",
+			i+1, prediction.TTFT, prediction.TPOT, prediction.Quantile)
+	}
+
+	// Test performance expectation
+	avgTimePerPrediction := bulkResponse.ProcessingTimeMs / float64(bulkResponse.SuccessfulPredictions)
+	t.Logf("Average time per prediction: %.2f ms", avgTimePerPrediction)
+
+	if avgTimePerPrediction > 100 { // Bulk should be more efficient
+		t.Logf("Note: Bulk prediction averaging %.2f ms per request (may be acceptable)", avgTimePerPrediction)
+	} else {
+		t.Logf("✓ Good bulk prediction performance: %.2f ms per request", avgTimePerPrediction)
+	}
+}
+
+func testBulkPredictionsStrict(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing strict bulk predictions...")
+
+	if !predictor.IsReady() {
+		t.Skip("Predictor not ready for strict bulk prediction testing")
+	}
+
+	// Create valid prediction requests
+	requests := make([]PredictionRequest, 5)
+	for i := 0; i < 5; i++ {
+		requests[i] = PredictionRequest{
+			KVCachePercentage:  0.6,
+			InputTokenLength:   300 + i*100,
+			NumRequestWaiting:  i,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 50,
+			PrefixCacheScore:   float64(i) / 4.0, // 0.0 to 1.0
+		}
+	}
+
+	t.Logf("Making strict bulk prediction request with %d requests", len(requests))
+
+	bulkResponse, err := predictor.PredictBulkStrict(ctx, requests)
+	if err != nil {
+		t.Fatalf("Strict bulk prediction failed: %v", err)
+	}
+
+	t.Logf("Strict Bulk Prediction Response:")
+	t.Logf("  Total Requests: %d", bulkResponse.TotalRequests)
+	t.Logf("  Successful: %d", bulkResponse.SuccessfulPredictions)
+	t.Logf("  Failed: %d", bulkResponse.FailedPredictions)
+	t.Logf("  Processing Time: %.2f ms", bulkResponse.ProcessingTimeMs)
+
+	// In strict mode, we expect all requests to succeed or the entire batch to fail
+	if bulkResponse.FailedPredictions > 0 {
+		t.Errorf("Strict bulk prediction should not have partial failures, got %d failed",
+			bulkResponse.FailedPredictions)
+	}
+
+	if bulkResponse.SuccessfulPredictions != len(requests) {
+		t.Errorf("Expected all %d requests to succeed, got %d",
+			len(requests), bulkResponse.SuccessfulPredictions)
+	}
+
+	// Test bulk size limits
+	t.Log("Testing bulk size limits...")
+	largeRequests := make([]PredictionRequest, 150) // Over the limit
+	for i := range largeRequests {
+		largeRequests[i] = requests[0] // Use a valid request template
+	}
+
+	_, err = predictor.PredictBulkStrict(ctx, largeRequests)
+	if err == nil {
+		t.Error("Expected error for oversized bulk request, but got none")
+	} else {
+		t.Logf("✓ Correctly rejected oversized bulk request: %v", err)
+	}
+}
+
+func testLightGBMSupport(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing LightGBM support...")
+
+	currentModelType := predictor.GetCurrentModelType()
+	t.Logf("Current model type: %s", currentModelType)
+
+	if currentModelType == gbmModelType {
+		t.Log("Testing LightGBM-specific functionality...")
+
+		// Test LightGBM readiness
+		isReady := predictor.IsLightGBMReady()
+		t.Logf("LightGBM ready: %t", isReady)
+
+		if isReady {
+			// Test LightGBM prediction
+			req := PredictionRequest{
+				KVCachePercentage:  0.7,
+				InputTokenLength:   400,
+				NumRequestWaiting:  2,
+				NumRequestRunning:  1,
+				NumTokensGenerated: 60,
+				PrefixCacheScore:   0.8,
+			}
+
+			response, err := predictor.Predict(ctx, req)
+			if err != nil {
+				t.Errorf("LightGBM prediction failed: %v", err)
+			} else {
+				t.Logf("LightGBM prediction successful: TTFT=%.2f, TPOT=%.2f",
+					response.TTFT, response.TPOT)
+
+				if response.ModelType != gbmModelType {
+					t.Errorf("Expected model type 'lightgbm', got '%s'", response.ModelType)
+				}
+			}
+		} else {
+			t.Log("LightGBM not ready, skipping LightGBM-specific tests")
+		}
+	} else {
+		t.Logf("Current model type is %s, not LightGBM. LightGBM-specific tests skipped.", currentModelType)
+	}
+
+	// Test that the client handles all model types properly
+	t.Log("Verifying model type handling...")
+	switch currentModelType {
+	case bayesianRidgeModelType:
+		t.Log("✓ Bayesian Ridge model type recognized")
+	case xgBoostModelType:
+		t.Log("✓ XGBoost model type recognized")
+	case gbmModelType:
+		t.Log("✓ LightGBM model type recognized")
+	default:
+		t.Logf("⚠ Unknown model type: %s", currentModelType)
+	}
+}
+
+func testPredictionWithPrefixCache(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing prefix cache score impact on predictions...")
+
+	if !predictor.IsReady() {
+		t.Skip("Predictor not ready for prefix cache testing")
+	}
+
+	// Test with different prefix cache scores to see impact
+	baseRequest := PredictionRequest{
+		KVCachePercentage:  0.6,
+		InputTokenLength:   500,
+		NumRequestWaiting:  3,
+		NumRequestRunning:  2,
+		NumTokensGenerated: 75,
+	}
+
+	prefixCacheScores := []float64{0.0, 0.2, 0.4, 0.6, 0.8, 1.0}
+	ttftResults := make([]float64, 0, len(prefixCacheScores))
+	quantileResults := make([]float64, 0, len(prefixCacheScores))
+
+	for _, prefixScore := range prefixCacheScores {
+		req := baseRequest
+		req.PrefixCacheScore = prefixScore
+
+		response, err := predictor.Predict(ctx, req)
+		if err != nil {
+			t.Errorf("Prediction failed for prefix cache score %.1f: %v", prefixScore, err)
+			continue
+		}
+
+		ttftResults = append(ttftResults, response.TTFT)
+		quantileResults = append(quantileResults, response.Quantile)
+		t.Logf("Prefix cache %.0f%%: TTFT=%.2f ms, TPOT=%.2f ms, quantile=%.2f",
+			prefixScore*100, response.TTFT, response.TPOT, response.Quantile)
+	}
+
+	// Verify quantile consistency
+	if len(quantileResults) > 1 {
+		firstQuantile := quantileResults[0]
+		for i, q := range quantileResults {
+			if abs(q-firstQuantile) > 0.01 {
+				t.Errorf("Quantile inconsistency: prediction %d has quantile %.2f, expected %.2f",
+					i, q, firstQuantile)
+			}
+		}
+		t.Log("✓ Quantile values consistent across predictions")
+	}
+
+	// Analyze the relationship between prefix cache and TTFT
+	if len(ttftResults) >= 2 {
+		t.Log("Prefix cache impact analysis:")
+		lowCacheTTFT := ttftResults[0]                   // 0% prefix cache
+		highCacheTTFT := ttftResults[len(ttftResults)-1] // 100% prefix cache
+		difference := highCacheTTFT - lowCacheTTFT
+
+		t.Logf("  TTFT at 0%% prefix cache: %.2f ms", lowCacheTTFT)
+		t.Logf("  TTFT at 100%% prefix cache: %.2f ms", highCacheTTFT)
+		t.Logf("  Difference: %.2f ms", difference)
+
+		if predictor.GetCurrentModelType() == bayesianRidgeModelType {
+			// For Bayesian Ridge, we expect to see the linear relationship
+			if difference > 5 {
+				t.Logf("✓ Detected prefix cache impact: %.2f ms difference", difference)
+			} else {
+				t.Logf("ℹ Small prefix cache impact: %.2f ms difference", difference)
+			}
+		}
+	}
+}
+
+func testHTTPFallbackPrediction(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing HTTP fallback prediction...")
+
+	modelType := predictor.GetCurrentModelType()
+	if modelType == bayesianRidgeModelType {
+		t.Skip("HTTP fallback test not applicable for Bayesian Ridge")
+	}
+
+	// Test prediction with HTTP fallback including prefix cache score
+	req := PredictionRequest{
+		KVCachePercentage:  0.8, // 80% as a fraction
+		InputTokenLength:   1024,
+		NumRequestWaiting:  5,
+		NumRequestRunning:  3,
+		NumTokensGenerated: 150,
+		PrefixCacheScore:   0.9, // 90% prefix cache hit rate
+	}
+
+	t.Logf("Making HTTP prediction request: %+v", req)
+
+	response, err := predictor.Predict(ctx, req)
+	if err != nil {
+		t.Fatalf("HTTP prediction failed: %v", err)
+	}
+
+	t.Logf("HTTP Prediction Response:")
+	t.Logf("  TTFT: %.2f ms", response.TTFT)
+	t.Logf("  TPOT: %.2f ms", response.TPOT)
+	t.Logf("  Model Type: %s", response.ModelType)
+	t.Logf("  Quantile: %.2f", response.Quantile)
+	t.Logf("  Prefix Cache Score Used: %.1f%%", req.PrefixCacheScore*100)
+
+	// Validate that we got a reasonable response
+	if response.TTFT <= 0 {
+		t.Error("TTFT should be positive")
+	}
+	if response.TPOT <= 0 {
+		t.Error("TPOT should be positive")
+	}
+
+	// The model type should indicate the correct type
+	if response.ModelType == "" {
+		t.Error("Model type should not be empty")
+	}
+
+	t.Logf("Successfully tested HTTP prediction with prefix cache")
+}
+
+func testPredictionPerformance(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing prediction performance (target: < 300ms) with prefix cache scores...")
+
+	// Ensure predictor is ready
+	if !predictor.IsReady() {
+		t.Skip("Predictor not ready for performance test")
+	}
+
+	req := PredictionRequest{
+		KVCachePercentage:  0.6, // 60% as a fraction
+		InputTokenLength:   768,
+		NumRequestWaiting:  2,
+		NumRequestRunning:  1,
+		NumTokensGenerated: 80,
+		PrefixCacheScore:   0.7, // 70% prefix cache hit rate
+	}
+
+	// Warm up with a few predictions
+	for i := 0; i < 3; i++ {
+		_, err := predictor.Predict(ctx, req)
+		if err != nil {
+			t.Fatalf("Warmup prediction %d failed: %v", i+1, err)
+		}
+	}
+
+	// Test multiple predictions and measure time
+	const numTests = 10
+	const avgDurationMs = 250
+
+	var totalDuration time.Duration
+	var maxSingleDuration time.Duration
+	var minSingleDuration = time.Hour // Initialize to large value
+
+	t.Logf("Running %d prediction performance tests...", numTests)
+
+	for i := 0; i < numTests; i++ {
+		// Vary prefix cache score for each test
+		testReq := req
+		testReq.PrefixCacheScore = float64(i) / float64(numTests-1) // 0.0 to 1.0
+
+		start := time.Now()
+
+		response, err := predictor.Predict(ctx, testReq)
+
+		duration := time.Since(start)
+		totalDuration += duration
+
+		if err != nil {
+			t.Errorf("Prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		// Track min/max durations
+		if duration > maxSingleDuration {
+			maxSingleDuration = duration
+		}
+		if duration < minSingleDuration {
+			minSingleDuration = duration
+		}
+
+		durationMs := float64(duration.Nanoseconds()) / 1e6
+		t.Logf("Prediction %d: %.2fms - TTFT: %.1fms, TPOT: %.1fms (prefix: %.0f%%, quantile: %.2f)",
+			i+1, durationMs, response.TTFT, response.TPOT, testReq.PrefixCacheScore*100, response.Quantile)
+	}
+
+	// Calculate statistics
+	avgDuration := totalDuration / numTests
+	avgMs := float64(avgDuration.Nanoseconds()) / 1e6
+	maxMs := float64(maxSingleDuration.Nanoseconds()) / 1e6
+	minMs := float64(minSingleDuration.Nanoseconds()) / 1e6
+
+	t.Logf("Performance Results:")
+	t.Logf("  Average: %.2fms", avgMs)
+	t.Logf("  Minimum: %.2fms", minMs)
+	t.Logf("  Maximum: %.2fms", maxMs)
+	t.Logf("  Target:  < %dms", avgDurationMs)
+
+	// Overall performance check
+	if avgMs > avgDurationMs {
+		t.Errorf("Average prediction time %.2fms exceeded target of %dms", avgMs, avgDurationMs)
+	} else {
+		t.Logf("✅ Performance target met: avg %.2fms < %dms", avgMs, avgDurationMs)
+	}
+
+	// Check for consistency (max shouldn't be too much higher than average)
+	if maxMs > avgMs*3 {
+		t.Logf("⚠️  High variance detected: max %.2fms is %.1fx the average", maxMs, maxMs/avgMs)
+	} else {
+		t.Logf("✅ Good consistency: max %.2fms is %.1fx the average", maxMs, maxMs/avgMs)
+	}
+}
+
+func testBulkPredictionPerformance(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing bulk prediction performance...")
+
+	if !predictor.IsReady() {
+		t.Skip("Predictor not ready for bulk performance test")
+	}
+
+	// Create batch of prediction requests
+	const batchSize = 20
+	requests := make([]PredictionRequest, batchSize)
+	for i := 0; i < batchSize; i++ {
+		requests[i] = PredictionRequest{
+			KVCachePercentage:  0.6 + float64(i%5)*0.05,           // Vary between 0.6 and 0.8
+			InputTokenLength:   300 + i*10,                        // Vary input length
+			NumRequestWaiting:  i % 4,                             // 0 to 3
+			NumRequestRunning:  (i % 2) + 1,                       // 1 to 2
+			NumTokensGenerated: 50 + i*2,                          // Vary generated tokens
+			PrefixCacheScore:   float64(i) / float64(batchSize-1), // 0.0 to 1.0
+		}
+	}
+
+	// Warm up
+	warmupRequests := requests[:3]
+	_, err := predictor.PredictBulk(ctx, warmupRequests)
+	if err != nil {
+		t.Fatalf("Warmup bulk prediction failed: %v", err)
+	}
+
+	// Performance test
+	const numTests = 5
+	var totalDuration time.Duration
+	var totalRequests int
+	var totalSuccessful int
+
+	t.Logf("Running %d bulk prediction performance tests with %d requests each...", numTests, batchSize)
+
+	for i := 0; i < numTests; i++ {
+		start := time.Now()
+
+		response, err := predictor.PredictBulk(ctx, requests)
+
+		duration := time.Since(start)
+		totalDuration += duration
+
+		if err != nil {
+			t.Errorf("Bulk prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		totalRequests += response.TotalRequests
+		totalSuccessful += response.SuccessfulPredictions
+
+		durationMs := float64(duration.Nanoseconds()) / 1e6
+		avgPerRequest := durationMs / float64(response.SuccessfulPredictions)
+
+		t.Logf("Bulk test %d: %.2fms total, %.2fms per request (%d/%d successful)",
+			i+1, durationMs, avgPerRequest, response.SuccessfulPredictions, response.TotalRequests)
+	}
+
+	// Calculate bulk performance statistics
+	avgTotalDuration := totalDuration / numTests
+	avgTotalMs := float64(avgTotalDuration.Nanoseconds()) / 1e6
+	avgPerRequest := avgTotalMs / float64(batchSize)
+
+	t.Logf("Bulk Performance Results:")
+	t.Logf("  Average total time: %.2fms", avgTotalMs)
+	t.Logf("  Average per request: %.2fms", avgPerRequest)
+	t.Logf("  Success rate: %.1f%%", float64(totalSuccessful)/float64(totalRequests)*100)
+
+	// Compare with single prediction performance target
+	singlePredictionTarget := 250.0                         // ms
+	bulkEfficiencyThreshold := singlePredictionTarget * 0.7 // Bulk should be more efficient
+
+	switch {
+	case avgPerRequest <= bulkEfficiencyThreshold:
+		t.Logf("✅ Bulk predictions are efficient: %.2fms per request < %.2fms threshold",
+			avgPerRequest, bulkEfficiencyThreshold)
+	case avgPerRequest <= singlePredictionTarget:
+		t.Logf("✓ Bulk predictions acceptable: %.2fms per request < %.2fms single target",
+			avgPerRequest, singlePredictionTarget)
+	default:
+		t.Errorf("❌ Bulk predictions slow: %.2fms per request > %.2fms target",
+			avgPerRequest, singlePredictionTarget)
+	}
+}
+
+func testHTTPOnlyPerformance(t *testing.T, ctx context.Context) {
+	t.Log("Testing HTTP-only prediction performance (no native XGBoost interference) with prefix cache...")
+
+	predictionURLs := os.Getenv("PREDICTION_SERVER_URL")
+	trainingURL := os.Getenv("TRAINING_SERVER_URL")
+	if predictionURLs == "" {
+		t.Skip("PREDICTION_SERVER_URL not set")
+	}
+	if trainingURL == "" {
+		// Use first prediction URL as fallback
+		urls := strings.Split(predictionURLs, ",")
+		if len(urls) > 0 {
+			trainingURL = strings.TrimSpace(urls[0])
+		} else {
+			t.Skip("No valid URLs available for testing")
+		}
+	}
+
+	// Parse prediction URLs
+	urls := strings.Split(predictionURLs, ",")
+	parsedPredictionURLs := make([]string, 0, len(urls))
+	for _, url := range urls {
+		parsedPredictionURLs = append(parsedPredictionURLs, strings.TrimSpace(url))
+	}
+
+	// Create a dedicated HTTP-only predictor for clean performance testing
+	zapLog, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatalf("Failed to create logger: %v", err)
+	}
+	logger := zapr.NewLogger(zapLog)
+
+	httpOnlyConfig := &Config{
+		TrainingURL:            trainingURL,
+		PredictionURLs:         parsedPredictionURLs,
+		MaxSampleSize:          1000,
+		FlushInterval:          1 * time.Second, // Long interval to avoid interference
+		MetricsRefreshInterval: 1 * time.Second, // Longer for metrics
+		UseNativeXGBoost:       false,           // Force HTTP-only
+		HTTPTimeout:            5 * time.Second, // Reasonable timeout
+		MaxBulkSize:            50,
+	}
+
+	httpPredictor := New(httpOnlyConfig, logger)
+	defer httpPredictor.Stop()
+
+	err = httpPredictor.Start(ctx)
+	if err != nil {
+		t.Fatalf("Failed to start HTTP-only predictor: %v", err)
+	}
+
+	// Wait for readiness
+	time.Sleep(1 * time.Second)
+
+	// Wait for coefficients to be cached
+	maxWaitTime := 10 * time.Second
+	waitInterval := 200 * time.Millisecond
+	elapsed := time.Duration(0)
+
+	for elapsed < maxWaitTime {
+		if httpPredictor.IsReady() {
+			break
+		}
+		time.Sleep(waitInterval)
+		elapsed += waitInterval
+	}
+
+	if !httpPredictor.IsReady() {
+		t.Skip("model not ready yet")
+	}
+
+	req := PredictionRequest{
+		KVCachePercentage:  0.65,
+		InputTokenLength:   512,
+		NumRequestWaiting:  1,
+		NumRequestRunning:  2,
+		NumTokensGenerated: 100,
+		PrefixCacheScore:   0.75, // 75% prefix cache hit rate
+	}
+
+	// Warm up
+	for i := 0; i < 2; i++ {
+		_, err := httpPredictor.Predict(ctx, req)
+		if err != nil {
+			t.Fatalf("HTTP warmup prediction %d failed: %v", i+1, err)
+		}
+	}
+
+	// Performance test
+	const numTests = 15
+	const targetMs = 250
+
+	var durations []time.Duration
+	var successful int
+
+	t.Logf("Running %d HTTP-only prediction tests...", numTests)
+
+	for i := 0; i < numTests; i++ {
+		// Vary prefix cache for each test
+		testReq := req
+		testReq.PrefixCacheScore = 0.5 + (float64(i)/float64(numTests-1))*0.5 // 0.5 to 1.0
+
+		start := time.Now()
+
+		response, err := httpPredictor.Predict(ctx, testReq)
+
+		duration := time.Since(start)
+		durations = append(durations, duration)
+
+		if err != nil {
+			t.Errorf("HTTP prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		successful++
+		durationMs := float64(duration.Nanoseconds()) / 1e6
+
+		status := "✅"
+
+		t.Logf("%s Test %d: %.1fms (TTFT: %.0fms, TPOT: %.0fms, prefix: %.0f%%, quantile: %.2f)",
+			status, i+1, durationMs, response.TTFT, response.TPOT, testReq.PrefixCacheScore*100, response.Quantile)
+	}
+
+	// Calculate statistics
+	if len(durations) == 0 {
+		t.Fatal("No successful predictions to analyze")
+	}
+
+	var total time.Duration
+	min := durations[0]
+	max := durations[0]
+
+	for _, d := range durations {
+		total += d
+		if d < min {
+			min = d
+		}
+		if d > max {
+			max = d
+		}
+	}
+
+	avg := total / time.Duration(len(durations))
+	avgMs := float64(avg.Nanoseconds()) / 1e6
+	minMs := float64(min.Nanoseconds()) / 1e6
+	maxMs := float64(max.Nanoseconds()) / 1e6
+
+	// Count fast predictions
+	fastCount := 0
+	for _, d := range durations {
+		if float64(d.Nanoseconds())/1e6 <= targetMs {
+			fastCount++
+		}
+	}
+
+	t.Logf("\n📊 HTTP-Only Performance Summary:")
+	t.Logf("  Success Rate: %d/%d (%.1f%%)", successful, numTests, float64(successful)/float64(numTests)*100)
+	t.Logf("  Average: %.1fms", avgMs)
+	t.Logf("  Minimum: %.1fms", minMs)
+	t.Logf("  Maximum: %.1fms", maxMs)
+	t.Logf("  Under %dms: %d/%d (%.1f%%)", targetMs, fastCount, len(durations), float64(fastCount)/float64(len(durations))*100)
+
+	// Performance assertions
+	if successful < numTests {
+		t.Errorf("Some predictions failed: %d/%d successful", successful, numTests)
+	}
+
+	if avgMs <= targetMs {
+		t.Logf("✅ PASS: Average response time %.1fms ≤ %dms target", avgMs, targetMs)
+	} else {
+		t.Errorf("❌ FAIL: Average response time %.1fms > %dms target", avgMs, targetMs)
+	}
+
+	// Check that at least 80% of requests are under target
+	fastPercentage := float64(fastCount) / float64(len(durations)) * 100
+	if fastPercentage >= 80 {
+		t.Logf("✅ PASS: %.1f%% of requests under %dms (≥80%% target)", fastPercentage, targetMs)
+	} else {
+		t.Errorf("❌ FAIL: Only %.1f%% of requests under %dms (<80%% target)", fastPercentage, targetMs)
+	}
+}
+
+func testHTTPOnlyPrediction(t *testing.T, ctx context.Context) {
+	t.Log("Testing HTTP-only prediction (bypassing native XGBoost) with prefix cache...")
+
+	// Create a predictor with native XGBoost disabled to force HTTP usage
+	predictionURLs := os.Getenv("PREDICTION_SERVER_URL")
+	trainingURL := os.Getenv("TRAINING_SERVER_URL")
+	if predictionURLs == "" {
+		t.Skip("PREDICTION_SERVER_URL not set")
+	}
+	if trainingURL == "" {
+		// Use first prediction URL as fallback
+		urls := strings.Split(predictionURLs, ",")
+		if len(urls) > 0 {
+			trainingURL = strings.TrimSpace(urls[0])
+		} else {
+			t.Skip("No valid URLs available for testing")
+		}
+	}
+
+	// Parse prediction URLs
+	urls := strings.Split(predictionURLs, ",")
+	parsedPredictionURLs := make([]string, 0, len(urls))
+	for _, url := range urls {
+		parsedPredictionURLs = append(parsedPredictionURLs, strings.TrimSpace(url))
+	}
+
+	zapLog, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatalf("Failed to create logger: %v", err)
+	}
+	logger := zapr.NewLogger(zapLog)
+
+	httpOnlyConfig := &Config{
+		TrainingURL:            trainingURL,
+		PredictionURLs:         parsedPredictionURLs,
+		MaxSampleSize:          1000,
+		FlushInterval:          1 * time.Second,
+		MetricsRefreshInterval: 1 * time.Second, // Longer for metrics
+		UseNativeXGBoost:       false,           // Force HTTP fallback
+		HTTPTimeout:            30 * time.Second,
+		MaxBulkSize:            25,
+	}
+
+	httpPredictor := New(httpOnlyConfig, logger)
+	defer httpPredictor.Stop()
+
+	err = httpPredictor.Start(ctx)
+	if err != nil {
+		t.Fatalf("Failed to start HTTP-only predictor: %v", err)
+	}
+
+	// Wait a moment for startup and coefficient caching
+	time.Sleep(3 * time.Second)
+
+	// Ensure coefficients are ready
+	maxWait := 10 * time.Second
+	waited := time.Duration(0)
+	for waited < maxWait {
+		if httpPredictor.IsReady() {
+			break
+		}
+		time.Sleep(500 * time.Millisecond)
+		waited += 500 * time.Millisecond
+	}
+
+	if !httpPredictor.IsReady() {
+		t.Skip("Model not ready yet")
+	}
+
+	// Test prediction using HTTP only with prefix cache
+	req := PredictionRequest{
+		KVCachePercentage:  0.6, // 60% as a fraction
+		InputTokenLength:   256,
+		NumRequestWaiting:  1,
+		NumRequestRunning:  2,
+		NumTokensGenerated: 75,
+		PrefixCacheScore:   0.85, // 85% prefix cache hit rate
+	}
+
+	t.Logf("Making HTTP-only prediction request: %+v", req)
+
+	response, err := httpPredictor.Predict(ctx, req)
+	if err != nil {
+		t.Fatalf("HTTP-only prediction failed: %v", err)
+	}
+
+	t.Logf("HTTP-Only Prediction Response:")
+	t.Logf("  TTFT: %.2f ms", response.TTFT)
+	t.Logf("  TPOT: %.2f ms", response.TPOT)
+	t.Logf("  Model Type: %s", response.ModelType)
+	t.Logf("  Quantile: %.2f", response.Quantile)
+	t.Logf("  TTFT Uncertainty: %.2f", response.TTFTUncertainty)
+	t.Logf("  TPOT Uncertainty: %.2f", response.TPOTUncertainty)
+	t.Logf("  Prefix Cache Score Used: %.1f%%", req.PrefixCacheScore*100)
+
+	// Validate response
+	if response.TTFT <= 0 {
+		t.Error("TTFT should be positive")
+	}
+	if response.TPOT <= 0 {
+		t.Error("TPOT should be positive")
+	}
+
+	// Test multiple HTTP-only predictions with varying prefix cache
+	t.Log("Testing multiple HTTP-only predictions with different prefix cache scores...")
+	for i := 0; i < 3; i++ {
+		testReq := PredictionRequest{
+			KVCachePercentage:  float64(30+i*20) / 100.0,
+			InputTokenLength:   128 + i*256,
+			NumRequestWaiting:  i,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 25 + i*50,
+			PrefixCacheScore:   float64(60+i*20) / 100.0, // 60%, 80%, 100%
+		}
+
+		resp, err := httpPredictor.Predict(ctx, testReq)
+		if err != nil {
+			t.Errorf("HTTP-only prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		t.Logf("HTTP-only prediction %d: TTFT=%.2f, TPOT=%.2f (prefix: %.0f%%, quantile: %.2f)",
+			i+1, resp.TTFT, resp.TPOT, testReq.PrefixCacheScore*100, resp.Quantile)
+	}
+
+	t.Log("Successfully tested HTTP-only predictions with prefix cache")
+}
+
+func testLoadBalancing(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing load balancing across multiple prediction URLs with prefix cache...")
+
+	predictionURLs := predictor.GetPredictionURLs()
+	if len(predictionURLs) <= 1 {
+		t.Skip("Need multiple prediction URLs to test load balancing")
+	}
+
+	t.Logf("Testing load balancing across %d prediction URLs: %v", len(predictionURLs), predictionURLs)
+
+	// Make multiple predictions to test load balancing
+	const numPredictions = 20
+	req := PredictionRequest{
+		KVCachePercentage:  0.7,
+		InputTokenLength:   512,
+		NumRequestWaiting:  2,
+		NumRequestRunning:  1,
+		NumTokensGenerated: 100,
+		PrefixCacheScore:   0.8, // 80% prefix cache hit rate
+	}
+
+	successfulPredictions := 0
+	for i := 0; i < numPredictions; i++ {
+		// Vary prefix cache score across requests
+		testReq := req
+		testReq.PrefixCacheScore = 0.5 + (float64(i)/float64(numPredictions-1))*0.5 // 0.5 to 1.0
+
+		response, err := predictor.Predict(ctx, testReq)
+		if err != nil {
+			t.Logf("Prediction %d failed: %v", i+1, err)
+			continue
+		}
+
+		successfulPredictions++
+		t.Logf("Prediction %d: TTFT=%.2f, TPOT=%.2f (prefix: %.0f%%, quantile: %.2f)",
+			i+1, response.TTFT, response.TPOT, testReq.PrefixCacheScore*100, response.Quantile)
+	}
+
+	successRate := float64(successfulPredictions) / float64(numPredictions) * 100
+	t.Logf("Load balancing test results: %d/%d successful (%.1f%%)", successfulPredictions, numPredictions, successRate)
+
+	if successRate < 80 {
+		t.Errorf("Low success rate in load balancing test: %.1f%% < 80%%", successRate)
+	} else {
+		t.Logf("✅ Load balancing test successful with %.1f%% success rate", successRate)
+	}
+}
+
+func testPrefixCacheValidation(t *testing.T, predictor *Predictor) {
+	t.Log("Testing prefix cache score validation...")
+
+	// Test valid prefix cache scores
+	validScores := []float64{0.0, 0.25, 0.5, 0.75, 1.0}
+	for _, score := range validScores {
+		req := PredictionRequest{
+			KVCachePercentage:  0.5,
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   score,
+		}
+
+		err := predictor.ValidatePredictionRequest(req)
+		if err != nil {
+			t.Errorf("Valid prefix cache score %.2f should not cause validation error: %v", score, err)
+		}
+	}
+
+	// Test invalid prefix cache scores
+	invalidScores := []float64{-0.1, -1.0, 1.1, 2.0}
+	for _, score := range invalidScores {
+		req := PredictionRequest{
+			KVCachePercentage:  0.5,
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   score,
+		}
+
+		err := predictor.ValidatePredictionRequest(req)
+		if err == nil {
+			t.Errorf("Invalid prefix cache score %.2f should cause validation error", score)
+		} else {
+			t.Logf("✓ Invalid prefix cache score %.2f correctly rejected: %v", score, err)
+		}
+	}
+
+	// Test training entry validation
+	validEntry := TrainingEntry{
+		KVCachePercentage:  0.6,
+		InputTokenLength:   200,
+		NumRequestWaiting:  2,
+		NumRequestRunning:  1,
+		NumTokensGenerated: 20,
+		ActualTTFT:         50.0,
+		ActualTPOT:         15.0,
+		PrefixCacheScore:   0.8,
+		Timestamp:          time.Now(),
+	}
+
+	err := predictor.ValidateTrainingEntry(validEntry)
+	if err != nil {
+		t.Errorf("Valid training entry should not cause validation error: %v", err)
+	}
+
+	// Test invalid training entry
+	invalidEntry := validEntry
+	invalidEntry.PrefixCacheScore = 1.5 // Invalid
+
+	err = predictor.ValidateTrainingEntry(invalidEntry)
+	if err == nil {
+		t.Error("Invalid training entry should cause validation error")
+	} else {
+		t.Logf("✓ Invalid training entry correctly rejected: %v", err)
+	}
+
+	t.Log("✅ Prefix cache validation tests completed")
+}
+
+func testPredictionConstructors(t *testing.T) {
+	t.Log("Testing prediction and training entry constructors with prefix cache...")
+
+	// Test valid prediction request constructor
+	req, err := NewPredictionRequest(
+		0.7,  // kv_cache_percentage
+		500,  // input_token_length
+		3,    // num_request_waiting
+		2,    // num_request_running
+		100,  // num_tokens_generated
+		0.85, // prefix_cache_score
+	)
+	if err != nil {
+		t.Errorf("Valid prediction request constructor failed: %v", err)
+	} else {
+		t.Logf("✓ Created prediction request: TTFT features with %.0f%% prefix cache", req.PrefixCacheScore*100)
+	}
+
+	// Test invalid prediction request constructor
+	_, err = NewPredictionRequest(
+		0.7, // kv_cache_percentage
+		500, // input_token_length
+		3,   // num_request_waiting
+		2,   // num_request_running
+		100, // num_tokens_generated
+		1.5, // prefix_cache_score (invalid)
+	)
+	if err == nil {
+		t.Error("Invalid prediction request constructor should have failed")
+	} else {
+		t.Logf("✓ Invalid prediction request correctly rejected: %v", err)
+	}
+
+	// Test valid training entry constructor
+	entry, err := NewTrainingEntry(
+		0.6,  // kv_cache_percentage
+		300,  // input_token_length
+		2,    // num_request_waiting
+		1,    // num_request_running
+		50,   // num_tokens_generated
+		45.5, // actual_ttft_ms
+		12.3, // actual_tpot_ms
+		0.75, // prefix_cache_score
+	)
+	if err != nil {
+		t.Errorf("Valid training entry constructor failed: %v", err)
+	} else {
+		t.Logf("✓ Created training entry: TTFT=%.1fms, TPOT=%.1fms, prefix cache=%.0f%%",
+			entry.ActualTTFT, entry.ActualTPOT, entry.PrefixCacheScore*100)
+	}
+
+	// Test invalid training entry constructor
+	_, err = NewTrainingEntry(
+		0.6,  // kv_cache_percentage
+		300,  // input_token_length
+		2,    // num_request_waiting
+		1,    // num_request_running
+		50,   // num_tokens_generated
+		45.5, // actual_ttft_ms
+		12.3, // actual_tpot_ms
+		-0.1, // prefix_cache_score (invalid)
+	)
+	if err == nil {
+		t.Error("Invalid training entry constructor should have failed")
+	} else {
+		t.Logf("✓ Invalid training entry correctly rejected: %v", err)
+	}
+
+	t.Log("✅ Constructor validation tests completed")
+}
+
+func testXGBoostJSONStructure(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing XGBoost JSON structure from server...")
+
+	if predictor.GetCurrentModelType() != xgBoostModelType {
+		t.Skip("This test is specific to XGBoost model type")
+	}
+
+	// Get raw trees to examine structure
+	trees, err := predictor.GetXGBoostTrees(ctx)
+	if err != nil {
+		t.Fatalf("Failed to get XGBoost trees: %v", err)
+	}
+
+	if len(trees.TTFTTrees) == 0 {
+		t.Fatal("No TTFT trees available")
+	}
+
+	// Examine the first tree structure
+	firstTree := trees.TTFTTrees[0]
+	t.Logf("First TTFT tree structure: %T", firstTree)
+
+	// Convert to map to examine fields
+	if treeMap, ok := firstTree.(map[string]interface{}); ok {
+		t.Log("First tree fields:")
+		for key, value := range treeMap {
+			switch key {
+			case "split":
+				t.Logf("  %s: %T = %v", key, value, value)
+			case "children":
+				if value != nil {
+					if children, ok := value.([]interface{}); ok {
+						// This is the unique execution path for valid children
+						t.Logf("  %s: []interface{} with %d children", key, len(children))
+						// Examine first child
+						if len(children) > 0 {
+							if childMap, ok := children[0].(map[string]interface{}); ok {
+								for childKey, childValue := range childMap {
+									if childKey == "split" {
+										t.Logf("    child[0].%s: %T = %v", childKey, childValue, childValue)
+									}
+								}
+							}
+						}
+					} else {
+						// Fallback if type assertion fails
+						t.Logf("  %s: %T = %v", key, value, value)
+					}
+				} else {
+					// Fallback if value is nil
+					t.Logf("  %s: %T = %v", key, value, value)
+				}
+			default:
+				t.Logf("  %s: %T = %v", key, value, value)
+			}
+		}
+	}
+
+	// Try to understand why the conversion is failing
+	t.Log("Analyzing conversion issue...")
+	if len(trees.TTFTTrees) > 0 {
+		// Test the conversion function manually
+		testConvertXGBoostJSON(t, trees.TTFTTrees[0])
+	}
+
+	t.Log("XGBoost JSON structure analysis complete")
+}
+
+// Helper function to test the conversion logic
+func testConvertXGBoostJSON(t *testing.T, tree interface{}) {
+	featureMap := map[string]int{
+		"kv_cache_percentage":  0,
+		"input_token_length":   1,
+		"num_request_waiting":  2,
+		"num_request_running":  3,
+		"num_tokens_generated": 4,
+		"prefix_cache_score":   5, // Added prefix cache score mapping
+	}
+
+	t.Log("Testing XGBoost JSON conversion...")
+
+	treeMap, ok := tree.(map[string]interface{})
+	if !ok {
+		t.Log("Tree is not a map[string]interface{}")
+		return
+	}
+
+	// Check if split field exists and what type it is
+	if split, exists := treeMap["split"]; exists {
+		t.Logf("Split field exists: %T = %v", split, split)
+
+		switch splitVal := split.(type) {
+		case string:
+			t.Logf("Split is string: '%s'", splitVal)
+			if featureIdx, found := featureMap[splitVal]; found {
+				t.Logf("Found feature index for '%s': %d", splitVal, featureIdx)
+			} else {
+				t.Logf("Feature '%s' not found in feature map", splitVal)
+			}
+		case float64:
+			t.Logf("Split is float64: %v (already numeric, no conversion needed)", splitVal)
+		case int:
+			t.Logf("Split is int: %v (already numeric, no conversion needed)", splitVal)
+		default:
+			t.Logf("Split is unexpected type: %T = %v", splitVal, splitVal)
+		}
+	} else {
+		t.Log("Split field does not exist")
+	}
+}
+
+func testMetricsRetrieval(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing metrics retrieval...")
+
+	modelType := predictor.GetCurrentModelType()
+	quantile := predictor.GetCurrentQuantile()
+	t.Logf("Testing metrics for model type: %s, quantile: %.2f", modelType, quantile)
+
+	switch modelType {
+	case bayesianRidgeModelType:
+		testBayesianRidgeMetrics(t, ctx, predictor)
+	case xgBoostModelType:
+		testXGBoostMetrics(t, ctx, predictor)
+	case gbmModelType:
+		testLightGBMMetrics(t, ctx, predictor)
+	default:
+		t.Logf("Unknown model type %s, testing cached metrics only", modelType)
+	}
+
+	// Test cached metrics
+	cachedMetrics, hasCached := predictor.GetCachedMetrics()
+	if hasCached {
+		t.Logf("Cached metrics available - Model Type: %s", cachedMetrics.ModelType)
+		if len(cachedMetrics.RawMetrics) > 0 {
+			t.Logf("Raw metrics length: %d characters", len(cachedMetrics.RawMetrics))
+		}
+	} else {
+		t.Log("No cached metrics available")
+	}
+
+	// Test readiness status
+	t.Logf("Predictor readiness status:")
+	t.Logf("  Overall Ready: %t", predictor.IsReady())
+	t.Logf("  XGBoost Ready: %t", predictor.IsXGBoostReady())
+	t.Logf("  LightGBM Ready: %t", predictor.IsLightGBMReady())
+	t.Logf("  Bayesian Ridge Ready: %t", predictor.IsBayesianRidgeReady())
+}
+
+func testBayesianRidgeMetrics(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing Bayesian Ridge specific metrics with prefix cache support...")
+
+	metrics, err := predictor.GetMetrics(ctx)
+	if err != nil {
+		t.Errorf("Failed to get Bayesian Ridge metrics: %v", err)
+		return
+	}
+
+	if metrics.Coefficients == nil {
+		t.Error("Bayesian Ridge coefficients should not be nil")
+		return
+	}
+
+	t.Logf("TTFT Coefficients (should include prefix_cache_score):")
+	t.Logf("  Intercept: %.6f", metrics.Coefficients.TTFTIntercept)
+	for feature, coeff := range metrics.Coefficients.TTFTCoeffs {
+		t.Logf("  %s: %.6f", feature, coeff)
+	}
+
+	t.Logf("TPOT Coefficients (should NOT include prefix_cache_score):")
+	t.Logf("  Intercept: %.6f", metrics.Coefficients.TPOTIntercept)
+	for feature, coeff := range metrics.Coefficients.TPOTCoeffs {
+		t.Logf("  %s: %.6f", feature, coeff)
+	}
+
+	// Validate prefix cache score is in TTFT but not TPOT
+	if _, hasPrefixCache := metrics.Coefficients.TTFTCoeffs["prefix_cache_score"]; hasPrefixCache {
+		t.Log("✓ TTFT model includes prefix_cache_score coefficient")
+	} else {
+		t.Log("ℹ TTFT model does not include prefix_cache_score coefficient (may not be trained yet)")
+	}
+
+	if _, hasPrefixCache := metrics.Coefficients.TPOTCoeffs["prefix_cache_score"]; hasPrefixCache {
+		t.Error("❌ TPOT model should NOT include prefix_cache_score coefficient")
+	} else {
+		t.Log("✓ TPOT model correctly excludes prefix_cache_score coefficient")
+	}
+
+	// Test individual coefficient and bucket retrieval
+	coeffs, err := predictor.GetModelCoefficients(ctx)
+	if err != nil {
+		t.Errorf("Failed to get model coefficients: %v", err)
+	} else {
+		t.Logf("Retrieved coefficients separately: %d TTFT, %d TPOT features",
+			len(coeffs.TTFTCoeffs), len(coeffs.TPOTCoeffs))
+	}
+
+	buckets, err := predictor.GetBucketCounts(ctx)
+	if err != nil {
+		t.Errorf("Failed to get bucket counts: %v", err)
+	} else {
+		t.Logf("Retrieved bucket counts: %d TTFT, %d TPOT buckets",
+			len(buckets.TTFTBuckets), len(buckets.TPOTBuckets))
+	}
+}
+
+func testXGBoostMetrics(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing XGBoost specific metrics...")
+
+	// Wait a bit for XGBoost models to potentially load
+	time.Sleep(3 * time.Second)
+
+	trees, err := predictor.GetXGBoostTrees(ctx)
+	if err != nil {
+		t.Errorf("Failed to get XGBoost trees: %v", err)
+		return
+	}
+
+	t.Logf("XGBoost Trees:")
+	t.Logf("  TTFT Trees: %d", len(trees.TTFTTrees))
+	t.Logf("  TPOT Trees: %d", len(trees.TPOTTrees))
+
+	if len(trees.TTFTTrees) == 0 {
+		t.Error("Expected at least one TTFT tree")
+	}
+	if len(trees.TPOTTrees) == 0 {
+		t.Error("Expected at least one TPOT tree")
+	}
+
+	// Test native XGBoost readiness
+	if predictor.IsXGBoostReady() {
+		t.Log("Native XGBoost models are ready for local prediction")
+	} else {
+		t.Log("Native XGBoost models not ready, will use HTTP fallback")
+	}
+}
+
+func testLightGBMMetrics(t *testing.T, ctx context.Context, predictor *Predictor) {
+	t.Log("Testing LightGBM specific metrics...")
+
+	// For LightGBM, we primarily use HTTP calls, so test the HTTP connectivity
+	if predictor.IsLightGBMReady() {
+		t.Log("LightGBM models are ready via HTTP")
+
+		// Test a simple prediction to ensure the HTTP endpoint works
+		req := PredictionRequest{
+			KVCachePercentage:  0.6,
+			InputTokenLength:   300,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 50,
+			PrefixCacheScore:   0.7,
+		}
+
+		_, err := predictor.Predict(ctx, req)
+		if err != nil {
+			t.Errorf("LightGBM test prediction failed: %v", err)
+		} else {
+			t.Log("✓ LightGBM HTTP prediction working")
+		}
+	} else {
+		t.Log("LightGBM models not ready")
+	}
+}
+
+// generateTrainingEntries creates random training data for testing with prefix cache scores
+func generateTrainingEntries(count int) []TrainingEntry {
+	entries := make([]TrainingEntry, count)
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	for i := 0; i < count; i++ {
+		// Generate TTFT and TPOT using a simple equation based on features, plus some noise
+		kv := rng.Float64() // 0.0 to 1.0
+		inputLen := rng.Intn(2048) + 1
+		waiting := rng.Intn(20)
+		running := rng.Intn(10) + 1
+		generated := rng.Intn(500) + 1
+		prefixCache := rng.Float64() // 0.0 to 1.0
+
+		// Updated equations to include prefix cache impact on TTFT:
+		// TTFT includes prefix cache, TPOT does not
+		ttft := 100 + 2*float64(inputLen) + 10*kv + 5*float64(waiting) + 30*prefixCache + rng.NormFloat64()*20
+		tpot := 20 + 0.5*float64(generated) + 2*float64(running) + rng.NormFloat64()*5 + 9*kv
+
+		entries[i] = TrainingEntry{
+			KVCachePercentage:  kv,
+			InputTokenLength:   inputLen,
+			NumRequestWaiting:  waiting,
+			NumRequestRunning:  running,
+			NumTokensGenerated: generated,
+			ActualTTFT:         ttft,
+			ActualTPOT:         tpot,
+			PrefixCacheScore:   prefixCache, // Added prefix cache score
+			Timestamp:          time.Now().Add(-time.Duration(rng.Intn(3600)) * time.Second),
+		}
+	}
+
+	return entries
+}
+
+// Benchmark test for prediction performance with prefix cache
+func BenchmarkPrediction(b *testing.B) {
+	predictionURLs := os.Getenv("PREDICTION_SERVER_URL")
+	trainingURL := os.Getenv("TRAINING_SERVER_URL")
+	if predictionURLs == "" {
+		b.Skip("PREDICTION_SERVER_URL not set, skipping benchmark")
+	}
+	if trainingURL == "" {
+		// Use first prediction URL as fallback
+		urls := strings.Split(predictionURLs, ",")
+		if len(urls) > 0 {
+			trainingURL = strings.TrimSpace(urls[0])
+		} else {
+			b.Skip("No valid URLs available for benchmarking")
+		}
+	}
+
+	// Parse prediction URLs
+	urls := strings.Split(predictionURLs, ",")
+	parsedPredictionURLs := make([]string, 0, len(urls))
+	for _, url := range urls {
+		parsedPredictionURLs = append(parsedPredictionURLs, strings.TrimSpace(url))
+	}
+
+	logger := logr.Discard() // Silent logger for benchmark
+	config := &Config{
+		TrainingURL:            trainingURL,
+		PredictionURLs:         parsedPredictionURLs,
+		MaxSampleSize:          1000,
+		FlushInterval:          1 * time.Second, // Long interval for benchmark
+		MetricsRefreshInterval: 1 * time.Second,
+		UseNativeXGBoost:       true,
+		HTTPTimeout:            10 * time.Second,
+		MaxBulkSize:            100,
+	}
+
+	predictor := New(config, logger)
+	defer predictor.Stop()
+
+	ctx := context.Background()
+	err := predictor.Start(ctx)
+	if err != nil {
+		b.Fatalf("Failed to start predictor: %v", err)
+	}
+
+	// Wait for predictor to be ready
+	for i := 0; i < 100; i++ {
+		if predictor.IsReady() {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	req := PredictionRequest{
+		KVCachePercentage:  0.75, // 75% as a fraction
+		InputTokenLength:   512,
+		NumRequestWaiting:  2,
+		NumRequestRunning:  1,
+		NumTokensGenerated: 100,
+		PrefixCacheScore:   0.8, // 80% prefix cache hit rate
+	}
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_, err := predictor.Predict(ctx, req)
+			if err != nil {
+				b.Errorf("Prediction failed: %v", err)
+			}
+		}
+	})
+}
+
+// Benchmark test for bulk prediction performance
+func BenchmarkBulkPrediction(b *testing.B) {
+	predictionURLs := os.Getenv("PREDICTION_SERVER_URL")
+	trainingURL := os.Getenv("TRAINING_SERVER_URL")
+	if predictionURLs == "" {
+		b.Skip("PREDICTION_SERVER_URL not set, skipping benchmark")
+	}
+	if trainingURL == "" {
+		urls := strings.Split(predictionURLs, ",")
+		if len(urls) > 0 {
+			trainingURL = strings.TrimSpace(urls[0])
+		} else {
+			b.Skip("No valid URLs available for benchmarking")
+		}
+	}
+
+	// Parse prediction URLs
+	urls := strings.Split(predictionURLs, ",")
+	parsedPredictionURLs := make([]string, 0, len(urls))
+	for _, url := range urls {
+		parsedPredictionURLs = append(parsedPredictionURLs, strings.TrimSpace(url))
+	}
+
+	logger := logr.Discard()
+	config := &Config{
+		TrainingURL:            trainingURL,
+		PredictionURLs:         parsedPredictionURLs,
+		MaxSampleSize:          1000,
+		FlushInterval:          1 * time.Second,
+		MetricsRefreshInterval: 1 * time.Second,
+		UseNativeXGBoost:       true,
+		HTTPTimeout:            10 * time.Second,
+		MaxBulkSize:            50,
+	}
+
+	predictor := New(config, logger)
+	defer predictor.Stop()
+
+	ctx := context.Background()
+	err := predictor.Start(ctx)
+	if err != nil {
+		b.Fatalf("Failed to start predictor: %v", err)
+	}
+
+	for i := 0; i < 100; i++ {
+		if predictor.IsReady() {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	// Create batch of requests
+	const batchSize = 20
+	requests := make([]PredictionRequest, batchSize)
+	for i := 0; i < batchSize; i++ {
+		requests[i] = PredictionRequest{
+			KVCachePercentage:  0.6 + float64(i%5)*0.05,
+			InputTokenLength:   300 + i*10,
+			NumRequestWaiting:  i % 4,
+			NumRequestRunning:  (i % 2) + 1,
+			NumTokensGenerated: 50 + i*2,
+			PrefixCacheScore:   float64(i) / float64(batchSize-1),
+		}
+	}
+
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_, err := predictor.PredictBulk(ctx, requests)
+			if err != nil {
+				b.Errorf("Bulk prediction failed: %v", err)
+			}
+		}
+	})
+}
+
+// Test to verify config loading from environment
+func TestConfigFromEnv(t *testing.T) {
+	// Save original env vars
+	originalLatencyURL := os.Getenv("PREDICTION_SERVER_URL")
+	originalTrainingURL := os.Getenv("TRAINING_SERVER_URL")
+	originalSample := os.Getenv("LATENCY_MAX_SAMPLE_SIZE")
+	originalInterval := os.Getenv("LATENCY_FLUSH_INTERVAL_SEC")
+	originalNative := os.Getenv("LATENCY_USE_NATIVE_XGBOOST")
+	originalTimeout := os.Getenv("LATENCY_HTTP_TIMEOUT_SEC")
+	originalBulkSize := os.Getenv("LATENCY_MAX_BULK_SIZE")
+
+	// Set test env vars
+	os.Setenv("PREDICTION_SERVER_URL", "http://pred1.example.com,http://pred2.example.com,http://pred3.example.com")
+	os.Setenv("TRAINING_SERVER_URL", "http://training.example.com")
+	os.Setenv("LATENCY_MAX_SAMPLE_SIZE", "500")
+	os.Setenv("LATENCY_FLUSH_INTERVAL_SEC", "5")
+	os.Setenv("LATENCY_USE_NATIVE_XGBOOST", "false")
+	os.Setenv("LATENCY_HTTP_TIMEOUT_SEC", "20")
+	os.Setenv("LATENCY_MAX_BULK_SIZE", "75")
+
+	defer func() {
+		// Restore original env vars (handle empty strings properly)
+		if originalLatencyURL != "" {
+			os.Setenv("PREDICTION_SERVER_URL", originalLatencyURL)
+		} else {
+			os.Unsetenv("PREDICTION_SERVER_URL")
+		}
+		if originalTrainingURL != "" {
+			os.Setenv("TRAINING_SERVER_URL", originalTrainingURL)
+		} else {
+			os.Unsetenv("TRAINING_SERVER_URL")
+		}
+		if originalSample != "" {
+			os.Setenv("LATENCY_MAX_SAMPLE_SIZE", originalSample)
+		} else {
+			os.Unsetenv("LATENCY_MAX_SAMPLE_SIZE")
+		}
+		if originalInterval != "" {
+			os.Setenv("LATENCY_FLUSH_INTERVAL_SEC", originalInterval)
+		} else {
+			os.Unsetenv("LATENCY_FLUSH_INTERVAL_SEC")
+		}
+		if originalNative != "" {
+			os.Setenv("LATENCY_USE_NATIVE_XGBOOST", originalNative)
+		} else {
+			os.Unsetenv("LATENCY_USE_NATIVE_XGBOOST")
+		}
+		if originalTimeout != "" {
+			os.Setenv("LATENCY_HTTP_TIMEOUT_SEC", originalTimeout)
+		} else {
+			os.Unsetenv("LATENCY_HTTP_TIMEOUT_SEC")
+		}
+		if originalBulkSize != "" {
+			os.Setenv("LATENCY_MAX_BULK_SIZE", originalBulkSize)
+		} else {
+			os.Unsetenv("LATENCY_MAX_BULK_SIZE")
+		}
+	}()
+
+	config := ConfigFromEnv()
+
+	// Test training URL
+	if config.TrainingURL != "http://training.example.com" {
+		t.Errorf("Expected TrainingURL to be 'http://training.example.com', got '%s'", config.TrainingURL)
+	}
+
+	// Test prediction URLs
+	expectedPredictionURLs := []string{
+		"http://pred1.example.com",
+		"http://pred2.example.com",
+		"http://pred3.example.com",
+	}
+	if len(config.PredictionURLs) != len(expectedPredictionURLs) {
+		t.Errorf("Expected %d prediction URLs, got %d", len(expectedPredictionURLs), len(config.PredictionURLs))
+	}
+	for i, expected := range expectedPredictionURLs {
+		if i >= len(config.PredictionURLs) || config.PredictionURLs[i] != expected {
+			t.Errorf("Expected PredictionURLs[%d] to be '%s', got '%s'", i, expected, config.PredictionURLs[i])
+		}
+	}
+
+	// Test other config values
+	if config.MaxSampleSize != 500 {
+		t.Errorf("Expected MaxSampleSize to be 500, got %d", config.MaxSampleSize)
+	}
+	if config.FlushInterval != 5*time.Second {
+		t.Errorf("Expected FlushInterval to be 5s, got %v", config.FlushInterval)
+	}
+	if config.MetricsRefreshInterval != 60*time.Second {
+		t.Errorf("Expected MetricsRefreshInterval to be 60s, got %v", config.MetricsRefreshInterval)
+	}
+	if config.UseNativeXGBoost != false {
+		t.Errorf("Expected UseNativeXGBoost to be false, got %t", config.UseNativeXGBoost)
+	}
+	if config.HTTPTimeout != 20*time.Second {
+		t.Errorf("Expected HTTPTimeout to be 20s, got %v", config.HTTPTimeout)
+	}
+	if config.MaxBulkSize != 75 {
+		t.Errorf("Expected MaxBulkSize to be 75, got %d", config.MaxBulkSize)
+	}
+}
+
+// Helper function for absolute value
+func abs(x float64) float64 {
+	if x < 0 {
+		return -x
+	}
+	return x
+}
+
+// Test comprehensive bulk prediction functionality
+func TestBulkPredictionValidation(t *testing.T) {
+	t.Log("Testing bulk prediction validation...")
+
+	predictor := &Predictor{
+		config: &Config{MaxBulkSize: 5},
+	}
+
+	// Test empty request list
+	_, err := predictor.PredictBulk(context.Background(), []PredictionRequest{})
+	if err == nil {
+		t.Error("Expected error for empty request list")
+	} else {
+		t.Logf("✓ Correctly rejected empty request list: %v", err)
+	}
+
+	// Test oversized request list
+	oversizedRequests := make([]PredictionRequest, 10) // Over the limit of 5
+	for i := range oversizedRequests {
+		oversizedRequests[i] = PredictionRequest{
+			KVCachePercentage:  0.5,
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   0.5,
+		}
+	}
+
+	_, err = predictor.PredictBulk(context.Background(), oversizedRequests)
+	if err == nil {
+		t.Error("Expected error for oversized request list")
+	} else {
+		t.Logf("✓ Correctly rejected oversized request list: %v", err)
+	}
+
+	// Test invalid request in the list
+	invalidRequests := []PredictionRequest{
+		{
+			KVCachePercentage:  0.5,
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   0.5,
+		},
+		{
+			KVCachePercentage:  1.5, // Invalid
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   0.5,
+		},
+	}
+
+	_, err = predictor.PredictBulk(context.Background(), invalidRequests)
+	if err == nil {
+		t.Error("Expected error for invalid request in list")
+	} else {
+		t.Logf("✓ Correctly rejected list with invalid request: %v", err)
+	}
+
+	t.Log("✅ Bulk prediction validation tests completed")
+}
+
+// Test comprehensive prefix cache integration
+func TestPrefixCacheIntegration(t *testing.T) {
+	t.Log("Testing comprehensive prefix cache integration...")
+
+	// Test all components work together with prefix cache
+	zapLog, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatalf("Failed to create logger: %v", err)
+	}
+	logger := zapr.NewLogger(zapLog)
+
+	// Use a minimal config that doesn't require network calls
+	config := &Config{
+		TrainingURL:            "http://mock-training.local",
+		PredictionURLs:         []string{"http://mock-prediction.local"},
+		MaxSampleSize:          100,
+		FlushInterval:          1 * time.Hour, // Very long interval to avoid network calls
+		MetricsRefreshInterval: 1 * time.Hour, // Very long interval to avoid network calls
+		UseNativeXGBoost:       false,
+		HTTPTimeout:            5 * time.Second,
+		MaxBulkSize:            10,
+	}
+
+	predictor := New(config, logger)
+
+	// Manually stop background processes without triggering final flush/refresh
+	defer func() {
+		// Stop background loop without calling Stop() which does final flush
+		close(predictor.done)
+		predictor.wg.Wait()
+		t.Log("Background processes stopped without network calls")
+	}()
+
+	// Test training entries with prefix cache can be created and validated
+	entries := make([]TrainingEntry, 5)
+	for i := 0; i < 5; i++ {
+		entry, err := NewTrainingEntry(
+			float64(i)/10.0,   // kv_cache_percentage
+			100+i*50,          // input_token_length
+			i%3,               // num_request_waiting
+			1,                 // num_request_running (always > 0)
+			10+i*5,            // num_tokens_generated
+			50.0+float64(i)*5, // actual_ttft_ms
+			10.0+float64(i)*2, // actual_tpot_ms
+			float64(i)/4.0,    // prefix_cache_score (0.0 to 1.0)
+		)
+		if err != nil {
+			t.Fatalf("Failed to create training entry %d: %v", i, err)
+		}
+		entries[i] = entry
+
+		t.Logf("Entry %d: prefix_cache=%.1f%%, ttft=%.1f, tpot=%.1f",
+			i, entry.PrefixCacheScore*100, entry.ActualTTFT, entry.ActualTPOT)
+	}
+
+	// Add training data to buffer (won't flush due to long interval)
+	err = predictor.AddTrainingDataBulk(entries)
+	if err != nil {
+		t.Fatalf("Failed to add training entries: %v", err)
+	}
+	t.Log("✓ Successfully added training entries with prefix cache scores to buffer")
+
+	// Test prediction requests with prefix cache can be created and validated
+	requests := make([]PredictionRequest, 3)
+	for i := 0; i < 3; i++ {
+		req, err := NewPredictionRequest(
+			float64(i*20)/100.0, // kv_cache_percentage
+			200+i*100,           // input_token_length
+			i%2,                 // num_request_waiting
+			1,                   // num_request_running (always > 0)
+			20+i*10,             // num_tokens_generated
+			float64(i)/2.0,      // prefix_cache_score
+		)
+		if err != nil {
+			t.Fatalf("Failed to create prediction request %d: %v", i, err)
+		}
+		requests[i] = req
+
+		err = predictor.ValidatePredictionRequest(req)
+		if err != nil {
+			t.Errorf("Valid prediction request %d failed validation: %v", i, err)
+		}
+
+		t.Logf("Request %d: prefix_cache=%.1f%%, kv_cache=%.1f%%, input_len=%d",
+			i, req.PrefixCacheScore*100, req.KVCachePercentage*100, req.InputTokenLength)
+	}
+	t.Log("✓ Successfully created and validated prediction requests with prefix cache")
+
+	// Test validation edge cases
+	edgeCases := []struct {
+		name        string
+		prefixCache float64
+		shouldPass  bool
+	}{
+		{"Zero prefix cache", 0.0, true},
+		{"Max prefix cache", 1.0, true},
+		{"Negative prefix cache", -0.1, false},
+		{"Over-max prefix cache", 1.1, false},
+	}
+
+	for _, tc := range edgeCases {
+		req := PredictionRequest{
+			KVCachePercentage:  0.5,
+			InputTokenLength:   100,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 10,
+			PrefixCacheScore:   tc.prefixCache,
+		}
+
+		err := predictor.ValidatePredictionRequest(req)
+		switch {
+		case tc.shouldPass && err != nil:
+			t.Errorf("Edge case '%s' should pass but failed: %v", tc.name, err)
+		case !tc.shouldPass && err == nil:
+			t.Errorf("Edge case '%s' should fail but passed", tc.name)
+		default:
+			t.Logf("✓ Edge case '%s' handled correctly", tc.name)
+		}
+	}
+
+	// Test that we can access the configuration
+	t.Logf("Configuration validation:")
+	t.Logf("  Training URL: %s", predictor.GetTrainingURL())
+	t.Logf("  Prediction URLs: %v", predictor.GetPredictionURLs())
+	t.Logf("  Max Bulk Size: %d", config.MaxBulkSize)
+
+	// Validate configuration consistency
+	if len(predictor.GetPredictionURLs()) != len(config.PredictionURLs) {
+		t.Errorf("Prediction URLs mismatch: expected %d, got %d",
+			len(config.PredictionURLs), len(predictor.GetPredictionURLs()))
+	}
+
+	if predictor.GetTrainingURL() != config.TrainingURL {
+		t.Errorf("Training URL mismatch: expected %s, got %s",
+			config.TrainingURL, predictor.GetTrainingURL())
+	}
+
+	// Test data structure integrity
+	t.Log("Validating data structure integrity...")
+
+	// Check that training entries maintain their prefix cache scores
+	for i, entry := range entries {
+		expectedPrefixCache := float64(i) / 4.0
+		if abs(entry.PrefixCacheScore-expectedPrefixCache) > 0.001 {
+			t.Errorf("Training entry %d prefix cache score mismatch: expected %.3f, got %.3f",
+				i, expectedPrefixCache, entry.PrefixCacheScore)
+		}
+	}
+
+	// Check that prediction requests maintain their prefix cache scores
+	for i, req := range requests {
+		expectedPrefixCache := float64(i) / 2.0
+		if abs(req.PrefixCacheScore-expectedPrefixCache) > 0.001 {
+			t.Errorf("Prediction request %d prefix cache score mismatch: expected %.3f, got %.3f",
+				i, expectedPrefixCache, req.PrefixCacheScore)
+		}
+	}
+
+	// Verify that training data is properly buffered (not flushed due to long interval)
+	predictor.bufferMu.Lock()
+	bufferedCount := len(predictor.pending)
+	predictor.bufferMu.Unlock()
+
+	if bufferedCount != len(entries) {
+		t.Errorf("Expected %d buffered entries, got %d", len(entries), bufferedCount)
+	} else {
+		t.Logf("✓ Training data properly buffered: %d entries", bufferedCount)
+	}
+
+	t.Log("✅ Comprehensive prefix cache integration test completed (offline mode)")
+}
+
+// Test offline validation functionality without network dependencies
+func TestOfflineValidation(t *testing.T) {
+	t.Log("Testing offline validation functionality...")
+
+	// Create a minimal predictor for validation testing
+	predictor := &Predictor{}
+
+	// Test prediction request validation
+	t.Run("PredictionRequestValidation", func(t *testing.T) {
+		validReq := PredictionRequest{
+			KVCachePercentage:  0.7,
+			InputTokenLength:   500,
+			NumRequestWaiting:  2,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 80,
+			PrefixCacheScore:   0.85,
+		}
+
+		err := predictor.ValidatePredictionRequest(validReq)
+		if err != nil {
+			t.Errorf("Valid prediction request failed validation: %v", err)
+		}
+
+		// Test invalid requests
+		invalidTests := []struct {
+			name string
+			req  PredictionRequest
+		}{
+			{
+				name: "Negative KV cache",
+				req: PredictionRequest{
+					KVCachePercentage:  -0.1,
+					InputTokenLength:   100,
+					NumRequestWaiting:  1,
+					NumRequestRunning:  1,
+					NumTokensGenerated: 10,
+					PrefixCacheScore:   0.5,
+				},
+			},
+			{
+				name: "Invalid prefix cache",
+				req: PredictionRequest{
+					KVCachePercentage:  0.5,
+					InputTokenLength:   100,
+					NumRequestWaiting:  1,
+					NumRequestRunning:  1,
+					NumTokensGenerated: 10,
+					PrefixCacheScore:   1.5,
+				},
+			},
+		}
+
+		for _, test := range invalidTests {
+			err := predictor.ValidatePredictionRequest(test.req)
+			if err == nil {
+				t.Errorf("Invalid request '%s' should have failed validation", test.name)
+			} else {
+				t.Logf("✓ '%s' correctly rejected: %v", test.name, err)
+			}
+		}
+	})
+
+	// Test training entry validation
+	t.Run("TrainingEntryValidation", func(t *testing.T) {
+		validEntry := TrainingEntry{
+			KVCachePercentage:  0.6,
+			InputTokenLength:   300,
+			NumRequestWaiting:  1,
+			NumRequestRunning:  1,
+			NumTokensGenerated: 50,
+			ActualTTFT:         45.5,
+			ActualTPOT:         12.3,
+			PrefixCacheScore:   0.75,
+			Timestamp:          time.Now(),
+		}
+
+		err := predictor.ValidateTrainingEntry(validEntry)
+		if err != nil {
+			t.Errorf("Valid training entry failed validation: %v", err)
+		}
+
+		// Test invalid entry
+		invalidEntry := validEntry
+		invalidEntry.PrefixCacheScore = -0.5
+
+		err = predictor.ValidateTrainingEntry(invalidEntry)
+		if err == nil {
+			t.Error("Invalid training entry should have failed validation")
+		} else {
+			t.Logf("✓ Invalid training entry correctly rejected: %v", err)
+		}
+	})
+
+	// Test constructor functions
+	t.Run("ConstructorFunctions", func(t *testing.T) {
+		// Test valid constructors
+		_, err := NewPredictionRequest(0.5, 100, 1, 1, 10, 0.8)
+		if err != nil {
+			t.Errorf("Valid prediction request constructor failed: %v", err)
+		}
+
+		_, err = NewTrainingEntry(0.5, 100, 1, 1, 10, 30.0, 8.0, 0.8)
+		if err != nil {
+			t.Errorf("Valid training entry constructor failed: %v", err)
+		}
+
+		// Test invalid constructors
+		_, err = NewPredictionRequest(0.5, 100, 1, 1, 10, 1.5) // Invalid prefix cache
+		if err == nil {
+			t.Error("Invalid prediction request constructor should have failed")
+		}
+
+		_, err = NewTrainingEntry(0.5, 100, 1, 1, 10, 30.0, 8.0, -0.1) // Invalid prefix cache
+		if err == nil {
+			t.Error("Invalid training entry constructor should have failed")
+		}
+	})
+
+	t.Log("✅ Offline validation tests completed")
+}
+
+// Test configuration handling without network calls
+func TestConfigurationHandling(t *testing.T) {
+	t.Log("Testing configuration handling...")
+
+	// Test default configuration
+	defaultConfig := DefaultConfig()
+	if defaultConfig.MaxBulkSize != 100 {
+		t.Errorf("Expected default MaxBulkSize to be 100, got %d", defaultConfig.MaxBulkSize)
+	}
+
+	if defaultConfig.UseNativeXGBoost != true {
+		t.Errorf("Expected default UseNativeXGBoost to be true, got %t", defaultConfig.UseNativeXGBoost)
+	}
+
+	// Test configuration with mock URLs (no network calls)
+	config := &Config{
+		TrainingURL:            "http://mock-training.local",
+		PredictionURLs:         []string{"http://mock1.local", "http://mock2.local"},
+		MaxSampleSize:          500,
+		FlushInterval:          2 * time.Second,
+		MetricsRefreshInterval: 5 * time.Second,
+		UseNativeXGBoost:       false,
+		HTTPTimeout:            10 * time.Second,
+		MaxBulkSize:            50,
+	}
+
+	// Create logger
+	zapLog, err := zap.NewDevelopment()
+	if err != nil {
+		t.Fatalf("Failed to create logger: %v", err)
+	}
+	logger := zapr.NewLogger(zapLog)
+
+	// Create predictor but don't start it (to avoid network calls)
+	predictor := New(config, logger)
+
+	// Test configuration access
+	if predictor.GetTrainingURL() != config.TrainingURL {
+		t.Errorf("Training URL mismatch: expected %s, got %s",
+			config.TrainingURL, predictor.GetTrainingURL())
+	}
+
+	predictionURLs := predictor.GetPredictionURLs()
+	if len(predictionURLs) != len(config.PredictionURLs) {
+		t.Errorf("Prediction URLs length mismatch: expected %d, got %d",
+			len(config.PredictionURLs), len(predictionURLs))
+	}
+
+	// Cleanup without starting background processes
+	close(predictor.done)
+	predictor.wg.Wait()
+
+	t.Log("✅ Configuration handling tests completed")
+}