# Machine Learning Infrastructure Example

Deploy infrastructure for machine learning workloads including training, inference, and feature engineering.

## ML Infrastructure Architecture

```text
Data Ingestion
    ↓
Feature Store (SurrealDB Vector)
    ↓
Training Jobs (GPU Clusters)
    ├─ TensorFlow
    ├─ PyTorch
    └─ XGBoost
    ↓
Model Registry (Artifact Storage)
    ↓
Inference Services
    ├─ Online Inference (API)
    ├─ Batch Inference
    └─ Real-time Inference (Streaming)
    ↓
Monitoring and Feedback Loop
```

## ML Infrastructure Configuration

```nickel
{
  ml_infrastructure = {
    features = "enabled",

    training = {
      cluster = {
        name = "ml-training",
        compute = {
          gpu_enabled = true,
          gpu_type = "a100",
          gpu_count = 8,
          instance_type = "p4de.24xlarge",
          instance_count = 5,
          auto_scaling = {
            enabled = true,
            min_instances = 1,
            max_instances = 20,
            target_gpu_utilization = 70
          }
        },

        storage = {
          datasets_gb = 10000,
          models_gb = 500,
          logs_gb = 100,
          storage_class = "gp3"
        },

        networking = {
          high_bandwidth = true,
          gpu_peer_to_peer = true,
          collective_communication = "nccl"
        }
      }
    },

    inference = {
      online = {
        instance_type = "inf2.2xlarge",
        instance_count = 5,
        auto_scaling = {
          enabled = true,
          target_rps = 1000,
          max_instances = 20
        }
      },

      batch = {
        instance_type = "t3.2xlarge",
        instance_count = 10,
        schedule = "0 2 * * *"  # 2 AM daily
      }
    },

    feature_store = {
      engine = "surrealdb",
      vector_embeddings = {
        enabled = true,
        dimensions = 768,
        index_type = "hnsw"
      }
    }
  }
}
```

## Training Pipeline

### Distributed Training

```python
# training/distributed_training.py
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
import provisioning_sdk

# Initialize Provisioning client
prov = provisioning_sdk.Client()

# Get training job metadata
job_config = prov.get_job_config()
rank = int(job_config["rank"])
world_size = int(job_config["world_size"])
master_addr = job_config["master_addr"]

# Initialize distributed training
dist.init_process_group(
    backend="nccl",
    init_method=f"tcp://{master_addr}:29500",
    rank=rank,
    world_size=world_size
)

# Load model
model = load_model()
ddp_model = DistributedDataParallel(model)

# Load distributed dataset
dataset = load_distributed_dataset(rank, world_size)
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    sampler=DistributedSampler(dataset)
)

# Training loop
optimizer = torch.optim.Adam(ddp_model.parameters())

for epoch in range(100):
    for batch in dataloader:
        outputs = ddp_model(batch)
        loss = compute_loss(outputs)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Report metrics
        prov.report_metric(
            job_id=job_config["job_id"],
            metric="loss",
            value=float(loss),
            step=epoch
        )

    # Checkpoint
    if rank == 0:
        torch.save(ddp_model.state_dict(), f"checkpoint-{epoch}.pt")
        prov.upload_model(
            job_id=job_config["job_id"],
            model_path=f"checkpoint-{epoch}.pt",
            metadata={"epoch": epoch}
        )

dist.destroy_process_group()
```

### Training Orchestration

```nushell
def submit-training-job [--config: string] {
    print "Submitting ML training job..."

    # 1. Create training job
    let job = (
        provisioning ml submit \
            --config $config \
            --gpus 8 \
            --batch-size 256
    )

    print $"Job submitted: ($job.job_id)"

    # 2. Monitor training
    let monitor_start = (date now)

    loop {
        let status = (
            provisioning ml job-status \
                --job-id $job.job_id
        )

        print $"Status: ($status.status) - Epoch: ($status.epoch)/100"
        print $"Loss: ($status.metrics.loss) - Accuracy: ($status.metrics.accuracy)%"

        # Check for convergence
        if $status.status == "completed" {
            print "✓ Training completed successfully"
            break
        }

        if $status.status == "failed" {
            print $"✗ Training failed: ($status.error)"
            return 1
        }

        sleep 30s
    }

    # 3. Download final model
    provisioning ml download-model \
        --job-id $job.job_id \
        --output "models/final-model.pt"
}
```

## Feature Store

### Feature Engineering

```nushell
def compute-features [--dataset: string] {
    print "Computing features for ML model..."

    # 1. Load raw data
    let raw_data = (provisioning data load $dataset)

    # 2. Compute derived features
    let features = {
        user_engagement: (
            $raw_data | each { | record |
                {
                    user_id: $record.user_id,
                    engagement_score: (
                        ($record.views * 0.3) +
                        ($record.clicks * 0.7)
                    )
                }
            }
        ),

        temporal_features: (
            $raw_data | each { | record |
                {
                    user_id: $record.user_id,
                    hour_of_day: (($record.timestamp | date to-record).hour),
                    day_of_week: (($record.timestamp | date to-record).day)
                }
            }
        ),

        user_embeddings: (
            provisioning ml compute-embeddings \
                --input $raw_data \
                --model "sentence-transformer"
        )
    }

    # 3. Store in feature store
    provisioning feature-store upsert \
        --table "user_features" \
        --data $features

    print "✓ Features computed and stored"
}
```

