Biswajeet Arukha

fix: production hardening (mlflow optional + 403 forbidden bypass)

2d94b4d 8 months ago

10.2 kB

import os
import gc
import glob
import json
import logging
import numpy as np
import pandas as pd
import torch
from torch.utils.data import IterableDataset, DataLoader
import xgboost as xgb
from sklearn.model_selection import train_test_split
try:
    import mlflow
    import mlflow.pytorch
    import mlflow.xgboost
    MLFLOW_AVAILABLE = True
except ImportError:
    logger.warning("⚠️ mlflow not found. Training will proceed without remote tracking.")
    MLFLOW_AVAILABLE = False
from backend.ml_engine.hybrid_model import HybridEnsemble, ModelConfig, CricsheetNormalizer

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Base directory for relative paths
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

class CricsheetIterableDataset(IterableDataset):
    """
    An IterableDataset that loads and yields JSON match data one by one 
    preventing entire dataset from being loaded into RAM simultaneously.
    """
    def __init__(self, file_paths, normalizer):
        self.file_paths = file_paths
        self.normalizer = normalizer

    def __iter__(self):
        for filepath in self.file_paths:
            try:
                with open(filepath, 'r') as file:
                    match_data = json.load(file)
                    winner = match_data.get('info', {}).get('outcome', {}).get('winner')

                # Load one match at a time
                df = self.normalizer.load_cricsheet_data(filepath)
                # Create match features
                static_features = self.normalizer.create_match_features(df)
                sequence_features = self.normalizer._create_sequences(df, sequence_length=18)
                
                # Yield ball-by-ball samples
                for idx in range(len(sequence_features)):
                    seq = sequence_features.iloc[idx]['sequence']
                    batting_team = df.iloc[idx]['batting_team']
                    
                    if pd.isna(winner) or not batting_team:
                        label = 0
                    else:
                        label = 1 if self.normalizer.normalize_team(batting_team) == self.normalizer.normalize_team(winner) else 0
                    
                    yield torch.FloatTensor(seq), torch.FloatTensor([label])
            except Exception as e:
                logger.warning(f"Failed processing {filepath}: {e}")
            finally:
                # Force garbage collection to free RAM
                for var_name in ['df', 'static_features', 'sequence_features']:
                    if var_name in locals():
                        del locals()[var_name]
                gc.collect()

def train_efficiently(data_dir='data'):
    """
    Trains the XGBoost and LSTM models incrementally using low-RAM techniques.
    """
    logger.info("Initializing Memory-Efficient Training Pipeline...")
    normalizer = CricsheetNormalizer()
    normalizer.load_player_stats('backend/ml_engine/player_stats_db.json')
    normalizer.load_enhanced_features(
        weather_path='dataset/enhanced/weather_dataset.csv',
        elo_path='dataset/enhanced/team_elo.csv'
    )
    config = ModelConfig()
    ensemble = HybridEnsemble(config)


