test_model_serving.py

#!/usr/bin/env python3
"""
Test script for MLflow model serving endpoint.

Usage:
    python test_model_serving.py [--port PORT] [--run-id RUN_ID]
"""

import requests
import json
import numpy as np
import argparse


def test_model_endpoint(port=5001, num_samples=5):
    """
    Test the MLflow model serving endpoint.
    
    Args:
        port: Port number where the model is served
        num_samples: Number of test samples to send
    """
    
    url = f"http://localhost:{port}/invocations"
    headers = {"Content-Type": "application/json"}
    
    print("=" * 80)
    print("  MLflow Model Serving Test")
    print("=" * 80)
    print(f"\nEndpoint: {url}")
    
    # Generate random test data (20 features to match the tutorial models)
    print(f"\nGenerating {num_samples} test samples with 20 features...")
    test_data = np.random.randn(num_samples, 20).tolist()
    
    payload = {
        "inputs": test_data
    }
    
    print(f"Data shape: ({num_samples}, 20)")
    print("\nSending request...")
    
    try:
        response = requests.post(
            url, 
            headers=headers, 
            data=json.dumps(payload),
            timeout=10
        )
        
        if response.status_code == 200:
            print("\n" + "=" * 80)
            print("  ✓ SUCCESS!")
            print("=" * 80)
            
            predictions = response.json()
            
            if isinstance(predictions, dict) and 'predictions' in predictions:
                preds = predictions['predictions']
            else:
                preds = predictions
            
            print(f"\nReceived {len(preds)} predictions:")
            print("-" * 80)
            
            for i, pred in enumerate(preds, 1):
                print(f"Sample {i}: {pred}")
            
            print("\n✓ Model is serving correctly!")
            return True
            
        else:
            print("\n" + "=" * 80)
            print("  ✗ ERROR")
            print("=" * 80)
            print(f"\nStatus Code: {response.status_code}")
            print(f"Response: {response.text}")
            return False
    
    except requests.exceptions.ConnectionError:
        print("\n" + "=" * 80)
        print("  ✗ CONNECTION ERROR")
        print("=" * 80)
        print("\nCould not connect to the server.")
        print("\nMake sure the model is serving:")
        print(f"  python -m mlflow models serve -m MODEL_URI -p {port} --no-conda")
        return False
        
    except requests.exceptions.Timeout:
        print("\n" + "=" * 80)
        print("  ✗ TIMEOUT ERROR")
        print("=" * 80)
        print("\nRequest timed out. The server might be starting up.")
        print("Wait a few seconds and try again.")
        return False
        
    except Exception as e:
        print("\n" + "=" * 80)
        print("  ✗ UNEXPECTED ERROR")
        print("=" * 80)
        print(f"\nError: {str(e)}")
        import traceback
        traceback.print_exc()
        return False


def get_health_check(port=5001):
    """Check if the server is running."""
    try:
        response = requests.get(f"http://localhost:{port}/ping", timeout=2)
        return response.status_code == 200
    except:
        return False


def main():
    """Main function."""
    parser = argparse.ArgumentParser(
        description="Test MLflow model serving endpoint"
    )
    parser.add_argument(
        "--port", 
        type=int, 
        default=5001,
        help="Port number (default: 5001)"
    )
    parser.add_argument(
        "--samples",
        type=int,
        default=5,
        help="Number of test samples (default: 5)"
    )
    
    args = parser.parse_args()
    
    # Check if server is up
    print("Checking if server is running...")
    if get_health_check(args.port):
        print("✓ Server is running\n")
    else:
        print("⚠ Server health check failed")
        print(f"Make sure model is serving on port {args.port}\n")
    
    # Test the endpoint
    success = test_model_endpoint(args.port, args.samples)
    
    if success:
        print("\n" + "=" * 80)
        print("  Test completed successfully!")
        print("=" * 80)
        print("\nYou can now make predictions using:")
        print(f"  curl -X POST http://localhost:{args.port}/invocations \\")
        print("       -H 'Content-Type: application/json' \\")
        print("       -d '{\"inputs\": [[...20 features...]]}'")
        return 0
    else:
        print("\n" + "=" * 80)
        print("  Test failed")
        print("=" * 80)
        return 1


if __name__ == "__main__":
    import sys
    sys.exit(main())