foxhunt/crates/training_uploader/src/main.rs

#![allow(unused_crate_dependencies)]

//! K8s sidecar that uploads training artifacts to S3 and reports completion via gRPC.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use bytes::Bytes;
use clap::Parser;
use object_store::aws::AmazonS3Builder;
use object_store::{Certificate, ClientOptions, ObjectStore, PutPayload};
use tracing::{error, info, warn};

mod proto {
    #![allow(clippy::enum_variant_names, clippy::mixed_attributes_style)]
    pub mod ml_training {
        tonic::include_proto!("ml_training");
    }
    pub use ml_training::*;
}

/// Training uploader sidecar: watches for training completion, uploads
/// artifacts to S3, and reports results back to ml_training_service via gRPC.
#[derive(Parser, Debug)]
#[command(name = "training-uploader")]
struct Args {
    /// Directory where training job writes its output artifacts
    #[arg(long, default_value = "/output")]
    output_dir: PathBuf,

    /// S3 bucket name for artifact storage
    #[arg(long, env = "S3_BUCKET")]
    s3_bucket: String,

    /// S3-compatible endpoint URL
    #[arg(long, env = "S3_ENDPOINT")]
    s3_endpoint: String,

    /// S3 region
    #[arg(long, env = "S3_REGION", default_value = "fr-par")]
    s3_region: String,

    /// S3 key prefix for this job's artifacts
    #[arg(long, env = "S3_PREFIX")]
    s3_prefix: String,

    /// gRPC endpoint of ml_training_service for completion callback
    #[arg(long, env = "CALLBACK_ENDPOINT")]
    callback_endpoint: String,

    /// Unique identifier for this training job
    #[arg(long, env = "JOB_ID")]
    job_id: String,

    /// Seconds between polls for the DONE/FAILED marker file
    #[arg(long, default_value = "10")]
    poll_interval: u64,

    /// Maximum seconds to wait before giving up
    #[arg(long, default_value = "3600")]
    max_wait: u64,
}

/// Poll the output directory for a `DONE` or `FAILED` marker file.
///
/// Returns `(success, status_message)`.
fn wait_for_completion(dir: &Path, poll_secs: u64, max_wait: u64) -> (bool, String) {
    let start = Instant::now();
    let poll = Duration::from_secs(poll_secs);
    let deadline = Duration::from_secs(max_wait);

    loop {
        let done_path = dir.join("DONE");
        let failed_path = dir.join("FAILED");

        if done_path.exists() {
            let msg = std::fs::read_to_string(&done_path).unwrap_or_default();
            return (true, msg);
        }
        if failed_path.exists() {
            let msg = std::fs::read_to_string(&failed_path).unwrap_or_default();
            return (false, msg);
        }

        if start.elapsed() >= deadline {
            return (
                false,
                format!("Timed out after {}s waiting for marker file", max_wait),
            );
        }

        std::thread::sleep(poll);
    }
}

/// Read `metrics.json` from the output directory.
///
/// Returns an empty map if the file is missing or cannot be parsed.
fn read_metrics(dir: &Path) -> HashMap<String, f64> {
    let path = dir.join("metrics.json");
    match std::fs::read_to_string(&path) {
        Ok(data) => serde_json::from_str(&data).unwrap_or_default(),
        Err(_) => {
            warn!("No metrics.json found at {}", path.display());
            HashMap::new()
        }
    }
}

