rustytorch/crates/production/rtx-serving-api/src/grpc.rs

//! gRPC service implementation for model serving

use crate::error::ApiResult;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status};

/// gRPC server configuration
#[derive(Debug, Clone)]
pub struct GrpcConfig {
    pub port: u16,
    pub max_message_size: usize,
    pub tls_cert: Option<String>,
    pub tls_key: Option<String>,
    pub enable_reflection: bool,
}

impl Default for GrpcConfig {
    fn default() -> Self {
        Self {
            port: 50051,
            max_message_size: 4 * 1024 * 1024, // 4MB
            tls_cert: None,
            tls_key: None,
            enable_reflection: true,
        }
    }
}

impl GrpcConfig {
    pub fn enable_tls(&mut self, cert: &str, key: &str) {
        self.tls_cert = Some(cert.to_string());
        self.tls_key = Some(key.to_string());
    }
}

/// Model information for gRPC responses
#[derive(Debug, Clone)]
pub struct ModelInfo {
    pub name: String,
    pub model_type: String,
    pub max_batch_size: usize,
    pub is_loaded: bool,
}

/// gRPC server for model inference
#[derive(Debug)]
pub struct GrpcServer {
    config: GrpcConfig,
    models: Arc<RwLock<HashMap<String, ModelInfo>>>,
    is_ready: Arc<RwLock<bool>>,
    start_time: std::time::Instant,
}

impl GrpcServer {
    /// Create a new gRPC server
    #[must_use]
    pub fn new(config: GrpcConfig) -> Self {
        Self {
            config,
            models: Arc::new(RwLock::new(HashMap::new())),
            is_ready: Arc::new(RwLock::new(true)),
            start_time: std::time::Instant::now(),
        }
    }

    /// Get the server port
    #[must_use]
    pub fn port(&self) -> u16 {
        self.config.port
    }

    /// Check if server is ready
    #[must_use]
    pub fn is_ready(&self) -> bool {
        futures::executor::block_on(async { *self.is_ready.read().await })
    }

    /// Check if TLS is enabled
    #[must_use]
    pub fn is_tls_enabled(&self) -> bool {
        self.config.tls_cert.is_some() && self.config.tls_key.is_some()
    }

    /// Register a model
    pub fn register_model(&mut self, name: &str, model_type: &str, max_batch_size: usize) {
        futures::executor::block_on(async {
            let mut models = self.models.write().await;
            models.insert(
                name.to_string(),
                ModelInfo {
                    name: name.to_string(),
                    model_type: model_type.to_string(),
                    max_batch_size,
                    is_loaded: true,
                },
            );
        });
    }

    /// Health check handler
    pub async fn health_check(
        &self,
        request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        let uptime = self.start_time.elapsed().as_secs();

        Ok(Response::new(HealthCheckResponse {
            status: "healthy".to_string(),
            uptime_seconds: uptime,
            service: request.into_inner().service,
        }))
    }

    /// Model info handler
    pub async fn model_info(
        &self,
        request: Request<ModelInfoRequest>,
    ) -> Result<Response<ModelInfoResponse>, Status> {
        let model_name = request.into_inner().model_name;
        let models = self.models.read().await;

        let info = models
            .get(&model_name)
            .ok_or_else(|| Status::not_found(format!("Model {model_name} not found")))?;

        Ok(Response::new(ModelInfoResponse {
            name: info.name.clone(),
            model_type: info.model_type.clone(),
            max_batch_size: info.max_batch_size as i32,
            is_loaded: info.is_loaded,
        }))
    }

    /// Predict handler for single inference
    pub async fn predict(
        &self,
        request: Request<InferenceRequest>,
    ) -> Result<Response<InferenceResponse>, Status> {
        let req = request.into_inner();
        let models = self.models.read().await;

        if !models.contains_key(&req.model_name) {
            return Err(Status::not_found(format!(
                "Model {} not found",
                req.model_name
            )));
        }

        // Simulate inference
        let start = std::time::Instant::now();
        let outputs = vec![0.1, 0.2, 0.3, 0.4]; // Placeholder outputs
        let latency_ms = start.elapsed().as_secs_f64() * 1000.0;

        Ok(Response::new(InferenceResponse {
            model_name: req.model_name,
            outputs,
            latency_ms,
            batch_size: req.batch_size,
        }))
    }

