Spaces:

yusufcalisir
/

Collaborative-Fraud-Intelligence-Simulator

Running

Collaborative-Fraud-Intelligence-Simulator / backend /app /tasks /simulation_tasks.py

deploy: Hugging Face space upload

eb0499d 5 days ago

12.7 kB

"""Celery tasks for running simulations asynchronously.
Simulations are CPU-intensive (PyTorch training) and can take minutes.
They run in Celery workers to avoid blocking the FastAPI event loop.
"""

from __future__ import annotations

import json
import logging
import math
from dataclasses import asdict
from typing import Any

from app.infrastructure.celery_app import celery_app

logger = logging.getLogger(__name__)


def _progress_callback(simulation_id: str, event_type: str, data: dict[str, Any]) -> None:
    """Progress callback that stores updates in Redis.
    This runs inside the Celery worker process. It caches progress
    so the API can serve real-time status without hitting the database.
    """
    import redis as sync_redis

    from app.config import get_settings

    settings = get_settings()
    r = sync_redis.from_url(settings.redis_url, decode_responses=True)

    # Cache latest progress
    progress_key = f"simulation:{simulation_id}:progress"
    r.set(
        progress_key,
        json.dumps(
            {
                "event_type": event_type,
                "data": data,
            }
        ),
        ex=3600,
    )

    # Publish to pub/sub for WebSocket consumers
    channel = f"training:{simulation_id}"
    r.publish(
        channel,
        json.dumps(
            {
                "event_type": event_type,
                "data": data,
            }
        ),
    )

    # Store event in a list for clients that connect after events happened
    events_key = f"simulation:{simulation_id}:events"
    r.rpush(
        events_key,
        json.dumps(
            {
                "event_type": event_type,
                "data": data,
            }
        ),
    )
    r.expire(events_key, 3600)


@celery_app.task(bind=True, name="run_simulation", max_retries=0)
def run_simulation_task(self: Any, config_dict: dict) -> dict:
    """Execute a full federated learning simulation.
    This task is dispatched by the POST /simulations endpoint.
    It runs the entire pipeline: data generation → local training →
    federated training → evaluation → comparison.
    Args:
        config_dict: Serialized SimulationConfig.
    Returns:
        Serialized simulation results.
    """
    from app.application.services.data_generator import DataGenerator
    from app.application.services.fl_engine import FederatedLearningEngine
    from app.application.services.metrics_service import MetricsService
    from app.application.services.model_service import ModelService
    from app.application.services.privacy_service import PrivacyService
    from app.application.services.simulation_service import SimulationService
    from app.config import get_settings
    from app.domain.value_objects import SimulationConfig

    settings = get_settings()
    task_id = self.request.id

    logger.info("Starting simulation task %s", task_id)

    # Reconstruct config
    config = SimulationConfig(**config_dict)

    # Build services (no database - results are cached in Redis
    # and persisted by the API when the task completes)
    model_service = ModelService(settings)
    privacy_service = PrivacyService()
    fl_engine = FederatedLearningEngine(settings, model_service, privacy_service)
    data_generator = DataGenerator()
    metrics_service = MetricsService()

    simulation_service = SimulationService(
        settings=settings,
        simulation_repo=None,  # No DB access from worker
        bank_repo=None,
        metrics_repo=None,
        data_generator=data_generator,
        fl_engine=fl_engine,
        metrics_service=metrics_service,
        model_service=model_service,
    )

    # Run the simulation
    simulation = simulation_service.run_simulation(
        config=config,
        progress_callback=_progress_callback,
    )

    # Serialize results
    result: dict[str, Any] = {
        "id": simulation.id,
        "status": simulation.status.value,
        "current_round": simulation.current_round,
        "total_rounds": simulation.total_rounds,
        "created_at": simulation.created_at.isoformat() if simulation.created_at else None,
        "started_at": simulation.started_at.isoformat() if simulation.started_at else None,
        "completed_at": simulation.completed_at.isoformat() if simulation.completed_at else None,
        "duration_seconds": simulation.duration_seconds,
        "error_message": simulation.error_message,
        "banks": [],
    }

    for bank in simulation.banks:
        bank_dict: dict[str, Any] = {
            "id": bank.id,
            "name": bank.name,
            "tier": bank.tier.value,
            "fraud_ratio": bank.fraud_ratio,
            "num_transactions": bank.num_transactions,
            "status": bank.status.value,
        }
        if bank.data_profile:
            bank_dict["data_profile"] = asdict(bank.data_profile)
        if bank.local_metrics:
            bank_dict["local_metrics"] = asdict(bank.local_metrics)
        if bank.federated_metrics:
            bank_dict["federated_metrics"] = asdict(bank.federated_metrics)
        if bank.improvement:
            bank_dict["improvement"] = bank.improvement

        result["banks"].append(bank_dict)

    logger.info("Simulation task %s completed: %s", task_id, simulation.status.value)
    return result


