#!/usr/bin/env python3 """ Unit tests for hyperparameter_tuner.py Run with: python3 -m pytest test_hyperparameter_tuner.py -v """ import json import os import sys import tempfile from unittest.mock import Mock, patch, MagicMock import pytest import yaml # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from hyperparameter_tuner import GPUMonitor, GRPCModelTrainer, HyperparameterTuner class TestGPUMonitor: """Tests for GPU memory monitoring.""" def test_gpu_monitor_initialization_without_gpu(self): """Test GPU monitor handles missing GPU gracefully.""" with patch('hyperparameter_tuner.GPU_AVAILABLE', False): monitor = GPUMonitor() assert not monitor.enabled def test_get_memory_usage_without_gpu(self): """Test memory usage returns zeros when GPU unavailable.""" with patch('hyperparameter_tuner.GPU_AVAILABLE', False): monitor = GPUMonitor() usage = monitor.get_memory_usage() assert usage == {"used_gb": 0.0, "total_gb": 0.0, "percent": 0.0} def test_check_memory_available_without_gpu(self): """Test memory check returns False when GPU unavailable.""" with patch('hyperparameter_tuner.GPU_AVAILABLE', False): monitor = GPUMonitor() assert not monitor.check_memory_available(required_gb=2.0) @patch('hyperparameter_tuner.pynvml') def test_get_memory_usage_with_gpu(self, mock_pynvml): """Test memory usage calculation with GPU.""" # Mock GPU memory info mock_handle = Mock() mock_mem_info = Mock() mock_mem_info.used = 2 * (1024 ** 3) # 2GB mock_mem_info.total = 4 * (1024 ** 3) # 4GB mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_handle mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_mem_info with patch('hyperparameter_tuner.GPU_AVAILABLE', True): monitor = GPUMonitor() monitor.enabled = True usage = monitor.get_memory_usage() assert usage["used_gb"] == pytest.approx(2.0, rel=0.01) assert usage["total_gb"] == pytest.approx(4.0, rel=0.01) assert usage["percent"] == pytest.approx(50.0, rel=0.01) @patch('hyperparameter_tuner.pynvml') def test_check_memory_available_sufficient(self, mock_pynvml): """Test memory check passes when sufficient VRAM available.""" mock_handle = Mock() mock_mem_info = Mock() mock_mem_info.used = 1 * (1024 ** 3) # 1GB used mock_mem_info.total = 4 * (1024 ** 3) # 4GB total (3GB available) mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_handle mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_mem_info with patch('hyperparameter_tuner.GPU_AVAILABLE', True): monitor = GPUMonitor() monitor.enabled = True assert monitor.check_memory_available(required_gb=2.0) @patch('hyperparameter_tuner.pynvml') def test_check_memory_available_insufficient(self, mock_pynvml): """Test memory check fails when insufficient VRAM available.""" mock_handle = Mock() mock_mem_info = Mock() mock_mem_info.used = 3.5 * (1024 ** 3) # 3.5GB used mock_mem_info.total = 4 * (1024 ** 3) # 4GB total (0.5GB available) mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_handle mock_pynvml.nvmlDeviceGetMemoryInfo.return_value = mock_mem_info with patch('hyperparameter_tuner.GPU_AVAILABLE', True): monitor = GPUMonitor() monitor.enabled = True assert not monitor.check_memory_available(required_gb=2.0) class TestGRPCModelTrainer: """Tests for gRPC client.""" def test_initialization(self): """Test gRPC client initialization.""" client = GRPCModelTrainer(grpc_host="localhost", grpc_port=50054) assert client.address == "localhost:50054" assert client.channel is None assert client.stub is None @patch('hyperparameter_tuner.grpc.insecure_channel') def test_connect_success(self, mock_channel): """Test successful gRPC connection.""" # Mock channel and stub mock_channel_instance = Mock() mock_channel.return_value = mock_channel_instance with patch('hyperparameter_tuner.ml_training_pb2_grpc'): with patch('hyperparameter_tuner.ml_training_pb2'): client = GRPCModelTrainer() # Connection tested via HealthCheck in actual code # Just verify channel creation assert True # Placeholder for connection test def test_train_model_builds_correct_request(self): """Test TrainModel request construction.""" client = GRPCModelTrainer() # Mock stub mock_stub = Mock() mock_response = Mock() mock_response.success = True mock_response.sharpe_ratio = 1.5 mock_response.training_loss = 0.05 mock_response.validation_metrics = {"accuracy": 0.85} mock_response.error_message = "" mock_response.training_duration_seconds = 120 mock_stub.TrainModel.return_value = mock_response client.stub = mock_stub # Call train_model result = client.train_model( model_type="TLOB", hyperparameters={"learning_rate": 0.001, "epochs": 10.0}, data_source={"file_path": "/tmp/data.parquet"}, use_gpu=True, trial_id="trial_1" ) assert result["success"] assert result["sharpe_ratio"] == 1.5 assert result["training_loss"] == 0.05 assert result["validation_metrics"]["accuracy"] == 0.85 assert result["training_duration_seconds"] == 120 class TestHyperparameterTuner: """Tests for main tuner class.""" @pytest.fixture def config_file(self): """Create temporary config file.""" config = { "global": { "optimization_direction": "maximize", "pruning_enabled": True, "median_pruner": { "n_startup_trials": 5, "n_warmup_steps": 0, "interval_steps": 1 }, "sampler": "TPE" }, "models": { "TLOB": { "epochs": { "type": "int", "low": 10, "high": 50, "step": 10 }, "learning_rate": { "type": "float", "low": 0.0001, "high": 0.01, "log": True }, "batch_size": { "type": "categorical", "choices": [32, 64, 128] }, "use_positional_encoding": { "type": "categorical", "choices": [True, False] } } } } with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: yaml.dump(config, f) config_path = f.name yield config_path # Cleanup os.unlink(config_path) @pytest.fixture def storage_file(self): """Create temporary storage file.""" with tempfile.NamedTemporaryFile(suffix='.log', delete=False) as f: storage_path = f.name yield storage_path # Cleanup if os.path.exists(storage_path): os.unlink(storage_path) def test_tuner_initialization(self, config_file, storage_file): """Test tuner initialization.""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=10, config_path=config_file, data_source={"file_path": "/tmp/data.parquet"}, use_gpu=False, storage_path=storage_file ) assert tuner.job_id == "test_job" assert tuner.model_type == "TLOB" assert tuner.num_trials == 10 assert not tuner.use_gpu assert tuner.config is not None assert "TLOB" in tuner.config["models"] def test_suggest_hyperparameters_int(self, config_file, storage_file): """Test integer hyperparameter sampling.""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) # Mock Optuna trial mock_trial = Mock() mock_trial.suggest_int.return_value = 30 params = tuner.suggest_hyperparameters(mock_trial) assert "epochs" in params assert params["epochs"] == 30.0 # Converted to float for gRPC def test_suggest_hyperparameters_float(self, config_file, storage_file): """Test float hyperparameter sampling.""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) mock_trial = Mock() mock_trial.suggest_float.return_value = 0.001 params = tuner.suggest_hyperparameters(mock_trial) assert "learning_rate" in params assert params["learning_rate"] == 0.001 def test_suggest_hyperparameters_categorical_numeric(self, config_file, storage_file): """Test categorical hyperparameter sampling (numeric).""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) mock_trial = Mock() mock_trial.suggest_categorical.return_value = 64 params = tuner.suggest_hyperparameters(mock_trial) assert "batch_size" in params assert params["batch_size"] == 64.0 def test_suggest_hyperparameters_categorical_boolean(self, config_file, storage_file): """Test categorical hyperparameter sampling (boolean).""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) mock_trial = Mock() mock_trial.suggest_categorical.return_value = True params = tuner.suggest_hyperparameters(mock_trial) assert "use_positional_encoding" in params assert params["use_positional_encoding"] == 1.0 # Boolean -> float def test_objective_handles_training_failure(self, config_file, storage_file): """Test objective function handles training failures gracefully.""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) # Mock failed training tuner.grpc_client = Mock() tuner.grpc_client.train_model.return_value = { "success": False, "sharpe_ratio": 0.0, "training_loss": float('inf'), "validation_metrics": {}, "error_message": "GPU OOM", "training_duration_seconds": 0 } mock_trial = Mock() mock_trial.number = 1 mock_trial.suggest_int.return_value = 10 mock_trial.suggest_float.return_value = 0.001 mock_trial.suggest_categorical.side_effect = [64, True] objective_value = tuner.objective(mock_trial) assert objective_value == -999.0 # Worst possible Sharpe ratio def test_objective_returns_sharpe_ratio(self, config_file, storage_file): """Test objective function returns Sharpe ratio on success.""" tuner = HyperparameterTuner( job_id="test_job", model_type="TLOB", num_trials=1, config_path=config_file, data_source={}, use_gpu=False, storage_path=storage_file ) # Mock successful training tuner.grpc_client = Mock() tuner.grpc_client.train_model.return_value = { "success": True, "sharpe_ratio": 2.5, "training_loss": 0.02, "validation_metrics": {"accuracy": 0.9}, "error_message": "", "training_duration_seconds": 180 } mock_trial = Mock() mock_trial.number = 1 mock_trial.suggest_int.return_value = 10 mock_trial.suggest_float.return_value = 0.001 mock_trial.suggest_categorical.side_effect = [64, True] mock_trial.set_user_attr = Mock() objective_value = tuner.objective(mock_trial) assert objective_value == 2.5 # Verify user attributes were set assert mock_trial.set_user_attr.call_count >= 2 if __name__ == "__main__": pytest.main([__file__, "-v"])