## Model Registry and Versioning

```nushell
def register-model [--name: string, --path: string] {
    print $"Registering model: ($name)..."

    # 1. Upload model to registry
    let model = (
        provisioning ml model register \
            --name $name \
            --model-path $path \
            --framework "pytorch" \
            --version "auto"
    )

    print $"Model registered: ($model.model_id) v($model.version)"

    # 2. Compute model metrics
    let test_results = (
        provisioning ml evaluate \
            --model-id $model.model_id \
            --test-dataset "test-2025-01"
    )

    print $"Accuracy: ($test_results.accuracy)%"
    print $"Precision: ($test_results.precision)%"
    print $"Recall: ($test_results.recall)%"
    print $"F1-Score: ($test_results.f1_score)"

    # 3. Deploy to inference
    if $test_results.accuracy > 95 {
        provisioning ml deploy \
            --model-id $model.model_id \
            --target "online-inference" \
            --replicas 3 \
            --canary-percentage 10
    }
}
```

## Inference Services

### Online Inference API

```rust
// inference/src/main.rs
use actix_web::{web, App, HttpServer, HttpResponse};
use serde::{Deserialize, Serialize};
use onnxruntime::environment::Environment;

#[derive(Deserialize)]
struct PredictionRequest {
    features: Vec<f32>,
}

#[derive(Serialize)]
struct PredictionResponse {
    prediction: f32,
    confidence: f32,
    model_version: String,
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let model_path = std::env::var("MODEL_PATH")
        .expect("MODEL_PATH not set");

    let env = Environment::builder()
        .with_name("inference")
        .build()
        .expect("Failed to build environment");

    let session = env
        .new_session_builder()
        .expect("Failed to create session builder")
        .with_optimization_level(ort::GraphOptimizationLevel::All)
        .expect("Failed to set optimization level")
        .with_number_threads(4)
        .expect("Failed to set threads")
        .build_from_file(&model_path)
        .expect("Failed to load model");

    let session = web::Data::new(session);

    HttpServer::new(move | | {
        App::new()
            .app_data(session.clone())
            .route("/health", web::get().to(health_check))
            .route("/predict", web::post().to(predict))
    })
    .bind("0.0.0.0:8080")?
    .run()
    .await
}

async fn health_check() -> HttpResponse {
    HttpResponse::Ok().json(serde_json::json!({
        "status": "healthy",
        "service": "ml-inference",
        "model_loaded": true
    }))
}

async fn predict(
    req: web::Json<PredictionRequest>,
) -> HttpResponse {
    let input_data = ort::DynValue::from_array_container(
        ndarray::Array1::from_vec(req.features.clone())
    ).expect("Failed to create input");

    let output = vec![input_data];

    let prediction = output[0]
        .extract_tensor()
        .expect("Failed to extract output")
        .to_vec();

    HttpResponse::Ok().json(PredictionResponse {
        prediction: prediction[0],
        confidence: prediction[1],
        model_version: "v1.2.3".to_string(),
    })
}
```

### Batch Inference

```nushell
def batch-inference [] {
    print "Running batch inference job..."

    let job = (
        provisioning ml batch-predict \
            --model-name "recommendation-model" \
            --input "s3://data/batch-input.parquet" \
            --output "s3://results/batch-output" \
            --batch-size 1000 \
            --instance-type "t3.2xlarge" \
            --parallel-workers 10
    )

    print $"Batch job started: ($job.job_id)"

    # Monitor progress
    loop {
        let progress = (
            provisioning ml batch-progress \
                --job-id $job.job_id
        )

        print $"Progress: ($progress.processed)/($progress.total) records"

        if $progress.status == "completed" {
            print $"✓ Batch inference complete"
            print $"Results: s3://results/batch-output"
            break
        }

        sleep 60s
    }
}
```

## Model Monitoring

```nushell
def monitor-model-drift [] {
    print "Monitoring model performance and data drift..."

    # 1. Monitor prediction metrics
    let model_metrics = (
        provisioning ml metrics get \
            --model-name "recommendation-model" \
            --window "7d"
    )

    print $"Inference Latency P99: ($model_metrics.latency_p99)ms"
    print $"Prediction Distribution: ($model_metrics.prediction_dist)"

    # 2. Check data drift
    let drift = (
        provisioning ml detect-drift \
            --model-name "recommendation-model" \
            --feature-store "user_features" \
            --baseline-date "2025-01-01"
    )

    if $drift.detected {
        print "⚠️  Data drift detected!"
        print $"Drift score: ($drift.score)"
        print "Recommendation: Retrain model"
    }

    # 3. Monitor model performance on live data
    let live_performance = (
        provisioning ml evaluate-live \
            --model-name "recommendation-model" \
            --metric "click_through_rate"
    )

    print $"Live CTR: ($live_performance.metric_value)%"
}
```

## See Also

- [AI and Machine Learning](../ai/README.md)
- [RAG and Knowledge System](../ai/rag-and-knowledge.md)
- [Cost Optimization](./cost-optimization-examples.md)
