File size: 2,931 Bytes
eb0499d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import os
import sys
# Force CPU execution to prevent PyTorch CUDA hangs
os.environ["CUDA_VISIBLE_DEVICES"] = ""
# Add backend root to path
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.abspath(os.path.join(SCRIPT_DIR, "..")))
from app.application.services.data_generator import DataGenerator # noqa: E402
from app.application.services.fl_engine import FederatedLearningEngine # noqa: E402
from app.application.services.metrics_service import MetricsService # noqa: E402
from app.application.services.model_service import ModelService # noqa: E402
from app.application.services.privacy_service import PrivacyService # noqa: E402
from app.application.services.simulation_service import SimulationService # noqa: E402
from app.config import get_settings # noqa: E402
from app.domain.value_objects import SimulationConfig # noqa: E402
def main():
print("Testing MLflow integration via SimulationService...")
settings = get_settings()
# Enable MLflow and customize experiment name
settings.mlflow_enabled = True
settings.mlflow_experiment_name = "Test-CFI-MLflow-Integration"
# Instantiate actual service layers
model_service = ModelService(settings)
privacy_service = PrivacyService()
fl_engine = FederatedLearningEngine(settings, model_service, privacy_service)
metrics_service = MetricsService()
data_generator = DataGenerator(seed=42)
# Instantiate SimulationService with mock repositories
simulation_service = SimulationService(
settings=settings,
simulation_repo=None,
bank_repo=None,
metrics_repo=None,
data_generator=data_generator,
fl_engine=fl_engine,
metrics_service=metrics_service,
model_service=model_service,
)
# Configure a tiny 2-round simulation for rapid local testing
config = SimulationConfig(
num_rounds=2,
local_epochs=1,
learning_rate=0.01,
batch_size=64,
min_clients_per_round=2,
bank_a_transactions=1000,
bank_b_transactions=800,
bank_c_transactions=600,
enable_differential_privacy=True,
dp_epsilon=3.0,
dp_delta=1e-5,
)
print("Starting simulation run...")
simulation = simulation_service.run_simulation(config)
print(f"Simulation completed with status: {simulation.status}")
# Verify that the MLflow tracking runs were generated
print("Verifying MLflow tracking files...")
mlruns_dir = os.path.join(SCRIPT_DIR, "..", "mlruns")
if os.path.exists(mlruns_dir):
print(f"Success: 'mlruns' directory created at: {os.path.abspath(mlruns_dir)}")
# Check files inside
experiments = os.listdir(mlruns_dir)
print(f"Active experiments in mlruns: {experiments}")
else:
print("Error: 'mlruns' directory not found!")
sys.exit(1)
if __name__ == "__main__":
main()
|