/// Upload all files from the output directory to S3, preserving the
/// subdirectory structure under the configured prefix.
///
/// Uses an iterative stack-based directory walk to avoid async-recursion.
async fn upload_artifacts(args: &Args) -> Result<String> {
    let mut client_options = ClientOptions::new();

    // Load custom CA certificate for self-signed MinIO TLS
    if let Ok(ca_path) = std::env::var("MINIO_CA_CERT_PATH") {
        let pem = std::fs::read(&ca_path)
            .with_context(|| format!("Failed to read CA cert from {ca_path}"))?;
        let cert = Certificate::from_pem(&pem)
            .with_context(|| format!("Failed to parse CA cert from {ca_path}"))?;
        client_options = client_options.with_root_certificate(cert);
        info!("Loaded MinIO CA certificate from {}", ca_path);
    }

    let mut builder = AmazonS3Builder::new()
        .with_bucket_name(&args.s3_bucket)
        .with_endpoint(&args.s3_endpoint)
        .with_region(&args.s3_region)
        .with_virtual_hosted_style_request(false)
        .with_allow_http(false)
        .with_client_options(client_options);

    // Explicitly pass credentials from env vars (MinIO has no instance metadata)
    if let (Ok(key_id), Ok(secret)) = (
        std::env::var("AWS_ACCESS_KEY_ID"),
        std::env::var("AWS_SECRET_ACCESS_KEY"),
    ) {
        builder = builder
            .with_access_key_id(key_id)
            .with_secret_access_key(secret);
    } else {
        warn!("AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY not set — S3 upload may fail");
    }

    let store = builder
        .build()
        .context("Failed to build S3 client")?;

    let mut uploaded = 0usize;
    let mut stack: Vec<(PathBuf, String)> =
        vec![(args.output_dir.clone(), args.s3_prefix.clone())];

    while let Some((local_dir, prefix)) = stack.pop() {
        let entries = std::fs::read_dir(&local_dir)
            .with_context(|| format!("Failed to read dir {}", local_dir.display()))?;

        for entry in entries.flatten() {
            let path = entry.path();
            let name = entry.file_name().to_string_lossy().to_string();

            if path.is_dir() {
                stack.push((path, format!("{}/{}", prefix, name)));
            } else {
                let s3_key = format!("{}/{}", prefix, name);
                let data = std::fs::read(&path)
                    .with_context(|| format!("Failed to read {}", path.display()))?;
                let s3_path = object_store::path::Path::from(s3_key.as_str());
                store
                    .put(&s3_path, PutPayload::from(Bytes::from(data)))
                    .await
                    .with_context(|| format!("Failed to upload {}", s3_key))?;
                uploaded += 1;
            }
        }
    }

    info!("Uploaded {} files to S3", uploaded);
    Ok(args.s3_prefix.clone())
}

/// Send a `JobCompletionReport` to the ml_training_service via gRPC.
async fn report_completion(
    args: &Args,
    success: bool,
    s3_path: &str,
    error_message: &str,
    metrics: HashMap<String, f64>,
) -> Result<()> {
    let mut client =
        proto::ml_training_service_client::MlTrainingServiceClient::connect(
            args.callback_endpoint.clone(),
        )
        .await
        .context("Failed to connect to ml_training_service")?;

    let report = proto::JobCompletionReport {
        job_id: args.job_id.clone(),
        s3_path: s3_path.to_owned(),
        success,
        error_message: error_message.to_owned(),
        metrics,
    };

    let resp = client
        .report_job_completion(report)
        .await
        .context("ReportJobCompletion RPC failed")?;

    let ack = resp.into_inner();
    info!(
        accepted = ack.accepted,
        promotion_status = %ack.promotion_status,
        "Completion report acknowledged"
    );

    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .json()
        .init();

    let args = Args::parse();
    info!(job_id = %args.job_id, output_dir = %args.output_dir.display(), "Training uploader started");

    // Phase 1: Wait for training container to finish
    info!("Waiting for DONE or FAILED marker in {}", args.output_dir.display());
    let (success, status_msg) =
        wait_for_completion(&args.output_dir, args.poll_interval, args.max_wait);

    if success {
        info!("Training completed successfully: {}", status_msg);
    } else {
        error!("Training failed or timed out: {}", status_msg);
    }

    // Phase 2: Read metrics
    let metrics = read_metrics(&args.output_dir);
    if !metrics.is_empty() {
        info!(count = metrics.len(), "Read training metrics");
    }

    // Phase 3: Upload artifacts to S3
    let (s3_path, error_message) = match upload_artifacts(&args).await {
        Ok(path) => (path, if success { String::new() } else { status_msg }),
        Err(e) => {
            error!("S3 upload failed: {:#}", e);
            (
                String::new(),
                format!("S3 upload failed: {e:#}; original status: {status_msg}"),
            )
        }
    };

    // Phase 4: Report completion via gRPC
    if let Err(e) = report_completion(&args, success, &s3_path, &error_message, metrics).await {
        error!("Failed to report completion: {:#}", e);
        return Err(e);
    }

    info!(job_id = %args.job_id, "Uploader finished, exiting");
    Ok(())
}