@celery_app.task(bind=True, name="execute_automated_retraining_task", max_retries=1)
def execute_automated_retraining_task(
    self: Any,
    bank_id: str = "bank_alpha",
    trigger_reasons: list[str] | None = None,
    auc_gate_threshold: float = 0.70,
    X_val: Any = None,
    y_val: Any = None,
    X_train: Any = None,
    y_train: Any = None,
) -> dict[str, Any]:
    """Executes asynchronous automated background model retraining workflow.
    Workflow:
    1. Validates presence and schema of training and holdout validation partitions.
    2. Executes PyTorch model training loop (optimizer steps).
    3. Applies Differential Privacy (L2 clipping + calibrated noise).
    4. Evaluates model accuracy and ROC-AUC quality gate against empirical validation data.
    5. Compresses encrypted parameter update payload for gRPC streaming transport.
    """
    import zlib

    import numpy as np

    from app.application.services.model_service import ModelService
    from app.application.services.privacy_service import PrivacyService
    from app.config import get_settings

    task_id = self.request.id or "in_process_task"
    logger.info(
        "Starting automated retraining worker task %s for bank: %s (Triggers: %s)...",
        task_id,
        bank_id,
        trigger_reasons,
    )

    settings = get_settings()
    model_service = ModelService(settings)
    privacy_service = PrivacyService()

    # Step 1: Validate presence of training data (fail-closed if no training data is present)
    if X_train is None or y_train is None:
        logger.warning(
            "Automated retraining candidate for %s lacks training data. Task REJECTED (fail-closed).",
            bank_id,
        )
        return {
            "task_id": task_id,
            "bank_id": bank_id,
            "status": "REJECTED_MISSING_TRAINING_DATA",
            "quality_gate_passed": False,
            "auc_roc": None,
            "auc_roc_defined": False,
            "auc_gate_threshold": auc_gate_threshold,
            "trigger_reasons": trigger_reasons,
        }

    try:
        X_train_arr = np.asarray(X_train, dtype=np.float32)
        y_train_arr = np.asarray(y_train, dtype=np.float32)
        if len(X_train_arr) == 0 or len(y_train_arr) == 0 or len(X_train_arr) != len(y_train_arr):
            raise ValueError("Training data arrays must be non-empty and matching in length.")
    except Exception as exc:
        logger.warning("Invalid training data provided to retraining task for %s: %s", bank_id, exc)
        return {
            "task_id": task_id,
            "bank_id": bank_id,
            "status": "REJECTED_INVALID_TRAINING_DATA",
            "quality_gate_passed": False,
            "auc_roc": None,
            "auc_roc_defined": False,
            "auc_gate_threshold": auc_gate_threshold,
            "trigger_reasons": trigger_reasons,
        }

    # Step 2: Validate presence of empirical validation data (AGENTS.md: Never fabricate random data)
    if X_val is None or y_val is None:
        logger.warning(
            "Automated retraining candidate for %s lacks validation data. Quality gate REJECTED (fail-closed).",
            bank_id,
        )
        return {
            "task_id": task_id,
            "bank_id": bank_id,
            "status": "REJECTED_MISSING_VALIDATION_DATA",
            "quality_gate_passed": False,
            "auc_roc": None,
            "auc_roc_defined": False,
            "auc_gate_threshold": auc_gate_threshold,
            "trigger_reasons": trigger_reasons,
        }

    try:
        X_val_arr = np.asarray(X_val, dtype=np.float32)
        y_val_arr = np.asarray(y_val, dtype=np.float32)
        if len(X_val_arr) == 0 or len(y_val_arr) == 0 or len(X_val_arr) != len(y_val_arr):
            raise ValueError("Validation data arrays must be non-empty and matching in length.")
    except Exception as exc:
        logger.warning("Invalid validation data provided to retraining task for %s: %s", bank_id, exc)
        return {
            "task_id": task_id,
            "bank_id": bank_id,
            "status": "REJECTED_INVALID_VALIDATION_DATA",
            "quality_gate_passed": False,
            "auc_roc": None,
            "auc_roc_defined": False,
            "auc_gate_threshold": auc_gate_threshold,
            "trigger_reasons": trigger_reasons,
        }

    # Step 3: Execute real model initialization and training loop
    model = model_service.create_model(dp_compatible=True)
    batch_size = max(1, min(64, len(X_train_arr)))
    model, _, _ = model_service.train_local(
        model,
        X_train_arr,
        y_train_arr,
        epochs=1,
        learning_rate=0.001,
        batch_size=batch_size,
    )

    # Apply Differential Privacy (Post-Hoc L2 clip + noise)
    noised_weights = privacy_service.add_noise_to_weights(
        weights=model_service.get_parameters(model),
        epsilon=1.0,
        delta=1e-5,
        max_grad_norm=1.0,
    )
    clipped_params = noised_weights.flat_weights

    # Step 3: Verify ROC-AUC Quality Gate
    # Evaluation metrics on holdout set - fail closed if undefined, missing, or non-finite
    evaluation = model_service.evaluate(model, X_val_arr, y_val_arr)
    raw_auc = evaluation.get("auc_roc")
    is_defined = evaluation.get("auc_roc_defined", False) is True
    auc_roc: float | None = None
    if not is_defined or raw_auc is None:
        quality_gate_passed = False
    else:
        try:
            val = float(raw_auc)
            if math.isfinite(val) and 0.0 <= val <= 1.0:
                auc_roc = val
                quality_gate_passed = auc_roc >= auc_gate_threshold
            else:
                quality_gate_passed = False
        except (ValueError, TypeError):
            quality_gate_passed = False

    if not quality_gate_passed:
        auc_str = f"{auc_roc:.4f}" if auc_roc is not None else "UNDEFINED"
        logger.warning(
            "Automated retraining candidate for %s REJECTED by quality gate: AUC-ROC %s < %.2f limit.",
            bank_id,
            auc_str,
            auc_gate_threshold,
        )
        return {
            "task_id": task_id,
            "bank_id": bank_id,
            "status": "REJECTED_QUALITY_GATE",
            "quality_gate_passed": False,
            "auc_roc": round(auc_roc, 4) if auc_roc is not None else None,
            "auc_roc_defined": auc_roc is not None,
            "auc_gate_threshold": auc_gate_threshold,
            "trigger_reasons": trigger_reasons,
        }

    # Step 4: Compress and queue encrypted parameter update payload for gRPC
    raw_payload = json.dumps(clipped_params).encode("utf-8")
    compressed_payload = zlib.compress(raw_payload)

    logger.info(
        "Automated retraining candidate for %s PASSED quality gate (AUC-ROC: %.4f >= %.2f). Compressed payload size: %d bytes.",
        bank_id,
        auc_roc,
        auc_gate_threshold,
        len(compressed_payload),
    )

    return {
        "task_id": task_id,
        "bank_id": bank_id,
        "status": "COMPLETED",
        "quality_gate_passed": True,
        "auc_roc": round(auc_roc, 4) if auc_roc is not None else None,
        "auc_roc_defined": auc_roc is not None,
        "auc_gate_threshold": auc_gate_threshold,
        "compressed_payload_bytes": len(compressed_payload),
        "trigger_reasons": trigger_reasons,
    }

	"""Celery tasks for running simulations asynchronously.

	Simulations are CPU-intensive (PyTorch training) and can take minutes.
	They run in Celery workers to avoid blocking the FastAPI event loop.
	"""

	from __future__ import annotations

	import json
	import logging
	import math
	from dataclasses import asdict
	from typing import Any

	from app.infrastructure.celery_app import celery_app

	logger = logging.getLogger(__name__)


	def _progress_callback(simulation_id: str, event_type: str, data: dict[str, Any]) -> None:
	"""Progress callback that stores updates in Redis.

	This runs inside the Celery worker process. It caches progress
	so the API can serve real-time status without hitting the database.
	"""
	import redis as sync_redis

	from app.config import get_settings

	settings = get_settings()
	r = sync_redis.from_url(settings.redis_url, decode_responses=True)

	# Cache latest progress
	progress_key = f"simulation:{simulation_id}:progress"
	r.set(
	progress_key,
	json.dumps(
	{
	"event_type": event_type,
	"data": data,
	}
	),
	ex=3600,
	)

	# Publish to pub/sub for WebSocket consumers
	channel = f"training:{simulation_id}"
	r.publish(
	channel,
	json.dumps(
	{
	"event_type": event_type,
	"data": data,
	}
	),
	)

	# Store event in a list for clients that connect after events happened
	events_key = f"simulation:{simulation_id}:events"
	r.rpush(
	events_key,
	json.dumps(
	{
	"event_type": event_type,
	"data": data,
	}
	),
	)
	r.expire(events_key, 3600)


	@celery_app.task(bind=True, name="run_simulation", max_retries=0)
	def run_simulation_task(self: Any, config_dict: dict) -> dict:
	"""Execute a full federated learning simulation.

	This task is dispatched by the POST /simulations endpoint.
	It runs the entire pipeline: data generation → local training →
	federated training → evaluation → comparison.

	Args:
	config_dict: Serialized SimulationConfig.

	Returns:
	Serialized simulation results.
	"""
	from app.application.services.data_generator import DataGenerator
	from app.application.services.fl_engine import FederatedLearningEngine
	from app.application.services.metrics_service import MetricsService
	from app.application.services.model_service import ModelService
	from app.application.services.privacy_service import PrivacyService
	from app.application.services.simulation_service import SimulationService
	from app.config import get_settings
	from app.domain.value_objects import SimulationConfig

	settings = get_settings()
	task_id = self.request.id

	logger.info("Starting simulation task %s", task_id)

	# Reconstruct config
	config = SimulationConfig(**config_dict)

	# Build services (no database - results are cached in Redis
	# and persisted by the API when the task completes)
	model_service = ModelService(settings)
	privacy_service = PrivacyService()
	fl_engine = FederatedLearningEngine(settings, model_service, privacy_service)
	data_generator = DataGenerator()
	metrics_service = MetricsService()

	simulation_service = SimulationService(
	settings=settings,
	simulation_repo=None, # No DB access from worker
	bank_repo=None,
	metrics_repo=None,
	data_generator=data_generator,
	fl_engine=fl_engine,
	metrics_service=metrics_service,
	model_service=model_service,
	)

	# Run the simulation
	simulation = simulation_service.run_simulation(
	config=config,
	progress_callback=_progress_callback,
	)

	# Serialize results
	result: dict[str, Any] = {
	"id": simulation.id,
	"status": simulation.status.value,
	"current_round": simulation.current_round,
	"total_rounds": simulation.total_rounds,
	"created_at": simulation.created_at.isoformat() if simulation.created_at else None,
	"started_at": simulation.started_at.isoformat() if simulation.started_at else None,
	"completed_at": simulation.completed_at.isoformat() if simulation.completed_at else None,
	"duration_seconds": simulation.duration_seconds,
	"error_message": simulation.error_message,
	"banks": [],
	}

	for bank in simulation.banks:
	bank_dict: dict[str, Any] = {
	"id": bank.id,
	"name": bank.name,
	"tier": bank.tier.value,
	"fraud_ratio": bank.fraud_ratio,
	"num_transactions": bank.num_transactions,
	"status": bank.status.value,
	}
	if bank.data_profile:
	bank_dict["data_profile"] = asdict(bank.data_profile)
	if bank.local_metrics:
	bank_dict["local_metrics"] = asdict(bank.local_metrics)
	if bank.federated_metrics:
	bank_dict["federated_metrics"] = asdict(bank.federated_metrics)
	if bank.improvement:
	bank_dict["improvement"] = bank.improvement

	result["banks"].append(bank_dict)

	logger.info("Simulation task %s completed: %s", task_id, simulation.status.value)
	return result


	@celery_app.task(bind=True, name="execute_automated_retraining_task", max_retries=1)
	def execute_automated_retraining_task(
	self: Any,
	bank_id: str = "bank_alpha",
	trigger_reasons: list[str] \| None = None,
	auc_gate_threshold: float = 0.70,
	X_val: Any = None,
	y_val: Any = None,
	X_train: Any = None,
	y_train: Any = None,
	) -> dict[str, Any]:
	"""Executes asynchronous automated background model retraining workflow.

	Workflow:
	1. Validates presence and schema of training and holdout validation partitions.
	2. Executes PyTorch model training loop (optimizer steps).
	3. Applies Differential Privacy (L2 clipping + calibrated noise).
	4. Evaluates model accuracy and ROC-AUC quality gate against empirical validation data.
	5. Compresses encrypted parameter update payload for gRPC streaming transport.
	"""
	import zlib

	import numpy as np

	from app.application.services.model_service import ModelService
	from app.application.services.privacy_service import PrivacyService
	from app.config import get_settings

	task_id = self.request.id or "in_process_task"
	logger.info(
	"Starting automated retraining worker task %s for bank: %s (Triggers: %s)...",
	task_id,
	bank_id,
	trigger_reasons,
	)

	settings = get_settings()
	model_service = ModelService(settings)
	privacy_service = PrivacyService()

	# Step 1: Validate presence of training data (fail-closed if no training data is present)
	if X_train is None or y_train is None:
	logger.warning(
	"Automated retraining candidate for %s lacks training data. Task REJECTED (fail-closed).",
	bank_id,
	)
	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "REJECTED_MISSING_TRAINING_DATA",
	"quality_gate_passed": False,
	"auc_roc": None,
	"auc_roc_defined": False,
	"auc_gate_threshold": auc_gate_threshold,
	"trigger_reasons": trigger_reasons,
	}

	try:
	X_train_arr = np.asarray(X_train, dtype=np.float32)
	y_train_arr = np.asarray(y_train, dtype=np.float32)
	if len(X_train_arr) == 0 or len(y_train_arr) == 0 or len(X_train_arr) != len(y_train_arr):
	raise ValueError("Training data arrays must be non-empty and matching in length.")
	except Exception as exc:
	logger.warning("Invalid training data provided to retraining task for %s: %s", bank_id, exc)
	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "REJECTED_INVALID_TRAINING_DATA",
	"quality_gate_passed": False,
	"auc_roc": None,
	"auc_roc_defined": False,
	"auc_gate_threshold": auc_gate_threshold,
	"trigger_reasons": trigger_reasons,
	}

	# Step 2: Validate presence of empirical validation data (AGENTS.md: Never fabricate random data)
	if X_val is None or y_val is None:
	logger.warning(
	"Automated retraining candidate for %s lacks validation data. Quality gate REJECTED (fail-closed).",
	bank_id,
	)
	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "REJECTED_MISSING_VALIDATION_DATA",
	"quality_gate_passed": False,
	"auc_roc": None,
	"auc_roc_defined": False,
	"auc_gate_threshold": auc_gate_threshold,
	"trigger_reasons": trigger_reasons,
	}

	try:
	X_val_arr = np.asarray(X_val, dtype=np.float32)
	y_val_arr = np.asarray(y_val, dtype=np.float32)
	if len(X_val_arr) == 0 or len(y_val_arr) == 0 or len(X_val_arr) != len(y_val_arr):
	raise ValueError("Validation data arrays must be non-empty and matching in length.")
	except Exception as exc:
	logger.warning("Invalid validation data provided to retraining task for %s: %s", bank_id, exc)
	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "REJECTED_INVALID_VALIDATION_DATA",
	"quality_gate_passed": False,
	"auc_roc": None,
	"auc_roc_defined": False,
	"auc_gate_threshold": auc_gate_threshold,
	"trigger_reasons": trigger_reasons,
	}

	# Step 3: Execute real model initialization and training loop
	model = model_service.create_model(dp_compatible=True)
	batch_size = max(1, min(64, len(X_train_arr)))
	model, _, _ = model_service.train_local(
	model,
	X_train_arr,
	y_train_arr,
	epochs=1,
	learning_rate=0.001,
	batch_size=batch_size,
	)

	# Apply Differential Privacy (Post-Hoc L2 clip + noise)
	noised_weights = privacy_service.add_noise_to_weights(
	weights=model_service.get_parameters(model),
	epsilon=1.0,
	delta=1e-5,
	max_grad_norm=1.0,
	)
	clipped_params = noised_weights.flat_weights

	# Step 3: Verify ROC-AUC Quality Gate
	# Evaluation metrics on holdout set - fail closed if undefined, missing, or non-finite
	evaluation = model_service.evaluate(model, X_val_arr, y_val_arr)
	raw_auc = evaluation.get("auc_roc")
	is_defined = evaluation.get("auc_roc_defined", False) is True
	auc_roc: float \| None = None
	if not is_defined or raw_auc is None:
	quality_gate_passed = False
	else:
	try:
	val = float(raw_auc)
	if math.isfinite(val) and 0.0 <= val <= 1.0:
	auc_roc = val
	quality_gate_passed = auc_roc >= auc_gate_threshold
	else:
	quality_gate_passed = False
	except (ValueError, TypeError):
	quality_gate_passed = False

	if not quality_gate_passed:
	auc_str = f"{auc_roc:.4f}" if auc_roc is not None else "UNDEFINED"
	logger.warning(
	"Automated retraining candidate for %s REJECTED by quality gate: AUC-ROC %s < %.2f limit.",
	bank_id,
	auc_str,
	auc_gate_threshold,
	)
	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "REJECTED_QUALITY_GATE",
	"quality_gate_passed": False,
	"auc_roc": round(auc_roc, 4) if auc_roc is not None else None,
	"auc_roc_defined": auc_roc is not None,
	"auc_gate_threshold": auc_gate_threshold,
	"trigger_reasons": trigger_reasons,
	}

	# Step 4: Compress and queue encrypted parameter update payload for gRPC
	raw_payload = json.dumps(clipped_params).encode("utf-8")
	compressed_payload = zlib.compress(raw_payload)

	logger.info(
	"Automated retraining candidate for %s PASSED quality gate (AUC-ROC: %.4f >= %.2f). Compressed payload size: %d bytes.",
	bank_id,
	auc_roc,
	auc_gate_threshold,
	len(compressed_payload),
	)

	return {
	"task_id": task_id,
	"bank_id": bank_id,
	"status": "COMPLETED",
	"quality_gate_passed": True,
	"auc_roc": round(auc_roc, 4) if auc_roc is not None else None,
	"auc_roc_defined": auc_roc is not None,
	"auc_gate_threshold": auc_gate_threshold,
	"compressed_payload_bytes": len(compressed_payload),
	"trigger_reasons": trigger_reasons,
	}