TensorFlow GPU Slim Images - Repository Skeleton¶

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import tensorflow as tf
import numpy as np
from typing import List, Dict, Any
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="TensorFlow GPU Inference API",
    description="Slim TensorFlow GPU container for inference",
    version="1.0.0"
)

# Global model variable
model = None

class PredictionRequest(BaseModel):
    data: List[float]
    batch_size: int = 1

class PredictionResponse(BaseModel):
    predictions: List[float]
    model_version: str
    gpu_available: bool

@app.on_event("startup")
async def load_model():
    """Load the TensorFlow model on startup"""
    global model
    try:
        model = tf.keras.models.load_model("/app/models/model.h5")
        logger.info(f"Model loaded successfully. Version: {tf.__version__}")

        # Log GPU information
        gpus = tf.config.list_physical_devices('GPU')
        if gpus:
            logger.info(f"GPU available: {[gpu.name for gpu in gpus]}")
        else:
            logger.warning("No GPU detected, running on CPU")

    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    gpus = tf.config.list_physical_devices('GPU')
    return {
        "status": "healthy",
        "tensorflow_version": tf.__version__,
        "gpu_available": len(gpus) > 0,
        "gpu_devices": [gpu.name for gpu in gpus] if gpus else [],
        "model_loaded": model is not None
    }

@app.get("/model/info")
async def model_info():
    """Get model information"""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    return {
        "model_summary": str(model.summary()),
        "input_shape": model.input_shape,
        "output_shape": model.output_shape,
        "total_params": model.count_params()
    }

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    """Make predictions using the loaded model"""
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        # Convert input to numpy array
        input_data = np.array(request.data).reshape(1, -1)

        # Make prediction
        prediction = model.predict(input_data, batch_size=request.batch_size)

        # Convert to list for JSON serialization
        predictions = prediction.tolist()[0] if len(prediction.shape) > 1 else prediction.tolist()

        return PredictionResponse(
            predictions=predictions,
            model_version=tf.__version__,
            gpu_available=len(tf.config.list_physical_devices('GPU')) > 0
        )

    except Exception as e:
        logger.error(f"Prediction failed: {e}")
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")

@app.get("/gpu/info")
async def gpu_info():
    """Get detailed GPU information"""
    gpus = tf.config.list_physical_devices('GPU')
    if not gpus:
        return {"gpu_available": False, "message": "No GPU detected"}

    gpu_info = []
    for gpu in gpus:
        gpu_info.append({
            "name": gpu.name,
            "device_type": gpu.device_type,
            "memory_limit": tf.config.experimental.get_memory_info(gpu.name) if hasattr(tf.config.experimental, 'get_memory_info') else None
        })

    return {
        "gpu_available": True,
        "gpu_count": len(gpus),
        "gpus": gpu_info
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

"""
Create a simple Keras model for demonstration purposes.
This script creates a small neural network and saves it as model.h5
"""

import tensorflow as tf
import numpy as np
from pathlib import Path

def create_simple_model():
    """Create a simple neural network for demonstration"""
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(
        optimizer='adam',
        loss='binary_crossentropy',
        metrics=['accuracy']
    )

    return model

def generate_training_data(n_samples=1000):
    """Generate synthetic training data"""
    np.random.seed(42)
    X = np.random.randn(n_samples, 10)
    y = (X.sum(axis=1) > 0).astype(int)
    return X, y

def main():
    """Create and train the model"""
    print("Creating model...")
    model = create_simple_model()

    print("Generating training data...")
    X, y = generate_training_data()

    print("Training model...")
    model.fit(X, y, epochs=10, batch_size=32, validation_split=0.2, verbose=1)

    # Create models directory
    models_dir = Path("models")
    models_dir.mkdir(exist_ok=True)

    # Save model
    model_path = models_dir / "model.h5"
    model.save(model_path)
    print(f"Model saved to {model_path}")

    # Test the model
    test_data = np.random.randn(1, 10)
    prediction = model.predict(test_data)
    print(f"Test prediction: {prediction[0][0]:.4f}")

if __name__ == "__main__":
    main()

"""
Data preprocessing utilities for the TensorFlow application
"""

import numpy as np
from typing import List, Union

def normalize_data(data: Union[List[float], np.ndarray]) -> np.ndarray:
    """Normalize input data to [0, 1] range"""
    data = np.array(data)
    return (data - data.min()) / (data.max() - data.min() + 1e-8)

def standardize_data(data: Union[List[float], np.ndarray]) -> np.ndarray:
    """Standardize input data to zero mean and unit variance"""
    data = np.array(data)
    return (data - data.mean()) / (data.std() + 1e-8)

def validate_input_shape(data: Union[List[float], np.ndarray], expected_shape: int = 10) -> bool:
    """Validate that input data has the expected shape"""
    data = np.array(data)
    return len(data) == expected_shape

def preprocess_for_prediction(data: List[float]) -> np.ndarray:
    """Complete preprocessing pipeline for prediction"""
    if not validate_input_shape(data):
        raise ValueError(f"Expected input shape of 10, got {len(data)}")

    # Normalize the data
    processed_data = normalize_data(data)

    # Reshape for model input
    return processed_data.reshape(1, -1)