    /// Streaming predict handler
    pub async fn predict_stream(
        &self,
        request: Request<InferenceRequest>,
    ) -> Result<Response<ReceiverStream<Result<InferenceResponse, Status>>>, Status> {
        let req = request.into_inner();
        let models = self.models.read().await;

        if !models.contains_key(&req.model_name) {
            return Err(Status::not_found(format!(
                "Model {} not found",
                req.model_name
            )));
        }

        let (tx, rx) = tokio::sync::mpsc::channel(128);
        let model_name = req.model_name.clone();

        // Spawn task to generate streaming tokens
        tokio::spawn(async move {
            for i in 0..5 {
                let response = InferenceResponse {
                    model_name: model_name.clone(),
                    outputs: vec![i as f32 * 0.1],
                    latency_ms: 10.0,
                    batch_size: 1,
                };

                if tx.send(Ok(response)).await.is_err() {
                    break;
                }

                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
            }
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }

    /// Shutdown the server
    pub async fn shutdown(&self) -> ApiResult<()> {
        let mut is_ready = self.is_ready.write().await;
        *is_ready = false;
        Ok(())
    }
}

// Proto message definitions
#[derive(Debug, Clone)]
pub struct HealthCheckRequest {
    pub service: String,
}

#[derive(Debug, Clone)]
pub struct HealthCheckResponse {
    pub status: String,
    pub uptime_seconds: u64,
    pub service: String,
}

#[derive(Debug, Clone)]
pub struct ModelInfoRequest {
    pub model_name: String,
}

#[derive(Debug, Clone)]
pub struct ModelInfoResponse {
    pub name: String,
    pub model_type: String,
    pub max_batch_size: i32,
    pub is_loaded: bool,
}

#[derive(Debug, Clone, Default)]
pub struct InferenceRequest {
    pub model_name: String,
    pub input_data: Vec<f32>,
    pub batch_size: i32,
    pub options: InferenceOptions,
}

#[derive(Debug, Clone, Default)]
pub struct InferenceOptions {
    pub temperature: Option<f32>,
    pub top_k: Option<i32>,
    pub top_p: Option<f32>,
    pub max_tokens: Option<i32>,
}

#[derive(Debug, Clone)]
pub struct InferenceResponse {
    pub model_name: String,
    pub outputs: Vec<f32>,
    pub latency_ms: f64,
    pub batch_size: i32,
}

/// Placeholder for the `InferenceService` trait
pub mod inference_service_server {
    use super::{
        HealthCheckRequest, HealthCheckResponse, InferenceRequest, InferenceResponse,
        ModelInfoRequest, ModelInfoResponse, ReceiverStream, Request, Response, Status,
    };

    #[async_trait::async_trait]
    pub trait InferenceService: Send + Sync + 'static {
        async fn health_check(
            &self,
            request: Request<HealthCheckRequest>,
        ) -> Result<Response<HealthCheckResponse>, Status>;

        async fn model_info(
            &self,
            request: Request<ModelInfoRequest>,
        ) -> Result<Response<ModelInfoResponse>, Status>;

        async fn predict(
            &self,
            request: Request<InferenceRequest>,
        ) -> Result<Response<InferenceResponse>, Status>;

        async fn predict_stream(
            &self,
            request: Request<InferenceRequest>,
        ) -> Result<Response<ReceiverStream<Result<InferenceResponse, Status>>>, Status>;
    }
}

#[async_trait::async_trait]
impl inference_service_server::InferenceService for GrpcServer {
    async fn health_check(
        &self,
        request: Request<HealthCheckRequest>,
    ) -> Result<Response<HealthCheckResponse>, Status> {
        self.health_check(request).await
    }

    async fn model_info(
        &self,
        request: Request<ModelInfoRequest>,
    ) -> Result<Response<ModelInfoResponse>, Status> {
        self.model_info(request).await
    }

    async fn predict(
        &self,
        request: Request<InferenceRequest>,
    ) -> Result<Response<InferenceResponse>, Status> {
        self.predict(request).await
    }

    async fn predict_stream(
        &self,
        request: Request<InferenceRequest>,
    ) -> Result<Response<ReceiverStream<Result<InferenceResponse, Status>>>, Status> {
        self.predict_stream(request).await
    }
}

/// Placeholder for `ScalingAction` enum
#[derive(Debug, Clone, PartialEq)]
pub enum ScalingAction {
    ScaleUp(usize),
    ScaleDown(usize),
    None,
}