run.py

import argparse
import json
import logging
import os
import time
import numpy as np
import pandas as pd
import yaml


def setup_logging(log_file: str):
    """Setup Python logging configuration [cite: 9]"""
    logging.basicConfig(
        filename=log_file,
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def parse_cli_args():
    """Parse and return command-line arguments [cite: 2]"""
    parser = argparse.ArgumentParser(
        description="MLOps Batch Job for Trading Signals"
    )
    parser.add_argument(
        "--input", required=True, help="Path to input data.csv"
    )
    parser.add_argument("--config", required=True, help="Path to config.yaml")
    parser.add_argument(
        "--output", required=True, help="Path to output metrics.json"
    )
    parser.add_argument(
        "--log-file", required=True, help="Path to output run.log"
    )
    return parser.parse_args()


def load_and_validate_config(config_path: str) -> dict:
    """Loads YAML config and validates presence of required fields [cite: 4]"""
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Config file not found: {config_path}")

    logging.info(f"Loading config from {config_path}")
    with open(config_path, "r") as f:
        try:
            config = yaml.safe_load(f)
        except yaml.YAMLError as e:
            raise ValueError(f"Invalid YAML format: {e}")

    # Explicit validation check [cite: 4]
    required_fields = ["seed", "window", "version"]
    if not config or not all(k in config for k in required_fields):
        raise ValueError(
            f"Config missing required fields. Expected: {required_fields}"
        )

    return config


def load_and_validate_dataset(input_path: str) -> pd.DataFrame:
    """
    Loads CSV dataset and validates required columns.
    Also handles malformed CSV files where each row is enclosed in quotes.
    """

    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input file not found: {input_path}")

    if os.path.getsize(input_path) == 0:
        raise ValueError("Input file is empty")

    logging.info(f"Loading dataset from {input_path}")

    try:
        # Try normal CSV loading first
        df = pd.read_csv(input_path)

        # If pandas reads the entire header as one column,
        # the CSV is malformed (entire row enclosed in quotes)
        if len(df.columns) == 1 and "," in df.columns[0]:
            raise ValueError("Malformed CSV detected")

    except Exception:
        logging.info("Malformed CSV detected. Applying custom parser...")

        try:
            raw_df = pd.read_csv(input_path, header=None)

            # Extract header
            headers = raw_df.iloc[0, 0].replace('"', "").split(",")

            # Extract rows
            rows = (
                raw_df.iloc[1:, 0]
                .str.replace('"', "", regex=False)
                .str.split(",")
            )

            df = pd.DataFrame(rows.tolist(), columns=headers)

        except Exception as e:
            raise ValueError(f"Invalid CSV format: {e}")

    # Remove unwanted quotes/spaces from column names
    df.columns = df.columns.str.replace('"', "", regex=False).str.strip()

    # Convert numeric columns where possible
    for col in df.columns:
        try:
            df[col] = pd.to_numeric(df[col])
        except (ValueError, TypeError):
            pass

    # Validate required column
    if "close" not in df.columns:
        raise ValueError(
            f"Missing required column: 'close'. Found columns: {list(df.columns)}"
        )

    return df


def compute_trading_signals(df: pd.DataFrame, window: int) -> float:
    """Computes rolling mean and generates binary trading signals [cite: 5, 7]"""
    logging.info("Processing steps: computing rolling mean...")
    df["rolling_mean"] = df["close"].rolling(window=window).mean()

    logging.info("Processing steps: generating binary signal...")
    df["signal"] = 0
    
    valid_mask = df["rolling_mean"].notna()
    df.loc[valid_mask & (df["close"] > df["rolling_mean"]), "signal"] = 1

    return float(df["signal"].mean())


def write_metrics(output_path: str, metrics: dict):
    """Writes metrics JSON to file and prints directly to stdout [cite: 8, 11]"""
    with open(output_path, "w") as f:
        json.dump(metrics, f, indent=4)
   
    print(json.dumps(metrics, indent=4))


def main():
    start_time = time.time()
    args = parse_cli_args()

    setup_logging(args.log_file)
    logging.info("Job started. [cite: 9]")

    version = "unknown"
    seed = None

    try:
        
        config = load_and_validate_config(args.config)
        version = config["version"]
        seed = config["seed"]
        window = config["window"]

        
        np.random.seed(seed)
        logging.info(
            f"Config loaded + validated. Version: {version}, Seed: {seed}, Window: {window} [cite: 9]"
        )


        df = load_and_validate_dataset(args.input)
        rows_processed = len(df)
        logging.info(f"Rows loaded successfully: {rows_processed} [cite: 9]")

        signal_rate = compute_trading_signals(df, window)

        latency_ms = int((time.time() - start_time) * 1000)
        metrics = {
            "version": version,
            "rows_processed": rows_processed,
            "metric": "signal_rate",
            "value": round(signal_rate, 4),
            "latency_ms": latency_ms,
            "seed": seed,
            "status": "success",
        }

        logging.info("Metrics summary calculated. [cite: 9]")
        logging.info("Job end + status: success [cite: 9]")
        write_metrics(args.output, metrics)

    except Exception as e:
        logging.error(f"Validation error or exception occurred: {str(e)} [cite: 9]")
        latency_ms = int((time.time() - start_time) * 1000)

        error_metrics = {
            "version": str(version),
            "status": "error",
            "error_message": str(e),
        }
        logging.info("Job end + status: error [cite: 9]")
        write_metrics(args.output, error_metrics)
        exit(1)  


if __name__ == "__main__":
    main()