def _run_training(config, ensemble, normalizer, train_files, data_dir):
    """Inner training function — runs inside MLflow context."""

    # Training limit configurable via env var (0 = no limit)
    max_matches = int(os.getenv('MAX_TRAIN_MATCHES', '0'))
    if max_matches > 0:
        train_files = train_files[:max_matches]
        logger.info(f"Training limited to {max_matches} matches (set MAX_TRAIN_MATCHES=0 for full)")

    # ---------------------------------------------------------
    # PART A: Incremental Training for XGBoost
    # ---------------------------------------------------------
    logger.info("Starting XGBoost Incremental Training...")
    xgb_model = None
    params = config.xgb_params
    
    # Process files in chunks of 50 matches (very low RAM usage)
    chunk_size = 50
    for i in range(0, len(train_files), chunk_size):
        chunk_files = train_files[i:i + chunk_size]
        logger.info(f"XGBoost - Processing chunk {i//chunk_size + 1}/{(len(train_files)//chunk_size) + 1}")
        
        chunk_dfs = []
        for f in chunk_files:
            try:
                with open(f, 'r') as file:
                    match_data = json.load(file)
                    winner = match_data.get('info', {}).get('outcome', {}).get('winner')
                
                match_df = normalizer.load_cricsheet_data(f)
                match_df['winner'] = winner
                chunk_dfs.append(match_df)
            except json.JSONDecodeError:
                continue
        
        if not chunk_dfs:
            continue
            
        combined_df = pd.concat(chunk_dfs, ignore_index=True)
        static_features = normalizer.create_match_features(combined_df)
        
        # DEBUG: Check types
        print("Feature Types:")
        print(static_features[normalizer.STATIC_FEATURE_COLS].dtypes)
        
        # Extract Actual Labels
        def get_label(row):
            w = row.get('winner')
            b = row.get('batting_team')
            if pd.isna(w) or not b: return 0
            return 1 if normalizer.normalize_team(b) == normalizer.normalize_team(w) else 0
            
        static_features['label'] = static_features.apply(get_label, axis=1)
        
        # Use Standardized Columns List
        X_chunk = static_features[normalizer.STATIC_FEATURE_COLS].fillna(0).values
        y_chunk = static_features['label'].values
        
        dtrain = xgb.DMatrix(X_chunk, label=y_chunk)
        
        # Train incrementally (xgb_model handles starting from previous state)
        if xgb_model is None:
            xgb_model = xgb.train(params, dtrain, num_boost_round=10)
        else:
            xgb_model = xgb.train(params, dtrain, num_boost_round=10, xgb_model=xgb_model)
            
        # Log XGBoost status to MLflow
        if MLFLOW_AVAILABLE:
            try:
                mlflow.log_metric("xgb_chunk_processed", i//chunk_size + 1)
            except Exception:
                pass
            
        # Free memory!
        del chunk_dfs, combined_df, static_features, X_chunk, y_chunk, dtrain
        gc.collect()

    ensemble.xgb_model = xgb_model

    # ---------------------------------------------------------
    # PART B: Streaming Training for PyTorch LSTM
    # ---------------------------------------------------------
    logger.info("Starting LSTM Streaming Training...")
    ensemble.lstm_model = ensemble.build_lstm_model()
    
    criterion = torch.nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(ensemble.lstm_model.parameters(), lr=config.learning_rate)
    
    # Iterable datasets automatically stream from disk
    train_dataset = CricsheetIterableDataset(train_files, normalizer)
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size)
    
    ensemble.lstm_model.train()
    batch_count = 0
    
    for seq, label in train_loader:
        seq = seq.to(ensemble.device)
        label = label.to(ensemble.device)
        
        optimizer.zero_grad()
        features = ensemble.lstm_model(seq)
        predictions = features[:, 0]  # First feature logic
        
        loss = criterion(predictions, label.squeeze())
        loss.backward()
        optimizer.step()
        
        batch_count += 1
        if batch_count % 100 == 0:
            logger.info(f"LSTM Batch {batch_count}: Loss = {loss.item():.4f}")
            if MLFLOW_AVAILABLE:
                try:
                    mlflow.log_metric("lstm_loss", loss.item(), step=batch_count)
                except Exception:
                    pass
            gc.collect() # Periodically clean garbage within epoch
            
    # ---------------------------------------------------------
    # PART C: Streaming Training for PyTorch Transformer
    # ---------------------------------------------------------
    logger.info("Starting Transformer Streaming Training...")
    ensemble.transformer_model = ensemble.build_transformer_model()
    
    t_optimizer = torch.optim.Adam(ensemble.transformer_model.parameters(), lr=config.learning_rate)
    ensemble.transformer_model.train()
    
    t_batch_count = 0
    t_loader = DataLoader(train_dataset, batch_size=config.batch_size)
    
    for seq, label in t_loader:
        seq, label = seq.to(ensemble.device), label.to(ensemble.device)
        t_optimizer.zero_grad()
        
        features = ensemble.transformer_model(seq)
        predictions = features[:, 0]
        
        loss = criterion(predictions, label.squeeze())
        loss.backward()
        t_optimizer.step()
        
        t_batch_count += 1
        if t_batch_count % 100 == 0:
            logger.info(f"Transformer Batch {t_batch_count}: Loss = {loss.item():.4f}")
            if MLFLOW_AVAILABLE:
                try:
                    mlflow.log_metric("transformer_loss", loss.item(), step=t_batch_count)
                except Exception:
                    pass
            gc.collect()

    # Save the models
    logger.info("Training complete! Saving models...")
    os.makedirs('models', exist_ok=True)
    ensemble.save_models('models/hybrid_ensemble')
    
    # Log models to MLflow
    if MLFLOW_AVAILABLE:
        try:
            mlflow.xgboost.log_model(ensemble.xgb_model, "xgb_static")
            mlflow.pytorch.log_model(ensemble.lstm_model, "lstm_sequence")
            mlflow.pytorch.log_model(ensemble.transformer_model, "transformer_attention")
        except Exception as e:
            logger.warning(f"MLflow model logging failed: {e}")
    
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Memory-Efficient Training Pipeline")
    parser.add_argument('--data_dir', type=str, default='data', help='Path to cricsheet JSONs')
    args = parser.parse_args()
    
    os.makedirs(args.data_dir, exist_ok=True)
    train_efficiently(args.data_dir)

	import os
	import gc
	import glob
	import json
	import logging
	import numpy as np
	import pandas as pd
	import torch
	from torch.utils.data import IterableDataset, DataLoader
	import xgboost as xgb
	from sklearn.model_selection import train_test_split
	try:
	import mlflow
	import mlflow.pytorch
	import mlflow.xgboost
	MLFLOW_AVAILABLE = True
	except ImportError:
	logger.warning("⚠️ mlflow not found. Training will proceed without remote tracking.")
	MLFLOW_AVAILABLE = False
	from backend.ml_engine.hybrid_model import HybridEnsemble, ModelConfig, CricsheetNormalizer

	logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
	logger = logging.getLogger(__name__)

	# Base directory for relative paths
	BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

	class CricsheetIterableDataset(IterableDataset):
	"""
	An IterableDataset that loads and yields JSON match data one by one
	preventing entire dataset from being loaded into RAM simultaneously.
	"""
	def __init__(self, file_paths, normalizer):
	self.file_paths = file_paths
	self.normalizer = normalizer

	def __iter__(self):
	for filepath in self.file_paths:
	try:
	with open(filepath, 'r') as file:
	match_data = json.load(file)
	winner = match_data.get('info', {}).get('outcome', {}).get('winner')

	# Load one match at a time
	df = self.normalizer.load_cricsheet_data(filepath)
	# Create match features
	static_features = self.normalizer.create_match_features(df)
	sequence_features = self.normalizer._create_sequences(df, sequence_length=18)

	# Yield ball-by-ball samples
	for idx in range(len(sequence_features)):
	seq = sequence_features.iloc[idx]['sequence']
	batting_team = df.iloc[idx]['batting_team']

	if pd.isna(winner) or not batting_team:
	label = 0
	else:
	label = 1 if self.normalizer.normalize_team(batting_team) == self.normalizer.normalize_team(winner) else 0

	yield torch.FloatTensor(seq), torch.FloatTensor([label])
	except Exception as e:
	logger.warning(f"Failed processing {filepath}: {e}")
	finally:
	# Force garbage collection to free RAM
	for var_name in ['df', 'static_features', 'sequence_features']:
	if var_name in locals():
	del locals()[var_name]
	gc.collect()

	def train_efficiently(data_dir='data'):
	"""
	Trains the XGBoost and LSTM models incrementally using low-RAM techniques.
	"""
	logger.info("Initializing Memory-Efficient Training Pipeline...")
	normalizer = CricsheetNormalizer()
	normalizer.load_player_stats('backend/ml_engine/player_stats_db.json')
	normalizer.load_enhanced_features(
	weather_path='dataset/enhanced/weather_dataset.csv',
	elo_path='dataset/enhanced/team_elo.csv'
	)
	config = ModelConfig()
	ensemble = HybridEnsemble(config)


	def _run_training(config, ensemble, normalizer, train_files, data_dir):
	"""Inner training function — runs inside MLflow context."""

	# Training limit configurable via env var (0 = no limit)
	max_matches = int(os.getenv('MAX_TRAIN_MATCHES', '0'))
	if max_matches > 0:
	train_files = train_files[:max_matches]
	logger.info(f"Training limited to {max_matches} matches (set MAX_TRAIN_MATCHES=0 for full)")

	# ---------------------------------------------------------
	# PART A: Incremental Training for XGBoost
	# ---------------------------------------------------------
	logger.info("Starting XGBoost Incremental Training...")
	xgb_model = None
	params = config.xgb_params

	# Process files in chunks of 50 matches (very low RAM usage)
	chunk_size = 50
	for i in range(0, len(train_files), chunk_size):
	chunk_files = train_files[i:i + chunk_size]
	logger.info(f"XGBoost - Processing chunk {i//chunk_size + 1}/{(len(train_files)//chunk_size) + 1}")

	chunk_dfs = []
	for f in chunk_files:
	try:
	with open(f, 'r') as file:
	match_data = json.load(file)
	winner = match_data.get('info', {}).get('outcome', {}).get('winner')

	match_df = normalizer.load_cricsheet_data(f)
	match_df['winner'] = winner
	chunk_dfs.append(match_df)
	except json.JSONDecodeError:
	continue

	if not chunk_dfs:
	continue

	combined_df = pd.concat(chunk_dfs, ignore_index=True)
	static_features = normalizer.create_match_features(combined_df)

	# DEBUG: Check types
	print("Feature Types:")
	print(static_features[normalizer.STATIC_FEATURE_COLS].dtypes)

	# Extract Actual Labels
	def get_label(row):
	w = row.get('winner')
	b = row.get('batting_team')
	if pd.isna(w) or not b: return 0
	return 1 if normalizer.normalize_team(b) == normalizer.normalize_team(w) else 0

	static_features['label'] = static_features.apply(get_label, axis=1)

	# Use Standardized Columns List
	X_chunk = static_features[normalizer.STATIC_FEATURE_COLS].fillna(0).values
	y_chunk = static_features['label'].values

	dtrain = xgb.DMatrix(X_chunk, label=y_chunk)

	# Train incrementally (xgb_model handles starting from previous state)
	if xgb_model is None:
	xgb_model = xgb.train(params, dtrain, num_boost_round=10)
	else:
	xgb_model = xgb.train(params, dtrain, num_boost_round=10, xgb_model=xgb_model)

	# Log XGBoost status to MLflow
	if MLFLOW_AVAILABLE:
	try:
	mlflow.log_metric("xgb_chunk_processed", i//chunk_size + 1)
	except Exception:
	pass

	# Free memory!
	del chunk_dfs, combined_df, static_features, X_chunk, y_chunk, dtrain
	gc.collect()

	ensemble.xgb_model = xgb_model

	# ---------------------------------------------------------
	# PART B: Streaming Training for PyTorch LSTM
	# ---------------------------------------------------------
	logger.info("Starting LSTM Streaming Training...")
	ensemble.lstm_model = ensemble.build_lstm_model()

	criterion = torch.nn.BCEWithLogitsLoss()
	optimizer = torch.optim.Adam(ensemble.lstm_model.parameters(), lr=config.learning_rate)

	# Iterable datasets automatically stream from disk
	train_dataset = CricsheetIterableDataset(train_files, normalizer)
	train_loader = DataLoader(train_dataset, batch_size=config.batch_size)

	ensemble.lstm_model.train()
	batch_count = 0

	for seq, label in train_loader:
	seq = seq.to(ensemble.device)
	label = label.to(ensemble.device)

	optimizer.zero_grad()
	features = ensemble.lstm_model(seq)
	predictions = features[:, 0] # First feature logic

	loss = criterion(predictions, label.squeeze())
	loss.backward()
	optimizer.step()

	batch_count += 1
	if batch_count % 100 == 0:
	logger.info(f"LSTM Batch {batch_count}: Loss = {loss.item():.4f}")
	if MLFLOW_AVAILABLE:
	try:
	mlflow.log_metric("lstm_loss", loss.item(), step=batch_count)
	except Exception:
	pass
	gc.collect() # Periodically clean garbage within epoch

	# ---------------------------------------------------------
	# PART C: Streaming Training for PyTorch Transformer
	# ---------------------------------------------------------
	logger.info("Starting Transformer Streaming Training...")
	ensemble.transformer_model = ensemble.build_transformer_model()

	t_optimizer = torch.optim.Adam(ensemble.transformer_model.parameters(), lr=config.learning_rate)
	ensemble.transformer_model.train()

	t_batch_count = 0
	t_loader = DataLoader(train_dataset, batch_size=config.batch_size)

	for seq, label in t_loader:
	seq, label = seq.to(ensemble.device), label.to(ensemble.device)
	t_optimizer.zero_grad()

	features = ensemble.transformer_model(seq)
	predictions = features[:, 0]

	loss = criterion(predictions, label.squeeze())
	loss.backward()
	t_optimizer.step()

	t_batch_count += 1
	if t_batch_count % 100 == 0:
	logger.info(f"Transformer Batch {t_batch_count}: Loss = {loss.item():.4f}")
	if MLFLOW_AVAILABLE:
	try:
	mlflow.log_metric("transformer_loss", loss.item(), step=t_batch_count)
	except Exception:
	pass
	gc.collect()

	# Save the models
	logger.info("Training complete! Saving models...")
	os.makedirs('models', exist_ok=True)
	ensemble.save_models('models/hybrid_ensemble')

	# Log models to MLflow
	if MLFLOW_AVAILABLE:
	try:
	mlflow.xgboost.log_model(ensemble.xgb_model, "xgb_static")
	mlflow.pytorch.log_model(ensemble.lstm_model, "lstm_sequence")
	mlflow.pytorch.log_model(ensemble.transformer_model, "transformer_attention")
	except Exception as e:
	logger.warning(f"MLflow model logging failed: {e}")

	if __name__ == "__main__":
	import argparse
	parser = argparse.ArgumentParser(description="Memory-Efficient Training Pipeline")
	parser.add_argument('--data_dir', type=str, default='data', help='Path to cricsheet JSONs')
	args = parser.parse_args()

	os.makedirs(args.data_dir, exist_ok=True)
	train_efficiently(args.data_dir)