[ENH] Improve logging and extract utilities

scripts/_io_utils.py

@@ -0,0 +1,76 @@
+"""Output paths and JSON helpers for the experiment scripts."""
+
+import json
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+import numpy as np
+
+PREFIX_MODEL = "model"
+PREFIX_RESULTS = "results"
+PREFIX_TENSORBOARD = "tensorboard"
+PREFIX_LOGS = "logs"
+
+
+@dataclass(frozen=True)
+class PathHelper:
+    output_dir: Path
+    date_tag: str
+    n_rounds: int
+    batch_size: int
+    n_updates: int
+    learning_rate: float
+    test_ratio: float
+    aggregator: str
+    fedprox_mu: float
+    alpha: float
+    secagg: bool
+    cv: bool
+
+    def get_fpath_out(
+        self,
+        task: str,
+        prefix: str,
+        extension: str = ".json",
+        held_out_node: str | None = None,
+    ) -> Path:
+        stem = "-".join(
+            [
+                prefix,
+                (
+                    self.aggregator
+                    if self.aggregator != "fedprox"
+                    else f"{self.aggregator}_{self.fedprox_mu}mu"
+                ),
+                f"{self.n_rounds}rounds",
+                f"{self.n_updates}updates",
+                f"{self.batch_size}",
+                f"{self.learning_rate}lr",
+                f"{self.test_ratio}test",
+                f"{self.alpha}alpha",
+                "secagg" if self.secagg else "no_secagg",
+                "cv" if self.cv else "no_cv",
+            ]
+        )
+        if held_out_node is not None:
+            stem = f"{stem}-heldout_{held_out_node}"
+        return self.output_dir / self.date_tag / task / f"{stem}{extension}"
+
+
+def serialize(obj: object) -> object:
+    if isinstance(obj, Enum):
+        return obj.name
+    if isinstance(obj, Path):
+        return str(obj)
+    if isinstance(obj, (np.ndarray, np.generic)):
+        return obj.tolist()
+    if isinstance(obj, set):
+        return sorted(obj)
+    raise TypeError(f"Type {type(obj)} not serializable")
+
+
+def save_json(fpath: Path, data: dict) -> None:
+    fpath.write_text(
+        json.dumps(data, indent=4, allow_nan=False, default=serialize) + "\n"
+    )

scripts/_log_utils.py

