main.py

import atexit
import time

from dotenv import load_dotenv
from opentelemetry import trace
from tabulate import tabulate

from src import determine_dataset_type
from src.applications.app_factory import app_factory
from src.registries.topologies.topology_registry import create_adjacency_matrix
from src.tracing.metrics import setup_metrics
from src.tracing.tracing import setup_tracing
from src.utils import (
    client_ids_list_generator,
    log_path,
    none_checker,
    var_name,
)
from src.utils.datasets.loader_storage import persist_loaders_to_disk
from src.utils.framework.framework_setup import FrameworkSetup
from src.utils.framework.log import Log
from src.utils.framework.ray import generate_unique_resource_label
from src.utils.framework.yaml_loader import load_objectified_yaml
from src.validators.config_validator import ConfigValidator
from src.validators.runtime_config import RuntimeConfig


@atexit.register
def flush_all_traces():
    try:
        trace.get_tracer_provider().shutdown()
        print("[otel] Tracer provider shutdown cleanly.")
    except Exception as e:
        print(f"[otel] Failed to flush traces at exit: {e}")


def main(config_yaml_path: str = "./config.yaml"):
    load_dotenv(override=True)
    config_yaml_path = none_checker(config_yaml_path, var_name(config_yaml_path))

    config_dict = load_objectified_yaml(config_yaml_path)

    config_dict = config_dict | {"desired_distribution": None}  # TODO: update

    config = ConfigValidator(**config_dict)

    setup_metrics(config)
    setup_tracing(config)

    log_file = log_path(
        model_type=config.MODEL_TYPE,
        dataset_type=config.DATASET_TYPE,
        data_distribution=config.DATA_DISTRIBUTION,
        distance_metric=config.DISTANCE_METRIC,
        sensitivity_percentage=config.SENSITIVITY_PERCENTAGE,
        fed_avg=config.FED_AVG,
        dynamic_sensitivity_percentage=config.DYNAMIC_SENSITIVITY_PERCENTAGE,
        distance_metric_on_parameters=config.DISTANCE_METRIC_ON_PARAMETERS,
        pre_computed_data_driven_clustering=config.PRE_COMPUTED_DATA_DRIVEN_CLUSTERING,
        remove_common_ids=config.REMOVE_COMMON_IDS,
    )

    log = Log(
        log_file,
        config.MODEL_TYPE,
        config.DISTANCE_METRIC,
        create_file=False,
    )

    table_data = [
        [key, value]
        for key, value in config.__dict__.items()
        if not key.startswith("_")
    ]
    log.info(tabulate(table_data, headers=["Config Key", "Value"], tablefmt="grid"))

    log.section("FRAMEWORK SETUP")

    FrameworkSetup.path_setup(config)

    train_loaders, test_loaders, machine_label = None, None, None
    clients_id_list = []
    if config.PRODUCTION_MODE:
        machine_label = generate_unique_resource_label(log)
    else:
        log.section("DATASETS DISTRIBUTION")

        train_loaders, test_loaders = determine_dataset_type(config, log)
        clients_id_list = client_ids_list_generator(config.NUMBER_OF_CLIENTS, log=log)

    if config.PRE_COMPUTED_DATA_DRIVEN_CLUSTERING:
        log.info("clients train loader label distribution")

    train_loader_paths = None
    test_loader_paths = None
    if train_loaders is not None and test_loaders is not None:
        manifest = persist_loaders_to_disk(
            train_loaders=train_loaders,
            test_loaders=test_loaders,
            federation_id=config.FEDERATION_ID,
            run_name="main",
        )
        train_loader_paths = manifest["train"]
        test_loader_paths = manifest["test"]

    log.section("RUNTIME CONFIGURATIONS")

    config.RUNTIME_CONFIG = RuntimeConfig(
        clients_id_list=clients_id_list,
        train_loaders=train_loaders,
        test_loaders=test_loaders,
        log=log,
        machine_label=machine_label,
        is_head=True,
        adjacency_matrix=create_adjacency_matrix(config),
        train_loader_paths=train_loader_paths,
        test_loader_paths=test_loader_paths,
    )

    log.section("APP FACTORY")

    app_fn = app_factory(config=config, log=log)

    s = time.perf_counter()
    app_fn(config, log)
    e = time.perf_counter()
    print(f"Whole FL run took {(e - s) / 60} minutes")

    try:
        from src.utils.logging.download_logs import (
            download_csv_metrics_from_ray_standalone,
        )

        result = download_csv_metrics_from_ray_standalone(config)
        if result:
            print("Successfully downloaded CSV metrics from Ray logs")
        else:
            print("No CSV metrics found to download or not running on head node")
    except Exception as e:
        print(f"Warning: Failed to download CSV metrics: {e}")


if __name__ == "__main__":
    main()