Machine Learning Workload

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from metaparticle_pkg import Containerize
import os

# Configuration
model_name = 'sentiment-classifier'
repository = 'docker.io/mlteam'
epochs = 10
batch_size = 32


def create_model(input_dim, num_classes):
    """Create a simple neural network model"""
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(input_dim,)),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(num_classes, activation='softmax')
    ])
    
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )
    
    return model


def save_model(model, path='./models'):
    """Save trained model to disk"""
    os.makedirs(path, exist_ok=True)
    model_path = os.path.join(path, f'{model_name}.h5')
    model.save(model_path)
    print(f"Model saved to {model_path}")
    return model_path


@Containerize(
    package={
        'name': model_name,
        'repository': repository,
        'publish': True,
        'dockerfile': 'Dockerfile.ml'  # Use custom Dockerfile with ML dependencies
    },
    runtime={
        'executor': 'docker',
        'gpu': True,  # Enable GPU support
        'memory': '4Gi',  # Request 4GB memory
        'cpu': '2',  # Request 2 CPU cores
    }
)
def train_model():
    """Main training function"""
    print("Starting ML training job...")
    print(f"TensorFlow version: {tf.__version__}")
    print(f"GPU Available: {tf.config.list_physical_devices('GPU')}")
    
    # Generate synthetic dataset
    print("\nGenerating dataset...")
    X, y = make_classification(
        n_samples=10000,
        n_features=20,
        n_informative=15,
        n_redundant=5,
        n_classes=3,
        random_state=42
    )
    
    # Split into train and validation sets
    X_train, X_val, y_train, y_val = train_test_split(
        X, y, test_size=0.2, random_state=42
    )
    
    print(f"Training samples: {len(X_train)}")
    print(f"Validation samples: {len(X_val)}")
    
    # Create and train model
    print("\nBuilding model...")
    model = create_model(input_dim=X_train.shape[1], num_classes=3)
    model.summary()
    
    print("\nTraining model...")
    history = model.fit(
        X_train, y_train,
        validation_data=(X_val, y_val),
        epochs=epochs,
        batch_size=batch_size,
        verbose=1
    )
    
    # Evaluate model
    print("\nEvaluating model...")
    val_loss, val_accuracy = model.evaluate(X_val, y_val, verbose=0)
    print(f"Validation Loss: {val_loss:.4f}")
    print(f"Validation Accuracy: {val_accuracy:.4f}")
    
    # Save model
    model_path = save_model(model)
    
    print("\nTraining completed successfully!")
    return model, history


if __name__ == '__main__':
    model, history = train_model()

FROM tensorflow/tensorflow:2.13.0-gpu

# Install additional dependencies
RUN pip install --no-cache-dir \
    scikit-learn \
    pandas \
    matplotlib \
    seaborn \
    jupyter

# Copy application code
COPY ./ /app/
WORKDIR /app

# Install project dependencies
RUN pip install --no-cache-dir -r requirements.txt

# Run training script
CMD ["python", "train.py"]

configurations = [
    {'learning_rate': 0.001, 'batch_size': 32},
    {'learning_rate': 0.0001, 'batch_size': 64},
    {'learning_rate': 0.01, 'batch_size': 16},
]

for i, config in enumerate(configurations):
    @Containerize(
        package={'name': f'{model_name}-{i}', 'repository': repository},
        runtime={'gpu': True, 'memory': '4Gi'}
    )
    def train_with_config():
        return train_model(**config)
    
    train_with_config()

@Containerize(
    package={'name': 'distributed-training', 'repository': repository},
    runtime={
        'executor': 'metaparticle',
        'replicas': 4,  # 4 training workers
        'gpu': True,
        'memory': '8Gi',
        'env': {
            'TF_CONFIG': json.dumps({
                'cluster': {'worker': ['worker-0:2222', 'worker-1:2222', 
                                       'worker-2:2222', 'worker-3:2222']},
                'task': {'type': 'worker', 'index': 0}
            })
        }
    }
)
def distributed_train():
    strategy = tf.distribute.MultiWorkerMirroredStrategy()
    with strategy.scope():
        model = create_model()
    # ... training code

from flask import Flask, request, jsonify

app = Flask(__name__)
model = tf.keras.models.load_model('./models/sentiment-classifier.h5')

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json['data']
    predictions = model.predict(data)
    return jsonify({'predictions': predictions.tolist()})

@Containerize(
    package={'name': 'ml-inference', 'repository': repository, 'publish': True},
    runtime={'ports': [5000], 'public': True, 'replicas': 3, 'memory': '2Gi'}
)
def serve():
    app.run(host='0.0.0.0', port=5000)

if __name__ == '__main__':
    serve()