@@ -0,0 +1,81 @@
+"""Logging helpers for the experiment scripts."""
+
+import logging
+import re
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+LOG_FORMAT_FILE = "%(asctime)s %(name)s %(levelname)-8s - %(message)s"
+LOG_FORMAT_SUMMARY = "%(levelname)-8s - %(message)s"
+_ANSI_ESCAPE = re.compile(r"\x1b\[[0-9;]*m")  # terminal formatting codes
+
+logger = logging.getLogger("brainflex")
+logger.setLevel(logging.INFO)
+fbm_logger = logging.getLogger("fedbiomed")
+
+
+class _ListHandler(logging.Handler):
+    """Keep the formatted log messages in memory."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.messages: list[str] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.messages.append(self.format(record))
+
+
+class _StripAnsiFormatter(logging.Formatter):
+    """Remove the terminal formatting codes."""
+
+    def format(self, record: logging.LogRecord) -> str:
+        return _ANSI_ESCAPE.sub("", super().format(record))
+
+
+def _get_node_log_level(record: logging.LogRecord) -> int | None:
+    """Get the original level of a node log message forwarded to the researcher."""
+    if record.funcName != "print_node_log_message":
+        return None
+    level_name = _ANSI_ESCAPE.sub("", record.getMessage().split("\n", 1)[0]).strip()
+    level = logging.getLevelName(level_name)
+    return level if isinstance(level, int) else None
+
+
+def _filter_log_file_record(record: logging.LogRecord) -> bool:
+    """Keep all script messages but only the Fed-BioMed warnings and errors."""
+    if record.name == fbm_logger.name or record.name.startswith(f"{fbm_logger.name}."):
+        level = _get_node_log_level(record) or record.levelno
+        # security audit events are excluded from Fed-BioMed's console handler too
+        return level >= logging.WARNING and not getattr(record, "is_security", False)
+    return True
+
+
+@contextmanager
+def log_to_file(fpath_log: Path) -> Generator[None, None, None]:
+    """Write the script messages and Fed-BioMed warnings/errors to a file."""
+    fpath_log.parent.mkdir(parents=True, exist_ok=True)
+    handler = logging.FileHandler(fpath_log, mode="w", encoding="utf-8")
+    handler.setFormatter(_StripAnsiFormatter(LOG_FORMAT_FILE))
+    handler.addFilter(_filter_log_file_record)
+    for logger_ in (logger, fbm_logger):
+        logger_.addHandler(handler)
+    try:
+        yield
+    finally:
+        for logger_ in (logger, fbm_logger):
+            logger_.removeHandler(handler)
+        handler.close()
+
+
+@contextmanager
+def print_logger_messages_at_exit() -> Generator[None, None, None]:
+    """Print the script messages."""
+    handler = _ListHandler()
+    handler.setFormatter(logging.Formatter(LOG_FORMAT_SUMMARY))
+    logger.addHandler(handler)
+    try:
+        yield
+    finally:
+        logger.removeHandler(handler)
+        print("\n".join(handler.messages))

scripts/_tasks.py

@@ -0,0 +1,94 @@
+"""Configuration of the experiment tasks."""
+
+import importlib.util
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from fedbiomed.common.training_plans import BaseTrainingPlan
+
+DPATH_TRAINING_PLANS = Path(__file__).resolve().parents[1] / "training_plans"
+
+COMMON_TAG = "brainflex"
+
+
+def _load_training_plan_class(
+    module_name: str, class_name: str
+) -> type[BaseTrainingPlan]:
+    spec = importlib.util.spec_from_file_location(
+        f"training_plans.{module_name}",
+        DPATH_TRAINING_PLANS / f"{module_name}.py",
+    )
+    module = importlib.util.module_from_spec(spec)
+    # registered so that Fed-BioMed can get the source code to send to nodes
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return getattr(module, class_name)
+
+
+@dataclass(frozen=True)
+class TaskConfig:
+    training_plan_class: type[BaseTrainingPlan]
+    target_column: str
+    # whether the model intercept is initialized to the global target mean
+    with_target_mean: bool = False
+    extra_model_args: dict = field(default_factory=dict)
+    # model parameters are clipped to [-x, x] before being encrypted for secure
+    # aggregation
+    secagg_clipping_range: int | None = None
+
+
+TASKS = {
+    "age": TaskConfig(
+        training_plan_class=_load_training_plan_class(
+            "tabular_regressor", "TabularRegressorTrainingPlan"
+        ),
+        target_column="nb:Age",
+        with_target_mean=True,
+        secagg_clipping_range=100,
+    ),
+    "sex": TaskConfig(
+        training_plan_class=_load_training_plan_class(
+            "tabular_classifier", "TabularClassifierTrainingPlan"
+        ),
+        target_column="nb:Sex",
+        extra_model_args={"n_classes": 2},
+    ),
+}
+
+
+def filter_feature_col(col: str) -> bool:
+    if col.startswith("nb:"):
+        return True
+    elif col.endswith("_thickness"):
+        # return False
+        return not (
+            col
+            in [
+                # exclude these
+                "lh_MeanThickness_thickness",
+                "rh_MeanThickness_thickness",
+                "rh_temporalpole_thickness",
+            ]
+        )
+    else:
+        return col in [
+            # include these
+            "EstimatedTotalIntraCranialVol",
+            "Left-Lateral-Ventricle",
+            "Left-Thalamus",
+            "Left-Caudate",
+            "Left-Putamen",
+            "Left-Pallidum",
+            "Left-Hippocampus",
+            "Left-Amygdala",
+            "Left-Accumbens-area",
+            "Right-Lateral-Ventricle",
+            "Right-Thalamus",
+            "Right-Caudate",
+            "Right-Putamen",
+            "Right-Pallidum",
+            "Right-Hippocampus",
+            "Right-Amygdala",
+            "Right-Accumbens-area",
+        ]

scripts/request_training_plan_approval.py

@@ -5,16 +5,14 @@
 import argparse
 import sys
 
+from _tasks import COMMON_TAG, TASKS
 from fedbiomed.common.constants import TrainingPlanApprovalStatus
 from fedbiomed.researcher.federated_workflows import Experiment
 from fedbiomed.researcher.federated_workflows.jobs import (
     TrainingPlanApproveJob,
     TrainingPlanCheckJob,
 )
 
-# importable because the script's directory is the first entry of sys.path
-from run_experiments import COMMON_TAG, TASKS
-
 # Fed-BioMed-defined status of a training plan that a node doesn't know
 TRAINING_PLAN_NOT_REGISTERED = "Not Registered"
 STATUS_SENT = "sent for approval"