run.py

import argparse
import yaml
import pandas as pd
import numpy as np
import logging
import sys
import json
import time
import traceback

def setup_logging(log_file: str):
    """Configures Python logging to output to a file and the console."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )

def parse_args():
    parser = argparse.ArgumentParser(description="MLOps Batch Job")
    parser.add_argument('--input', required=True, help="Path to input CSV data")
    parser.add_argument('--config', required=True, help="Path to YAML config")
    parser.add_argument('--output', required=True, help="Path to metrics JSON")
    parser.add_argument('--log-file', required=True, help="Path to write logs")
    return parser.parse_args()

def write_metrics(output_path: str, payload: dict):
    """Writes the JSON payload and prints it to stdout."""
    with open(output_path, 'w') as f:
        json.dump(payload, f, indent=4)
    # The assessment specifically asks to print final metrics JSON to stdout
    print(json.dumps(payload, indent=4))

def load_and_validate_config(config_path: str) -> dict:
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)

    required_keys = ['seed', 'window', 'version']
    for key in required_keys:
        if key not in config:
            raise ValueError(f"Config is missing '{key}'")
            
    np.random.seed(config['seed'])
    logging.info(f"Config loaded and validated. Seed: {config['seed']}, Window: {config['window']}")
    return config

def process_data(filepath: str, window: int) -> pd.DataFrame:
    df = pd.read_csv(filepath)
    if 'close' not in df.columns or df.empty:
        raise ValueError("Invalid dataset: Missing 'close' column or empty data.")
    
    logging.info(f"Successfully loaded {len(df)} rows.")
    
    logging.info("Computing rolling mean...")
    df['rolling_mean'] = df['close'].rolling(window=window).mean()
    
    logging.info("Generating signals...")
    df['signal'] = (df['close'] > df['rolling_mean']).astype(int)
    
    return df

def main():
    start_time = time.time()
    
    # 1. We must parse arguments before we can set up the log file
    args = parse_args()
    setup_logging(args.log_file)
    logging.info("Job started.")
    
    version = "unknown"
    
    try:
        # 2. Load Config
        config = load_and_validate_config(args.config)
        version = config.get('version', 'unknown')
        
        # 3. Process Data
        df = process_data(args.input, config['window'])
        
        # 4. Compute Metrics
        latency_ms = int((time.time() - start_time) * 1000)
        signal_rate = df['signal'].mean()
        
        success_metrics = {
            "version": version,
            "rows_processed": len(df),
            "metric": "signal_rate",
            "value": round(float(signal_rate), 4),
            "latency_ms": latency_ms,
            "seed": config['seed'],
            "status": "success"
        }
        
        # 5. Write Output
        logging.info("Writing success metrics...")
        write_metrics(args.output, success_metrics)
        logging.info("Job completed successfully.")
        sys.exit(0)
        
    except Exception as e:
        # Catch ANY error, log the stack trace, and write the error metrics JSON
        logging.error(f"Job failed: {str(e)}")
        logging.debug(traceback.format_exc())
        
        error_metrics = {
            "version": version,
            "status": "error",
            "error_message": str(e)
        }
        write_metrics(args.output, error_metrics)
        sys.exit(1) # Must exit with a non-zero failure code

if __name__ == "__main__":
    main()