From f4a8d005070b7a681877e65c50c5afc858cb2e05 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Fri, 30 Oct 2026 08:39:06 -0400 Subject: [PATCH 1/6] add logging to dedicated output file --- scripts/run_experiments.py | 191 ++++++++++++++++++-------- tests/scripts/test_run_experiments.py | 166 ++++++++++++++++------ 2 files changed, 257 insertions(+), 100 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 3b6f32a..73234a2 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -5,11 +5,13 @@ import argparse import importlib.util import json +import logging import math import shutil import sys -import warnings from collections import defaultdict +from collections.abc import Generator +from contextlib import contextmanager from dataclasses import dataclass, field from datetime import date from enum import Enum @@ -30,6 +32,8 @@ PREFIX_MODEL = "model" PREFIX_RESULTS = "results" PREFIX_TENSORBOARD = "tensorboard" +PREFIX_LOGS = "logs" + COMMON_TAG = "brainflex" AGGREGATORS = ["fedavg", "fedprox", "scaffold"] @@ -44,6 +48,13 @@ DEFAULT_ALPHA = 0.1 DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d") +LOG_FORMAT_FILE = "%(asctime)s %(name)s %(levelname)s - %(message)s" +LOG_FORMAT_SUMMARY = "%(levelname)s - %(message)s" + +logger = logging.getLogger("brainflex.run_experiments") +logger.setLevel(logging.INFO) +fbm_logger = logging.getLogger("fedbiomed") + @dataclass(frozen=True) class TaskConfig: @@ -270,6 +281,57 @@ def _save_json(fpath: Path, data: dict) -> None: ) +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)) + + +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}."): + # keep security audit events even if they are below WARNING level + return record.levelno >= 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(logging.Formatter(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)) + + def _load_training_plan_class(module_name: str, class_name: str) -> type: spec = importlib.util.spec_from_file_location( f"training_plans.{module_name}", @@ -324,10 +386,6 @@ def _get_feature_stats( columns = sorted(columns) - if excluded_columns: - for reason, cols in excluded_columns.items(): - print(f"Excluding {len(cols)} features ({reason}): {cols}") - # cast to float so that the model args are JSON-serializable return ( columns, @@ -473,6 +531,8 @@ def _run_task( """ config = TASKS[task] is_cv = node_data_held_out is not None + if is_cv: + logger.info(f"[{task}] Holding out node {', '.join(node_data_held_out)}") exp = Experiment( tags=[COMMON_TAG, task], nodes=training_nodes, @@ -502,23 +562,24 @@ def _run_task( exp.set_training_data(None, from_tags=True) if not is_cv: - warnings.warn( - "Feature z-scoring statistics are computed on all node data, including" - " the test data (data leakage). Use --cv to avoid this", - stacklevel=2, + logger.warning( + f"[{task}] Feature z-scoring statistics are computed on all node data," + " including the test data (data leakage). Use --cv to avoid this" ) if config.with_target_mean: - warnings.warn( - "The intercept is initialized with the target mean of all node data," - " including the test data (data leakage). Use --cv to avoid this", - stacklevel=2, + logger.warning( + f"[{task}] The intercept is initialized with the target mean of all" + " node data, including the test data (data leakage). Use --cv to" + " avoid this" ) # the training plan needs the feature stats to be instantiated columns, mean, std, excluded_columns, target_mean = _get_feature_stats( exp.analytics, exp.training_data().data(), config.target_column ) - print(f"[{task}] Using {len(columns)} features") + for reason, cols in excluded_columns.items(): + logger.info(f"[{task}] Excluding {len(cols)} features ({reason}): {cols}") + logger.info(f"[{task}] Using {len(columns)} features") model_args = _build_model_args(task, columns, mean, std, alpha, target_mean) exp.set_model_args(model_args) exp.set_training_plan_class( @@ -548,7 +609,7 @@ def _run_task( ) for round_ in range(1, n_rounds + 1): - print(f"[{task}] Round {round_}/{n_rounds}") + logger.info(f"[{task}] Round {round_}/{n_rounds}") if is_cv: _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments) exp.run_once() @@ -563,15 +624,15 @@ def _run_task( _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments) json_data["results"]["node_metrics"] = exp.monitor()._metric_store _save_json(fpath_results, json_data) - print(f"[{task}] Saved settings and results to {fpath_results}") + logger.info(f"[{task}] Saved settings and results to {fpath_results}") if dpath_tensorboard is not None: _move_tensorboard_logs(exp, dpath_tensorboard) - print(f"[{task}] Saved TensorBoard logs to {dpath_tensorboard}") + logger.info(f"[{task}] Saved TensorBoard logs to {dpath_tensorboard}") fpath_model.parent.mkdir(parents=True, exist_ok=True) exp.training_plan().export_model(str(fpath_model)) - print(f"[{task}] Saved global model to {fpath_model}") + logger.info(f"[{task}] Saved global model to {fpath_model}") return exp.training_plan() @@ -644,53 +705,61 @@ def run_experiments( " Use --overwrite to overwrite them" ) - for task, held_out_node in runs: - if held_out_node is None: - # use all nodes for training - training_nodes = None - node_data_held_out = None - else: - print(f"[{task}] Holding out node {held_out_node}") - training_nodes = sorted(set(node_data[task]) - {held_out_node}) - node_data_held_out = {held_out_node: node_data[task][held_out_node]} - - _run_task( - task=task, - n_rounds=n_rounds, - batch_size=batch_size, - n_updates=n_updates, - learning_rate=learning_rate, - test_ratio=test_ratio, - aggregator=aggregator, - fedprox_mu=fedprox_mu, - alpha=alpha, - settings=settings, - fpath_results=path_helper.get_fpath_out( - task=task, - prefix=PREFIX_RESULTS, - extension=".json", - held_out_node=held_out_node, - ), - fpath_model=path_helper.get_fpath_out( + with _print_logger_messages_at_exit(): + for task, held_out_node in runs: + if held_out_node is None: + # use all nodes for training + training_nodes = None + node_data_held_out = None + else: + training_nodes = sorted(set(node_data[task]) - {held_out_node}) + node_data_held_out = {held_out_node: node_data[task][held_out_node]} + + fpath_log = path_helper.get_fpath_out( task=task, - prefix=PREFIX_MODEL, - extension=".pkl", + prefix=PREFIX_LOGS, + extension=".log", held_out_node=held_out_node, - ), - dpath_tensorboard=( - path_helper.get_fpath_out( + ) + with _log_to_file(fpath_log): + _run_task( task=task, - prefix=PREFIX_TENSORBOARD, - extension="", - held_out_node=held_out_node, + n_rounds=n_rounds, + batch_size=batch_size, + n_updates=n_updates, + learning_rate=learning_rate, + test_ratio=test_ratio, + aggregator=aggregator, + fedprox_mu=fedprox_mu, + alpha=alpha, + settings=settings, + fpath_results=path_helper.get_fpath_out( + task=task, + prefix=PREFIX_RESULTS, + extension=".json", + held_out_node=held_out_node, + ), + fpath_model=path_helper.get_fpath_out( + task=task, + prefix=PREFIX_MODEL, + extension=".pkl", + held_out_node=held_out_node, + ), + dpath_tensorboard=( + path_helper.get_fpath_out( + task=task, + prefix=PREFIX_TENSORBOARD, + extension="", + held_out_node=held_out_node, + ) + if tensorboard + else None + ), + training_nodes=training_nodes, + node_data_held_out=node_data_held_out, + secagg=secagg, ) - if tensorboard - else None - ), - training_nodes=training_nodes, - node_data_held_out=node_data_held_out, - secagg=secagg, - ) + logger.info(f"[{task}] Saved log to {fpath_log}") def main(argv: list[str] | None = None) -> None: diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 02fe975..46ee784 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -1,7 +1,7 @@ import importlib.util import json +import logging import sys -import warnings from enum import Enum from pathlib import Path from types import ModuleType @@ -516,28 +516,78 @@ def test_overwrite_node_dir( ] +class TestLogToFile: + def test_log_to_file(self, module: ModuleType, tmp_path: Path): + fpath_log = tmp_path / "logs" / "logs.log" + fbm_logger = logging.getLogger("fedbiomed") + with module._log_to_file(fpath_log): + module.logger.info("script info") + module.logger.warning("script warning") + fbm_logger.info("fbm info") + fbm_logger.warning("fbm warning") + fbm_logger.error("fbm security", extra={"is_security": True}) + lines = fpath_log.read_text().splitlines() + assert len(lines) == 3 + assert lines[0].endswith("brainflex.run_experiments INFO - script info") + assert lines[1].endswith("brainflex.run_experiments WARNING - script warning") + assert lines[2].endswith("fedbiomed WARNING - fbm warning") + + def test_handler_removed_on_error(self, module: ModuleType, tmp_path: Path): + fbm_logger = logging.getLogger("fedbiomed") + handlers_before = (list(module.logger.handlers), list(fbm_logger.handlers)) + with pytest.raises(RuntimeError), module._log_to_file(tmp_path / "logs.log"): + raise RuntimeError + assert (list(module.logger.handlers), list(fbm_logger.handlers)) == ( + handlers_before + ) + + +class TestPrintLoggerMessagesAtExit: + def test_print_logger_messages_at_exit( + self, module: ModuleType, capsys: pytest.CaptureFixture + ): + with module._print_logger_messages_at_exit(): + module.logger.info("[age] first") + logging.getLogger("fedbiomed").warning("not printed") + assert capsys.readouterr().out == "" + module.logger.warning("[age] second") + assert capsys.readouterr().out == ( + "INFO - [age] first\nWARNING - [age] second\n" + ) + assert not any( + isinstance(handler, module._ListHandler) + for handler in module.logger.handlers + ) + + def test_printed_on_error(self, module: ModuleType, capsys: pytest.CaptureFixture): + with pytest.raises(RuntimeError), module._print_logger_messages_at_exit(): + module.logger.info("[age] before error") + raise RuntimeError + assert capsys.readouterr().out == "INFO - [age] before error\n" + + class TestRunTask: - @pytest.mark.filterwarnings("ignore:.*data leakage") def test_run_task( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): settings = {"n_rounds": 3} fpath_results = tmp_path / "age" / "results.json" fpath_model = tmp_path / "age" / "model.pkl" - module._run_task( - "age", - n_rounds=3, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=1.0, - alpha=0.1, - settings=settings, - fpath_results=fpath_results, - fpath_model=fpath_model, - ) + with module._log_to_file(tmp_path / "logs.log"): + module._run_task( + "age", + n_rounds=3, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="scaffold", + fedprox_mu=1.0, + alpha=0.1, + settings=settings, + fpath_results=fpath_results, + fpath_model=fpath_model, + ) assert mock_experiment.call_args.kwargs["tags"] == ["brainflex", "age"] assert mock_experiment.call_args.kwargs["round_limit"] == 3 @@ -564,6 +614,7 @@ def test_run_task( exp.set_training_plan_class.call_args.args[0].__name__ == "TabularRegressorTrainingPlan" ) + assert "[age] Using 3 features" in (tmp_path / "logs.log").read_text() json_data = json.loads(fpath_results.read_text()) assert json_data["settings"]["runtime_args"] == settings @@ -587,7 +638,6 @@ def test_run_task( ) ) - @pytest.mark.filterwarnings("ignore:.*data leakage") def test_settings_saved_if_training_fails( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): @@ -613,7 +663,6 @@ def test_settings_saved_if_training_fails( assert mock_experiment.call_args.kwargs["agg_optimizer"] is None assert "results" not in json_data - @pytest.mark.filterwarnings("ignore:.*data leakage") def test_tensorboard( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): @@ -636,7 +685,6 @@ def test_tensorboard( assert mock_experiment.call_args.kwargs["tensorboard"] is True assert (dpath_tensorboard / "NODE_2").is_dir() - @pytest.mark.filterwarnings("ignore:.*data leakage") def test_no_secagg( self, module: ModuleType, @@ -678,10 +726,11 @@ def test_leakage_warnings( mock_experiment: MagicMock, mock_evaluate_held_out: MagicMock, tmp_path: Path, + caplog: pytest.LogCaptureFixture, task: str, n_warnings: int, ): - with pytest.warns(UserWarning, match="data leakage") as record: + with module._log_to_file(tmp_path / "logs.log"): module._run_task( task, n_rounds=1, @@ -696,7 +745,12 @@ def test_leakage_warnings( fpath_results=tmp_path / "results.json", fpath_model=tmp_path / "model.pkl", ) - assert len(record) == n_warnings + assert [ + record.levelno + for record in caplog.records + if "data leakage" in record.getMessage() + ] == [logging.WARNING] * n_warnings + assert "data leakage" in (tmp_path / "logs.log").read_text() assert mock_experiment.call_args.kwargs["nodes"] is None mock_experiment.return_value.set_training_data.assert_not_called() training_args = mock_experiment.call_args.kwargs["training_args"] @@ -710,27 +764,26 @@ def test_held_out_node( mock_evaluate_held_out: MagicMock, mock_setup_held_out_secagg: MagicMock, tmp_path: Path, + caplog: pytest.LogCaptureFixture, ): fpath_results = tmp_path / "results.json" node_data_held_out = {"NODE_3": NODE_DATA["NODE_3"]} - with warnings.catch_warnings(): - warnings.simplefilter("error") - module._run_task( - "age", - n_rounds=3, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=1.0, - alpha=0.1, - settings={}, - fpath_results=fpath_results, - fpath_model=tmp_path / "model.pkl", - training_nodes=["NODE_1", "NODE_2"], - node_data_held_out=node_data_held_out, - ) + module._run_task( + "age", + n_rounds=3, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="scaffold", + fedprox_mu=1.0, + alpha=0.1, + settings={}, + fpath_results=fpath_results, + fpath_model=tmp_path / "model.pkl", + training_nodes=["NODE_1", "NODE_2"], + node_data_held_out=node_data_held_out, + ) assert mock_experiment.call_args.kwargs["nodes"] == ["NODE_1", "NODE_2"] mock_experiment.return_value.set_training_data.assert_called_once_with( None, from_tags=True @@ -751,6 +804,8 @@ def test_held_out_node( assert json_data["settings"]["training_nodes"] == ["NODE_1", "NODE_2"] assert json_data["settings"]["held_out_node"] == node_data_held_out assert "node_metrics" in json_data["results"] + assert not any("data leakage" in r.getMessage() for r in caplog.records) + assert "[age] Holding out node NODE_3" in caplog.messages def _write_results( @@ -804,6 +859,13 @@ def test_tasks_in_order( assert kwargs["fpath_results"].name.startswith("results-scaffold-2rounds-") assert kwargs["fpath_model"].name.startswith("model-scaffold-2rounds-") assert kwargs["fpath_model"].suffix == ".pkl" + # created by _log_to_file even though _run_task is mocked + fpath_log = ( + dpath_date + / "age" + / kwargs["fpath_results"].name.replace("results-", "logs-") + ) + assert fpath_log.with_suffix(".log").exists() assert kwargs["dpath_tensorboard"].parent == dpath_date / "age" assert kwargs["dpath_tensorboard"].name.startswith( "tensorboard-scaffold-2rounds-" @@ -819,6 +881,28 @@ def test_tasks_in_order( assert settings["secagg"] is True assert kwargs["secagg"] is True + def test_messages_printed_at_end( + self, + module: ModuleType, + mock_run_task: MagicMock, + tmp_path: Path, + capsys: pytest.CaptureFixture, + ): + def run_task(task: str, **kwargs): + module.logger.info(f"[{task}] done") + assert capsys.readouterr().out == "" + if task == "age": + raise RuntimeError + + mock_run_task.side_effect = run_task + with pytest.raises(RuntimeError): + module.main([str(tmp_path), "--task", "sex", "--task", "age", "--no-cv"]) + lines = capsys.readouterr().out.splitlines() + assert lines[0] == "INFO - [sex] done" + assert lines[1].startswith("INFO - [sex] Saved log to ") + # no "Saved log" message for the failed run + assert lines[2:] == ["INFO - [age] done"] + def test_date_tag( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): @@ -915,6 +999,10 @@ def test_cv( assert kwargs["node_data_held_out"] == {"NODE_1": NODE_DATA["NODE_1"]} assert kwargs["fpath_results"].stem.endswith("-heldout_NODE_1") assert kwargs["fpath_model"].stem.endswith("-heldout_NODE_1") + fpath_log = kwargs["fpath_results"].with_name( + kwargs["fpath_results"].name.replace("results-", "logs-") + ) + assert fpath_log.with_suffix(".log").exists() assert kwargs["dpath_tensorboard"].name.endswith("-heldout_NODE_1") assert kwargs["settings"]["cv"] is True From b35128c97c697368f80deb7e3b37523ba30694b9 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Fri, 30 Oct 2026 10:01:30 -0400 Subject: [PATCH 2/6] set default tasks and include cv/secagg args in output file paths --- scripts/run_experiments.py | 43 +++++++++++++++++++++++++------------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 73234a2..36f1a11 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -38,6 +38,7 @@ AGGREGATORS = ["fedavg", "fedprox", "scaffold"] +DEFAULT_TASKS = ["age", "sex"] DEFAULT_N_ROUNDS = 30 DEFAULT_BATCH_SIZE = 1000 # full batch DEFAULT_N_UPDATES = 50 @@ -46,12 +47,15 @@ DEFAULT_AGGREGATOR = "scaffold" DEFAULT_FEDPROX_MU = 1.0 DEFAULT_ALPHA = 0.1 +DEFAULT_TENSORBOARD = True +DEFAULT_SECAGG = True +DEFAULT_CV = True DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d") LOG_FORMAT_FILE = "%(asctime)s %(name)s %(levelname)s - %(message)s" LOG_FORMAT_SUMMARY = "%(levelname)s - %(message)s" -logger = logging.getLogger("brainflex.run_experiments") +logger = logging.getLogger("brainflex") logger.setLevel(logging.INFO) fbm_logger = logging.getLogger("fedbiomed") @@ -81,6 +85,8 @@ class PathHelper: test_ratio: float aggregator: str alpha: float + secagg: bool + cv: bool def get_fpath_out( self, @@ -92,13 +98,19 @@ def get_fpath_out( stem = "-".join( [ prefix, - self.aggregator, + ( + self.aggregator + if self.aggregator != "fedprox" + else f"{self.aggregator}_{DEFAULT_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: @@ -142,7 +154,6 @@ def _get_parser() -> argparse.ArgumentParser: "--task", dest="tasks", action="append", - required=True, choices=list(TASKS), help="Task to run. Can be repeated.", ) @@ -194,35 +205,35 @@ def _get_parser() -> argparse.ArgumentParser: default=DEFAULT_ALPHA, help=f"L2 regularization strength. Default: {DEFAULT_ALPHA}.", ) - parser.add_argument( - "--overwrite", - action="store_true", - help="Overwrite existing results.", - ) parser.add_argument( "--tensorboard", action=argparse.BooleanOptionalAction, - default=True, - help="Save node metrics as TensorBoard logs. Default: True.", + default=DEFAULT_TENSORBOARD, + help=f"Save node metrics as TensorBoard logs. Default: {DEFAULT_TENSORBOARD}.", ) parser.add_argument( "--secagg", action=argparse.BooleanOptionalAction, - default=True, + default=DEFAULT_SECAGG, help=( "Secure aggregation of the model parameters and the feature statistics." - " Needs at least 2 training nodes. Default: True." + f" Needs at least 2 training nodes. Default: {DEFAULT_SECAGG}." ), ) parser.add_argument( "--cv", action=argparse.BooleanOptionalAction, - default=True, + default=DEFAULT_CV, help=( "Leave-one-node-out cross-validation: train on all nodes but one and" - " evaluate on the held-out node. Default: True." + f" evaluate on the held-out node. Default: {DEFAULT_CV}." ), ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite existing results.", + ) return parser @@ -685,6 +696,8 @@ def run_experiments( test_ratio=test_ratio, aggregator=aggregator, alpha=alpha, + secagg=secagg, + cv=cv, ) # error out if one of the result files already exists @@ -767,7 +780,7 @@ def main(argv: list[str] | None = None) -> None: try: run_experiments( args.output_dir, - args.tasks, + args.tasks if args.tasks is not None else DEFAULT_TASKS, date_tag=args.date_tag, n_rounds=args.n_rounds, batch_size=args.batch_size, From f513bc3e6eef736e75ddec09167d7fbcd838ff62 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Fri, 30 Oct 2026 10:06:45 -0400 Subject: [PATCH 3/6] fix tests and bug --- scripts/run_experiments.py | 4 +- tests/scripts/test_run_experiments.py | 54 ++++++++++++++++++++------- 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 36f1a11..bb92730 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -84,6 +84,7 @@ class PathHelper: learning_rate: float test_ratio: float aggregator: str + fedprox_mu: float alpha: float secagg: bool cv: bool @@ -101,7 +102,7 @@ def get_fpath_out( ( self.aggregator if self.aggregator != "fedprox" - else f"{self.aggregator}_{DEFAULT_FEDPROX_MU}mu" + else f"{self.aggregator}_{self.fedprox_mu}mu" ), f"{self.n_rounds}rounds", f"{self.n_updates}updates", @@ -695,6 +696,7 @@ def run_experiments( learning_rate=learning_rate, test_ratio=test_ratio, aggregator=aggregator, + fedprox_mu=fedprox_mu, alpha=alpha, secagg=secagg, cv=cv, diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 46ee784..a9ad2d5 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -1,3 +1,4 @@ +import dataclasses import importlib.util import json import logging @@ -162,7 +163,10 @@ def path_helper(module: ModuleType, tmp_path: Path) -> object: learning_rate=0.001, test_ratio=0.2, aggregator="scaffold", + fedprox_mu=2.0, alpha=0.1, + secagg=True, + cv=True, ) @@ -172,9 +176,22 @@ def test_get_fpath_out(self, path_helper: object, tmp_path: Path): tmp_path / "2026_01_01" / "age" - / "results-scaffold-10rounds-50updates-32-0.001lr-0.2test-0.1alpha.json" + / ( + "results-scaffold-10rounds-50updates-32-0.001lr-0.2test-0.1alpha" + "-secagg-cv.json" + ) ) + def test_no_secagg_no_cv(self, path_helper: object): + path_helper = dataclasses.replace(path_helper, secagg=False, cv=False) + fpath = path_helper.get_fpath_out("age", "results") + assert fpath.stem.endswith("-0.1alpha-no_secagg-no_cv") + + def test_fedprox(self, path_helper: object): + path_helper = dataclasses.replace(path_helper, aggregator="fedprox") + fpath = path_helper.get_fpath_out("age", "results") + assert fpath.name.startswith("results-fedprox_2.0mu-") + def test_extension(self, path_helper: object, tmp_path: Path): fpath = path_helper.get_fpath_out("sex", "model", ".pkl") assert fpath.parent == tmp_path / "2026_01_01" / "sex" @@ -184,7 +201,7 @@ def test_held_out_node(self, path_helper: object): fpath = path_helper.get_fpath_out( "age", "model", ".pkl", held_out_node="NODE_1" ) - assert fpath.name.endswith("-0.1alpha-heldout_NODE_1.pkl") + assert fpath.name.endswith("-0.1alpha-secagg-cv-heldout_NODE_1.pkl") class TestGetNodeData: @@ -528,8 +545,8 @@ def test_log_to_file(self, module: ModuleType, tmp_path: Path): fbm_logger.error("fbm security", extra={"is_security": True}) lines = fpath_log.read_text().splitlines() assert len(lines) == 3 - assert lines[0].endswith("brainflex.run_experiments INFO - script info") - assert lines[1].endswith("brainflex.run_experiments WARNING - script warning") + assert lines[0].endswith("brainflex INFO - script info") + assert lines[1].endswith("brainflex WARNING - script warning") assert lines[2].endswith("fedbiomed WARNING - fbm warning") def test_handler_removed_on_error(self, module: ModuleType, tmp_path: Path): @@ -816,7 +833,10 @@ def _write_results( held_out_node: str | None = None, date_tag: str | None = None, ) -> None: - """Write a results file for the default CLI arguments.""" + """Write a results file for the default CLI arguments. + + CV is assumed to be on if and only if ``held_out_node`` is given. + """ fpath = module.PathHelper( output_dir=dpath, date_tag=date_tag if date_tag is not None else module.DEFAULT_DATE_TAG, @@ -826,7 +846,10 @@ def _write_results( learning_rate=module.DEFAULT_LEARNING_RATE, test_ratio=module.DEFAULT_TEST_RATIO, aggregator=module.DEFAULT_AGGREGATOR, + fedprox_mu=module.DEFAULT_FEDPROX_MU, alpha=module.DEFAULT_ALPHA, + secagg=module.DEFAULT_SECAGG, + cv=held_out_node is not None, ).get_fpath_out(task, module.PREFIX_RESULTS, held_out_node=held_out_node) fpath.parent.mkdir(parents=True) fpath.write_text(json.dumps(data)) @@ -939,8 +962,8 @@ def test_aggregator_options( assert kwargs["aggregator"] == "fedprox" assert kwargs["fedprox_mu"] == 5.0 assert kwargs["alpha"] == 0.5 - assert kwargs["fpath_results"].name.startswith("results-fedprox-") - assert kwargs["fpath_results"].stem.endswith("-0.5alpha") + assert kwargs["fpath_results"].name.startswith("results-fedprox_5.0mu-") + assert "-0.5alpha-" in kwargs["fpath_results"].stem def test_no_tensorboard( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path @@ -1066,14 +1089,17 @@ def test_cv_existing_results( module.main([str(tmp_path), "--task", "age"]) mock_run_task.assert_not_called() - @pytest.mark.parametrize("args", [[], ["--task", "height"]]) + def test_default_tasks( + self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path + ): + module.main([str(tmp_path), "--no-cv"]) + assert [ + call.kwargs["task"] for call in mock_run_task.call_args_list + ] == module.DEFAULT_TASKS + def test_invalid_task( - self, - module: ModuleType, - mock_run_task: MagicMock, - tmp_path: Path, - args: list[str], + self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): with pytest.raises(SystemExit): - module.main([str(tmp_path), *args]) + module.main([str(tmp_path), "--task", "height"]) mock_run_task.assert_not_called() From c2b19094f50da2b95973edb1a99ec60e0fa548aa Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Fri, 30 Oct 2026 14:35:23 -0400 Subject: [PATCH 4/6] merge log and results files --- scripts/run_experiments.py | 167 +++++++------ tests/scripts/test_run_experiments.py | 339 +++++++++++++++----------- 2 files changed, 284 insertions(+), 222 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index bb92730..701a820 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -22,7 +22,6 @@ from fedbiomed.common.metrics import MetricTypes from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_args import TrainingArgs -from fedbiomed.common.training_plans import BaseTrainingPlan from fedbiomed.researcher.datasets import FederatedDataset from fedbiomed.researcher.federated_workflows import Experiment, FederatedAnalytics from fedbiomed.researcher.federated_workflows.jobs import TrainingJob @@ -344,6 +343,12 @@ def _print_logger_messages_at_exit() -> Generator[None, None, None]: print("\n".join(handler.messages)) +def _get_log_prefix(task: str, held_out_node: str | None = None) -> str: + if held_out_node is None: + return f"[{task}]" + return f"[{task}|held_out={held_out_node}]" + + def _load_training_plan_class(module_name: str, class_name: str) -> type: spec = importlib.util.spec_from_file_location( f"training_plans.{module_name}", @@ -528,23 +533,26 @@ def _run_task( aggregator: str, fedprox_mu: float, alpha: float, - settings: dict, - fpath_results: Path, fpath_model: Path, dpath_tensorboard: Path | None = None, training_nodes: list[str] | None = None, node_data_held_out: dict[str, dict] | None = None, secagg: bool = True, -) -> BaseTrainingPlan: +) -> Generator[tuple[str, dict], None, None]: """Train a model for a task and save it in a task-specific output directory. If ``node_data_held_out`` is given, only ``training_nodes`` are used for training and the held-out node is used for testing. + + Yields the fold key (the training nodes joined by ``__``) and the fold settings + and results. Saves the model and TensorBoard logs after the last yield. """ config = TASKS[task] is_cv = node_data_held_out is not None + held_out_node = next(iter(node_data_held_out)) if is_cv else None + prefix = _get_log_prefix(task, held_out_node) if is_cv: - logger.info(f"[{task}] Holding out node {', '.join(node_data_held_out)}") + logger.info(f"{prefix} Holding out node {held_out_node}") exp = Experiment( tags=[COMMON_TAG, task], nodes=training_nodes, @@ -572,15 +580,16 @@ def _run_task( # redo the search to drop the held-out node # otherwise, it would be included in the federated analytics exp.set_training_data(None, from_tags=True) + fold_key = "__".join(sorted(exp.training_data().data())) if not is_cv: logger.warning( - f"[{task}] Feature z-scoring statistics are computed on all node data," + f"{prefix} Feature z-scoring statistics are computed on all node data," " including the test data (data leakage). Use --cv to avoid this" ) if config.with_target_mean: logger.warning( - f"[{task}] The intercept is initialized with the target mean of all" + f"{prefix} The intercept is initialized with the target mean of all" " node data, including the test data (data leakage). Use --cv to" " avoid this" ) @@ -590,8 +599,8 @@ def _run_task( exp.analytics, exp.training_data().data(), config.target_column ) for reason, cols in excluded_columns.items(): - logger.info(f"[{task}] Excluding {len(cols)} features ({reason}): {cols}") - logger.info(f"[{task}] Using {len(columns)} features") + logger.info(f"{prefix} Excluding {len(cols)} features ({reason}): {cols}") + logger.info(f"{prefix} Using {len(columns)} features") model_args = _build_model_args(task, columns, mean, std, alpha, target_mean) exp.set_model_args(model_args) exp.set_training_plan_class( @@ -601,9 +610,8 @@ def _run_task( ) # saved before training so that they are available even if training fails - json_data = { + fold = { "settings": { - "runtime_args": settings, "training_nodes": training_nodes, "held_out_node": node_data_held_out, "excluded_columns": excluded_columns, @@ -611,8 +619,7 @@ def _run_task( "training_args": exp.training_args(), } } - fpath_results.parent.mkdir(parents=True, exist_ok=True) - _save_json(fpath_results, json_data) + yield fold_key, fold held_out_secagg_arguments = None if is_cv and secagg: @@ -621,31 +628,29 @@ def _run_task( ) for round_ in range(1, n_rounds + 1): - logger.info(f"[{task}] Round {round_}/{n_rounds}") + logger.info(f"{prefix} Round {round_}/{n_rounds}") if is_cv: _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments) exp.run_once() - json_data["results"] = { + fold["results"] = { # node training losses and test metrics (no public getter) "node_metrics": exp.monitor()._metric_store, "training_replies": exp.save_training_replies(), } - _save_json(fpath_results, json_data) + yield fold_key, fold if is_cv: # final global model _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments) - json_data["results"]["node_metrics"] = exp.monitor()._metric_store - _save_json(fpath_results, json_data) - logger.info(f"[{task}] Saved settings and results to {fpath_results}") + fold["results"]["node_metrics"] = exp.monitor()._metric_store + yield fold_key, fold if dpath_tensorboard is not None: _move_tensorboard_logs(exp, dpath_tensorboard) - logger.info(f"[{task}] Saved TensorBoard logs to {dpath_tensorboard}") + logger.info(f"{prefix} Saved TensorBoard logs to {dpath_tensorboard}") fpath_model.parent.mkdir(parents=True, exist_ok=True) exp.training_plan().export_model(str(fpath_model)) - logger.info(f"[{task}] Saved global model to {fpath_model}") - return exp.training_plan() + logger.info(f"{prefix} Saved global model to {fpath_model}") def run_experiments( @@ -669,12 +674,12 @@ def run_experiments( output_dir = output_dir.resolve() settings = locals().copy() - # (task, held-out node) pairs, with None for no held-out node - runs: list[tuple[str, str | None]] = [] + # held-out nodes for each task, with None for no held-out node + runs: dict[str, list[str | None]] = {} node_data: dict[str, dict[str, dict]] = {} for task in tasks: if not cv: - runs.append((task, None)) + runs[task] = [None] continue node_data[task] = _get_node_data(task) # secure aggregation needs at least 2 training nodes @@ -685,7 +690,7 @@ def run_experiments( f"{' with secure aggregation' if secagg else ''}," f" but {task} has {len(node_data[task])}" ) - runs.extend((task, node) for node in sorted(node_data[task])) + runs[task] = sorted(node_data[task]) path_helper = PathHelper( output_dir=output_dir, @@ -704,14 +709,14 @@ def run_experiments( # error out if one of the result files already exists fpaths_results_existing = [] - for task, held_out_node in runs: + for task in runs: fpath = path_helper.get_fpath_out( - task=task, - prefix=PREFIX_RESULTS, - extension=".json", - held_out_node=held_out_node, + task=task, prefix=PREFIX_RESULTS, extension=".json" ) - if fpath.exists() and "results" in json.loads(fpath.read_text()): + if fpath.exists() and any( + "results" in fold + for fold in json.loads(fpath.read_text()).get("folds", {}).values() + ): fpaths_results_existing.append(fpath) if fpaths_results_existing and not overwrite: @@ -721,59 +726,63 @@ def run_experiments( ) with _print_logger_messages_at_exit(): - for task, held_out_node in runs: - if held_out_node is None: - # use all nodes for training - training_nodes = None - node_data_held_out = None - else: - training_nodes = sorted(set(node_data[task]) - {held_out_node}) - node_data_held_out = {held_out_node: node_data[task][held_out_node]} - + for task, held_out_nodes in runs.items(): + fpath_results = path_helper.get_fpath_out( + task=task, prefix=PREFIX_RESULTS, extension=".json" + ) fpath_log = path_helper.get_fpath_out( - task=task, - prefix=PREFIX_LOGS, - extension=".log", - held_out_node=held_out_node, + task=task, prefix=PREFIX_LOGS, extension=".log" ) + json_data = {"settings": {"runtime_args": settings}, "folds": {}} + fpath_results.parent.mkdir(parents=True, exist_ok=True) with _log_to_file(fpath_log): - _run_task( - task=task, - n_rounds=n_rounds, - batch_size=batch_size, - n_updates=n_updates, - learning_rate=learning_rate, - test_ratio=test_ratio, - aggregator=aggregator, - fedprox_mu=fedprox_mu, - alpha=alpha, - settings=settings, - fpath_results=path_helper.get_fpath_out( - task=task, - prefix=PREFIX_RESULTS, - extension=".json", - held_out_node=held_out_node, - ), - fpath_model=path_helper.get_fpath_out( + for held_out_node in held_out_nodes: + if held_out_node is None: + # use all nodes for training + training_nodes = None + node_data_held_out = None + else: + training_nodes = sorted(set(node_data[task]) - {held_out_node}) + node_data_held_out = { + held_out_node: node_data[task][held_out_node] + } + + for fold_key, fold in _run_task( task=task, - prefix=PREFIX_MODEL, - extension=".pkl", - held_out_node=held_out_node, - ), - dpath_tensorboard=( - path_helper.get_fpath_out( + n_rounds=n_rounds, + batch_size=batch_size, + n_updates=n_updates, + learning_rate=learning_rate, + test_ratio=test_ratio, + aggregator=aggregator, + fedprox_mu=fedprox_mu, + alpha=alpha, + fpath_model=path_helper.get_fpath_out( task=task, - prefix=PREFIX_TENSORBOARD, - extension="", + prefix=PREFIX_MODEL, + extension=".pkl", held_out_node=held_out_node, - ) - if tensorboard - else None - ), - training_nodes=training_nodes, - node_data_held_out=node_data_held_out, - secagg=secagg, - ) + ), + dpath_tensorboard=( + path_helper.get_fpath_out( + task=task, + prefix=PREFIX_TENSORBOARD, + extension="", + held_out_node=held_out_node, + ) + if tensorboard + else None + ), + training_nodes=training_nodes, + node_data_held_out=node_data_held_out, + secagg=secagg, + ): + json_data["folds"][fold_key] = fold + _save_json(fpath_results, json_data) + logger.info( + f"{_get_log_prefix(task, held_out_node)} Saved settings and" + f" results to {fpath_results}" + ) logger.info(f"[{task}] Saved log to {fpath_log}") diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index a9ad2d5..3934277 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -3,6 +3,7 @@ import json import logging import sys +from collections.abc import Generator from enum import Enum from pathlib import Path from types import ModuleType @@ -105,9 +106,20 @@ def mock_experiment( return mock_experiment +def _fake_run_task( + training_nodes: list[str] | None = None, **kwargs +) -> Generator[tuple[str, dict], None, None]: + """Yield a fold like _run_task, without training.""" + fold_key = "__".join(training_nodes or NODE_DATA) + fold = {"settings": {"training_nodes": training_nodes}} + yield fold_key, fold + fold["results"] = {"node_metrics": {}} + yield fold_key, fold + + @pytest.fixture def mock_run_task(module: ModuleType, mocker: MockerFixture) -> MagicMock: - return mocker.patch.object(module, "_run_task") + return mocker.patch.object(module, "_run_task", side_effect=_fake_run_task) @pytest.fixture @@ -587,23 +599,21 @@ class TestRunTask: def test_run_task( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): - settings = {"n_rounds": 3} - fpath_results = tmp_path / "age" / "results.json" fpath_model = tmp_path / "age" / "model.pkl" with module._log_to_file(tmp_path / "logs.log"): - module._run_task( - "age", - n_rounds=3, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=1.0, - alpha=0.1, - settings=settings, - fpath_results=fpath_results, - fpath_model=fpath_model, + folds = list( + module._run_task( + "age", + n_rounds=3, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="scaffold", + fedprox_mu=1.0, + alpha=0.1, + fpath_model=fpath_model, + ) ) assert mock_experiment.call_args.kwargs["tags"] == ["brainflex", "age"] @@ -633,10 +643,12 @@ def test_run_task( ) assert "[age] Using 3 features" in (tmp_path / "logs.log").read_text() - json_data = json.loads(fpath_results.read_text()) - assert json_data["settings"]["runtime_args"] == settings - assert json_data["settings"]["excluded_columns"] == {} - model_args = json_data["settings"]["model_args"] + # before training and after each round + assert len(folds) == 4 + assert {fold_key for fold_key, _ in folds} == {"NODE_1__NODE_2__NODE_3"} + fold = json.loads(json.dumps(folds[-1][1], default=module._serialize)) + assert fold["settings"]["excluded_columns"] == {} + model_args = fold["settings"]["model_args"] assert model_args["target_column"] == "nb:Age" assert model_args["alpha"] == 0.1 assert model_args["target_mean"] == 30.0 @@ -646,10 +658,10 @@ def test_run_task( "nb:Sex", ] exp.set_model_args.assert_called_once_with(model_args) - training_args = json_data["settings"]["training_args"] + training_args = fold["settings"]["training_args"] assert training_args["test_metric"] == "MEAN_SQUARE_ERROR" # JSON object keys are strings - assert json_data["results"] == json.loads( + assert fold["results"] == json.loads( json.dumps( {"node_metrics": NODE_METRICS, "training_replies": TRAINING_REPLIES} ) @@ -659,45 +671,45 @@ def test_settings_saved_if_training_fails( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): mock_experiment.return_value.run_once.side_effect = RuntimeError - fpath_results = tmp_path / "sex" / "results.json" + folds = module._run_task( + "sex", + n_rounds=3, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="fedavg", + fedprox_mu=1.0, + alpha=0.1, + fpath_model=tmp_path / "sex" / "model.pkl", + ) + # yielded before training so that the settings can be saved + fold_key, fold = next(folds) + assert fold_key == "NODE_1__NODE_2__NODE_3" + assert fold["settings"]["model_args"]["n_classes"] == 2 + assert mock_experiment.call_args.kwargs["agg_optimizer"] is None with pytest.raises(RuntimeError): + next(folds) + assert "results" not in fold + + def test_tensorboard( + self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path + ): + dpath_tensorboard = tmp_path / "age" / "tensorboard" + list( module._run_task( - "sex", - n_rounds=3, + "age", + n_rounds=1, batch_size=8, n_updates=10, learning_rate=0.1, test_ratio=0.2, - aggregator="fedavg", + aggregator="scaffold", fedprox_mu=1.0, alpha=0.1, - settings={}, - fpath_results=fpath_results, - fpath_model=tmp_path / "sex" / "model.pkl", + fpath_model=tmp_path / "age" / "model.pkl", + dpath_tensorboard=dpath_tensorboard, ) - json_data = json.loads(fpath_results.read_text()) - assert json_data["settings"]["model_args"]["n_classes"] == 2 - assert mock_experiment.call_args.kwargs["agg_optimizer"] is None - assert "results" not in json_data - - def test_tensorboard( - self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path - ): - dpath_tensorboard = tmp_path / "age" / "tensorboard" - module._run_task( - "age", - n_rounds=1, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=1.0, - alpha=0.1, - settings={}, - fpath_results=tmp_path / "age" / "results.json", - fpath_model=tmp_path / "age" / "model.pkl", - dpath_tensorboard=dpath_tensorboard, ) assert mock_experiment.call_args.kwargs["tensorboard"] is True assert (dpath_tensorboard / "NODE_2").is_dir() @@ -710,22 +722,22 @@ def test_no_secagg( mock_setup_held_out_secagg: MagicMock, tmp_path: Path, ): - module._run_task( - "sex", - n_rounds=1, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="fedavg", - fedprox_mu=1.0, - alpha=0.1, - settings={}, - fpath_results=tmp_path / "results.json", - fpath_model=tmp_path / "model.pkl", - training_nodes=["NODE_1", "NODE_2"], - node_data_held_out={"NODE_3": NODE_DATA["NODE_3"]}, - secagg=False, + list( + module._run_task( + "sex", + n_rounds=1, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="fedavg", + fedprox_mu=1.0, + alpha=0.1, + fpath_model=tmp_path / "model.pkl", + training_nodes=["NODE_1", "NODE_2"], + node_data_held_out={"NODE_3": NODE_DATA["NODE_3"]}, + secagg=False, + ) ) secagg = mock_experiment.call_args.kwargs["secagg"] assert secagg.active is False @@ -748,19 +760,19 @@ def test_leakage_warnings( n_warnings: int, ): with module._log_to_file(tmp_path / "logs.log"): - module._run_task( - task, - n_rounds=1, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="fedavg", - fedprox_mu=1.0, - alpha=0.1, - settings={}, - fpath_results=tmp_path / "results.json", - fpath_model=tmp_path / "model.pkl", + list( + module._run_task( + task, + n_rounds=1, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="fedavg", + fedprox_mu=1.0, + alpha=0.1, + fpath_model=tmp_path / "model.pkl", + ) ) assert [ record.levelno @@ -783,23 +795,27 @@ def test_held_out_node( tmp_path: Path, caplog: pytest.LogCaptureFixture, ): - fpath_results = tmp_path / "results.json" node_data_held_out = {"NODE_3": NODE_DATA["NODE_3"]} - module._run_task( - "age", - n_rounds=3, - batch_size=8, - n_updates=10, - learning_rate=0.1, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=1.0, - alpha=0.1, - settings={}, - fpath_results=fpath_results, - fpath_model=tmp_path / "model.pkl", - training_nodes=["NODE_1", "NODE_2"], - node_data_held_out=node_data_held_out, + # the held-out node is dropped by the new search + exp = mock_experiment.return_value + exp.training_data.return_value.data.return_value = { + node: NODE_DATA[node] for node in ["NODE_1", "NODE_2"] + } + folds = list( + module._run_task( + "age", + n_rounds=3, + batch_size=8, + n_updates=10, + learning_rate=0.1, + test_ratio=0.2, + aggregator="scaffold", + fedprox_mu=1.0, + alpha=0.1, + fpath_model=tmp_path / "model.pkl", + training_nodes=["NODE_1", "NODE_2"], + node_data_held_out=node_data_held_out, + ) ) assert mock_experiment.call_args.kwargs["nodes"] == ["NODE_1", "NODE_2"] mock_experiment.return_value.set_training_data.assert_called_once_with( @@ -817,12 +833,28 @@ def test_held_out_node( mock_evaluate_held_out.assert_called_with( exp, node_data_held_out, mock_setup_held_out_secagg.return_value ) - json_data = json.loads(fpath_results.read_text()) - assert json_data["settings"]["training_nodes"] == ["NODE_1", "NODE_2"] - assert json_data["settings"]["held_out_node"] == node_data_held_out - assert "node_metrics" in json_data["results"] + # before training, after each round, and after the final evaluation + assert len(folds) == 5 + fold_key, fold = folds[-1] + assert fold_key == "NODE_1__NODE_2" + assert fold["settings"]["training_nodes"] == ["NODE_1", "NODE_2"] + assert fold["settings"]["held_out_node"] == node_data_held_out + assert "node_metrics" in fold["results"] assert not any("data leakage" in r.getMessage() for r in caplog.records) - assert "[age] Holding out node NODE_3" in caplog.messages + assert "[age|held_out=NODE_3] Holding out node NODE_3" in caplog.messages + assert "[age|held_out=NODE_3] Round 1/3" in caplog.messages + + +def _read_results(dpath_task: Path) -> tuple[Path, dict]: + """Read the only results file of a task.""" + (fpath,) = dpath_task.glob("results-*.json") + return fpath, json.loads(fpath.read_text()) + + +RESULTS_WITH_FOLD = { + "settings": {"runtime_args": {}}, + "folds": {"NODE_1__NODE_2": {"settings": {}, "results": {}}}, +} def _write_results( @@ -830,13 +862,10 @@ def _write_results( dpath: Path, task: str, data: dict, - held_out_node: str | None = None, + cv: bool = False, date_tag: str | None = None, ) -> None: - """Write a results file for the default CLI arguments. - - CV is assumed to be on if and only if ``held_out_node`` is given. - """ + """Write a results file for the default CLI arguments.""" fpath = module.PathHelper( output_dir=dpath, date_tag=date_tag if date_tag is not None else module.DEFAULT_DATE_TAG, @@ -849,8 +878,8 @@ def _write_results( fedprox_mu=module.DEFAULT_FEDPROX_MU, alpha=module.DEFAULT_ALPHA, secagg=module.DEFAULT_SECAGG, - cv=held_out_node is not None, - ).get_fpath_out(task, module.PREFIX_RESULTS, held_out_node=held_out_node) + cv=cv, + ).get_fpath_out(task, module.PREFIX_RESULTS) fpath.parent.mkdir(parents=True) fpath.write_text(json.dumps(data)) @@ -878,23 +907,25 @@ def test_tasks_in_order( kwargs = mock_run_task.call_args.kwargs assert kwargs["n_rounds"] == 2 dpath_date = tmp_path.resolve() / module.DEFAULT_DATE_TAG - assert kwargs["fpath_results"].parent == dpath_date / "age" - assert kwargs["fpath_results"].name.startswith("results-scaffold-2rounds-") + fpath_results, json_data = _read_results(dpath_date / "age") + assert fpath_results.name.startswith("results-scaffold-2rounds-") assert kwargs["fpath_model"].name.startswith("model-scaffold-2rounds-") assert kwargs["fpath_model"].suffix == ".pkl" # created by _log_to_file even though _run_task is mocked - fpath_log = ( - dpath_date - / "age" - / kwargs["fpath_results"].name.replace("results-", "logs-") - ) + fpath_log = dpath_date / "age" / fpath_results.name.replace("results-", "logs-") assert fpath_log.with_suffix(".log").exists() assert kwargs["dpath_tensorboard"].parent == dpath_date / "age" assert kwargs["dpath_tensorboard"].name.startswith( "tensorboard-scaffold-2rounds-" ) - settings = kwargs["settings"] - assert settings["output_dir"] == tmp_path.resolve() + assert json_data["folds"] == { + "NODE_1__NODE_2__NODE_3": { + "settings": {"training_nodes": None}, + "results": {"node_metrics": {}}, + } + } + settings = json_data["settings"]["runtime_args"] + assert settings["output_dir"] == str(tmp_path.resolve()) assert settings["tasks"] == ["sex", "age"] assert settings["date_tag"] == module.DEFAULT_DATE_TAG assert settings["n_rounds"] == 2 @@ -911,35 +942,39 @@ def test_messages_printed_at_end( tmp_path: Path, capsys: pytest.CaptureFixture, ): - def run_task(task: str, **kwargs): + def run_task(task: str, **kwargs) -> Generator[tuple[str, dict], None, None]: module.logger.info(f"[{task}] done") assert capsys.readouterr().out == "" if task == "age": raise RuntimeError + yield from _fake_run_task(**kwargs) mock_run_task.side_effect = run_task with pytest.raises(RuntimeError): module.main([str(tmp_path), "--task", "sex", "--task", "age", "--no-cv"]) lines = capsys.readouterr().out.splitlines() assert lines[0] == "INFO - [sex] done" - assert lines[1].startswith("INFO - [sex] Saved log to ") - # no "Saved log" message for the failed run - assert lines[2:] == ["INFO - [age] done"] + assert lines[1].startswith("INFO - [sex] Saved settings and results to ") + assert lines[2].startswith("INFO - [sex] Saved log to ") + # no "Saved" messages for the failed run + assert lines[3:] == ["INFO - [age] done"] def test_date_tag( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): module.main([str(tmp_path), "--task", "age", "--date-tag", "mytag", "--no-cv"]) - kwargs = mock_run_task.call_args.kwargs - assert kwargs["fpath_results"].parent == tmp_path.resolve() / "mytag" / "age" - assert kwargs["settings"]["date_tag"] == "mytag" + _, json_data = _read_results(tmp_path.resolve() / "mytag" / "age") + assert json_data["settings"]["runtime_args"]["date_tag"] == "mytag" def test_no_secagg( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): module.main([str(tmp_path), "--task", "age", "--no-secagg", "--no-cv"]) assert mock_run_task.call_args.kwargs["secagg"] is False - assert mock_run_task.call_args.kwargs["settings"]["secagg"] is False + _, json_data = _read_results( + tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age" + ) + assert json_data["settings"]["runtime_args"]["secagg"] is False def test_aggregator_options( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path @@ -962,20 +997,26 @@ def test_aggregator_options( assert kwargs["aggregator"] == "fedprox" assert kwargs["fedprox_mu"] == 5.0 assert kwargs["alpha"] == 0.5 - assert kwargs["fpath_results"].name.startswith("results-fedprox_5.0mu-") - assert "-0.5alpha-" in kwargs["fpath_results"].stem + fpath_results, _ = _read_results( + tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age" + ) + assert fpath_results.name.startswith("results-fedprox_5.0mu-") + assert "-0.5alpha-" in fpath_results.stem def test_no_tensorboard( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): module.main([str(tmp_path), "--task", "age", "--no-tensorboard", "--no-cv"]) assert mock_run_task.call_args.kwargs["dpath_tensorboard"] is None - assert mock_run_task.call_args.kwargs["settings"]["tensorboard"] is False + _, json_data = _read_results( + tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age" + ) + assert json_data["settings"]["runtime_args"]["tensorboard"] is False def test_existing_results( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): - _write_results(module, tmp_path.resolve(), "age", {"results": {}}) + _write_results(module, tmp_path.resolve(), "age", RESULTS_WITH_FOLD) with pytest.raises(SystemExit, match="--overwrite"): module.main([str(tmp_path), "--task", "sex", "--task", "age", "--no-cv"]) mock_run_task.assert_not_called() @@ -984,7 +1025,12 @@ def test_existing_settings_only( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): # a run that failed before any results were saved can be rerun - _write_results(module, tmp_path.resolve(), "age", {"settings": {}}) + _write_results( + module, + tmp_path.resolve(), + "age", + {"settings": {}, "folds": {"NODE_1__NODE_2": {"settings": {}}}}, + ) module.main([str(tmp_path), "--task", "age", "--no-cv"]) mock_run_task.assert_called_once() @@ -992,7 +1038,7 @@ def test_existing_results_other_date_tag( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): _write_results( - module, tmp_path.resolve(), "age", {"results": {}}, date_tag="old" + module, tmp_path.resolve(), "age", RESULTS_WITH_FOLD, date_tag="old" ) module.main([str(tmp_path), "--task", "age", "--no-cv"]) mock_run_task.assert_called_once() @@ -1000,7 +1046,7 @@ def test_existing_results_other_date_tag( def test_overwrite( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path ): - _write_results(module, tmp_path.resolve(), "age", {"results": {}}) + _write_results(module, tmp_path.resolve(), "age", RESULTS_WITH_FOLD) module.main([str(tmp_path), "--task", "age", "--overwrite", "--no-cv"]) mock_run_task.assert_called_once() @@ -1020,14 +1066,23 @@ def test_cv( kwargs = calls[0] assert kwargs["training_nodes"] == ["NODE_2", "NODE_3"] assert kwargs["node_data_held_out"] == {"NODE_1": NODE_DATA["NODE_1"]} - assert kwargs["fpath_results"].stem.endswith("-heldout_NODE_1") assert kwargs["fpath_model"].stem.endswith("-heldout_NODE_1") - fpath_log = kwargs["fpath_results"].with_name( - kwargs["fpath_results"].name.replace("results-", "logs-") - ) - assert fpath_log.with_suffix(".log").exists() assert kwargs["dpath_tensorboard"].name.endswith("-heldout_NODE_1") - assert kwargs["settings"]["cv"] is True + # one results file and one log file per task, with all the folds + dpath_date = tmp_path.resolve() / module.DEFAULT_DATE_TAG + for task in ["sex", "age"]: + fpath_results, json_data = _read_results(dpath_date / task) + assert fpath_results.name.endswith("-secagg-cv.json") + assert json_data["settings"]["runtime_args"]["cv"] is True + assert list(json_data["folds"]) == [ + "NODE_2__NODE_3", + "NODE_1__NODE_3", + "NODE_1__NODE_2", + ] + assert all("results" in fold for fold in json_data["folds"].values()) + assert [p.name for p in (dpath_date / task).glob("logs-*")] == [ + fpath_results.name.replace("results-", "logs-").replace(".json", ".log") + ] def test_no_cv( self, @@ -1082,10 +1137,8 @@ def test_cv_existing_results( mock_get_node_data: MagicMock, tmp_path: Path, ): - _write_results( - module, tmp_path.resolve(), "age", {"results": {}}, held_out_node="NODE_2" - ) - with pytest.raises(SystemExit, match="heldout_NODE_2"): + _write_results(module, tmp_path.resolve(), "age", RESULTS_WITH_FOLD, cv=True) + with pytest.raises(SystemExit, match="-cv.json"): module.main([str(tmp_path), "--task", "age"]) mock_run_task.assert_not_called() From 2a3711dc8a6532e21e3fd408f107bb324d15c0ba Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Sun, 1 Nov 2026 20:51:19 -0500 Subject: [PATCH 5/6] update log prefix/sectioning logic/format and make sure to log node warnings as well --- scripts/run_experiments.py | 63 +++++++++++++++------------ tests/scripts/test_run_experiments.py | 56 +++++++++++++++++++----- 2 files changed, 80 insertions(+), 39 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 701a820..a052b82 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -7,6 +7,7 @@ import json import logging import math +import re import shutil import sys from collections import defaultdict @@ -51,8 +52,9 @@ DEFAULT_CV = True DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d") -LOG_FORMAT_FILE = "%(asctime)s %(name)s %(levelname)s - %(message)s" -LOG_FORMAT_SUMMARY = "%(levelname)s - %(message)s" +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) @@ -303,13 +305,28 @@ 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}."): - # keep security audit events even if they are below WARNING level - return record.levelno >= logging.WARNING and not getattr( - record, "is_security", False - ) + 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 @@ -318,7 +335,7 @@ 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(logging.Formatter(LOG_FORMAT_FILE)) + handler.setFormatter(_StripAnsiFormatter(LOG_FORMAT_FILE)) handler.addFilter(_filter_log_file_record) for logger_ in (logger, fbm_logger): logger_.addHandler(handler) @@ -343,12 +360,6 @@ def _print_logger_messages_at_exit() -> Generator[None, None, None]: print("\n".join(handler.messages)) -def _get_log_prefix(task: str, held_out_node: str | None = None) -> str: - if held_out_node is None: - return f"[{task}]" - return f"[{task}|held_out={held_out_node}]" - - def _load_training_plan_class(module_name: str, class_name: str) -> type: spec = importlib.util.spec_from_file_location( f"training_plans.{module_name}", @@ -550,9 +561,10 @@ def _run_task( config = TASKS[task] is_cv = node_data_held_out is not None held_out_node = next(iter(node_data_held_out)) if is_cv else None - prefix = _get_log_prefix(task, held_out_node) + section_name = task.upper() if is_cv: - logger.info(f"{prefix} Holding out node {held_out_node}") + section_name += f": held_out_node={held_out_node}" + logger.info(f"===== {section_name} =====") exp = Experiment( tags=[COMMON_TAG, task], nodes=training_nodes, @@ -584,12 +596,12 @@ def _run_task( if not is_cv: logger.warning( - f"{prefix} Feature z-scoring statistics are computed on all node data," + "Feature z-scoring statistics are computed on all node data," " including the test data (data leakage). Use --cv to avoid this" ) if config.with_target_mean: logger.warning( - f"{prefix} The intercept is initialized with the target mean of all" + "The intercept is initialized with the target mean of all" " node data, including the test data (data leakage). Use --cv to" " avoid this" ) @@ -599,8 +611,8 @@ def _run_task( exp.analytics, exp.training_data().data(), config.target_column ) for reason, cols in excluded_columns.items(): - logger.info(f"{prefix} Excluding {len(cols)} features ({reason}): {cols}") - logger.info(f"{prefix} Using {len(columns)} features") + logger.info(f"Excluding {len(cols)} features ({reason}): {cols}") + logger.info(f"Using {len(columns)} features") model_args = _build_model_args(task, columns, mean, std, alpha, target_mean) exp.set_model_args(model_args) exp.set_training_plan_class( @@ -628,7 +640,7 @@ def _run_task( ) for round_ in range(1, n_rounds + 1): - logger.info(f"{prefix} Round {round_}/{n_rounds}") + logger.info(f"Round {round_}/{n_rounds}") if is_cv: _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments) exp.run_once() @@ -646,11 +658,11 @@ def _run_task( if dpath_tensorboard is not None: _move_tensorboard_logs(exp, dpath_tensorboard) - logger.info(f"{prefix} Saved TensorBoard logs to {dpath_tensorboard}") + logger.info(f"Saved TensorBoard logs to {dpath_tensorboard}") fpath_model.parent.mkdir(parents=True, exist_ok=True) exp.training_plan().export_model(str(fpath_model)) - logger.info(f"{prefix} Saved global model to {fpath_model}") + logger.info(f"Saved global model to {fpath_model}") def run_experiments( @@ -779,11 +791,8 @@ def run_experiments( ): json_data["folds"][fold_key] = fold _save_json(fpath_results, json_data) - logger.info( - f"{_get_log_prefix(task, held_out_node)} Saved settings and" - f" results to {fpath_results}" - ) - logger.info(f"[{task}] Saved log to {fpath_log}") + logger.info(f"Saved settings and results to {fpath_results}") + logger.info(f"Saved log to {fpath_log}") def main(argv: list[str] | None = None) -> None: diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 3934277..98b01ca 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -557,9 +557,40 @@ def test_log_to_file(self, module: ModuleType, tmp_path: Path): fbm_logger.error("fbm security", extra={"is_security": True}) lines = fpath_log.read_text().splitlines() assert len(lines) == 3 - assert lines[0].endswith("brainflex INFO - script info") - assert lines[1].endswith("brainflex WARNING - script warning") - assert lines[2].endswith("fedbiomed WARNING - fbm warning") + assert lines[0].endswith("brainflex INFO - script info") + assert lines[1].endswith("brainflex WARNING - script warning") + assert lines[2].endswith("fedbiomed WARNING - fbm warning") + + @pytest.mark.parametrize( + ("node_level", "kept"), [("WARNING", True), ("ERROR", True), ("INFO", False)] + ) + def test_node_log_message( + self, module: ModuleType, tmp_path: Path, node_level: str, kept: bool + ): + fpath_log = tmp_path / "logs.log" + + # same name and message format as in Fed-BioMed's Requests class + def print_node_log_message() -> None: + logging.getLogger("fedbiomed").info( + f"\033[1m{node_level}\033[0m\n" + "\t\t\t\t\t\033[1m NODE\033[0m NODE_1\n" + "\t\t\t\t\t\033[1m MESSAGE:\033[0m node message\033[0m\n" + ) + + with module._log_to_file(fpath_log): + print_node_log_message() + text = fpath_log.read_text() + assert ("node message" in text) is kept + if kept: + assert f"fedbiomed INFO - {node_level}\n" in text + assert "\033" not in text + + def test_info_not_from_node(self, module: ModuleType, tmp_path: Path): + # the first line is only parsed for messages forwarded from the nodes + fpath_log = tmp_path / "logs.log" + with module._log_to_file(fpath_log): + logging.getLogger("fedbiomed").info("WARNING\nnot from a node") + assert fpath_log.read_text() == "" def test_handler_removed_on_error(self, module: ModuleType, tmp_path: Path): fbm_logger = logging.getLogger("fedbiomed") @@ -581,7 +612,7 @@ def test_print_logger_messages_at_exit( assert capsys.readouterr().out == "" module.logger.warning("[age] second") assert capsys.readouterr().out == ( - "INFO - [age] first\nWARNING - [age] second\n" + "INFO - [age] first\nWARNING - [age] second\n" ) assert not any( isinstance(handler, module._ListHandler) @@ -592,7 +623,7 @@ def test_printed_on_error(self, module: ModuleType, capsys: pytest.CaptureFixtur with pytest.raises(RuntimeError), module._print_logger_messages_at_exit(): module.logger.info("[age] before error") raise RuntimeError - assert capsys.readouterr().out == "INFO - [age] before error\n" + assert capsys.readouterr().out == "INFO - [age] before error\n" class TestRunTask: @@ -641,7 +672,9 @@ def test_run_task( exp.set_training_plan_class.call_args.args[0].__name__ == "TabularRegressorTrainingPlan" ) - assert "[age] Using 3 features" in (tmp_path / "logs.log").read_text() + log_text = (tmp_path / "logs.log").read_text() + assert "INFO - ===== AGE =====" in log_text + assert "INFO - Using 3 features" in log_text # before training and after each round assert len(folds) == 4 @@ -841,8 +874,7 @@ def test_held_out_node( assert fold["settings"]["held_out_node"] == node_data_held_out assert "node_metrics" in fold["results"] assert not any("data leakage" in r.getMessage() for r in caplog.records) - assert "[age|held_out=NODE_3] Holding out node NODE_3" in caplog.messages - assert "[age|held_out=NODE_3] Round 1/3" in caplog.messages + assert "===== AGE: held_out=NODE_3 =====" in caplog.messages def _read_results(dpath_task: Path) -> tuple[Path, dict]: @@ -953,11 +985,11 @@ def run_task(task: str, **kwargs) -> Generator[tuple[str, dict], None, None]: with pytest.raises(RuntimeError): module.main([str(tmp_path), "--task", "sex", "--task", "age", "--no-cv"]) lines = capsys.readouterr().out.splitlines() - assert lines[0] == "INFO - [sex] done" - assert lines[1].startswith("INFO - [sex] Saved settings and results to ") - assert lines[2].startswith("INFO - [sex] Saved log to ") + assert lines[0] == "INFO - [sex] done" + assert lines[1].startswith("INFO - Saved settings and results to ") + assert lines[2].startswith("INFO - Saved log to ") # no "Saved" messages for the failed run - assert lines[3:] == ["INFO - [age] done"] + assert lines[3:] == ["INFO - [age] done"] def test_date_tag( self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path From 1f8141a81cb10670915301da0fb374b402dab547 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Sun, 1 Nov 2026 21:16:55 -0500 Subject: [PATCH 6/6] extract utilities into separate files --- scripts/_io_utils.py | 76 ++++++ scripts/_log_utils.py | 81 ++++++ scripts/_tasks.py | 94 +++++++ scripts/request_training_plan_approval.py | 6 +- scripts/run_experiments.py | 255 ++---------------- tests/scripts/conftest.py | 36 +++ tests/scripts/test_io_utils.py | 84 ++++++ tests/scripts/test_log_utils.py | 98 +++++++ .../test_request_training_plan_approval.py | 28 +- tests/scripts/test_run_experiments.py | 230 +--------------- tests/scripts/test_tasks.py | 50 ++++ 11 files changed, 550 insertions(+), 488 deletions(-) create mode 100644 scripts/_io_utils.py create mode 100644 scripts/_log_utils.py create mode 100644 scripts/_tasks.py create mode 100644 tests/scripts/conftest.py create mode 100644 tests/scripts/test_io_utils.py create mode 100644 tests/scripts/test_log_utils.py create mode 100644 tests/scripts/test_tasks.py diff --git a/scripts/_io_utils.py b/scripts/_io_utils.py new file mode 100644 index 0000000..cc06f49 --- /dev/null +++ b/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" + ) diff --git a/scripts/_log_utils.py b/scripts/_log_utils.py new file mode 100644 index 0000000..28f69f3 --- /dev/null +++ b/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)) diff --git a/scripts/_tasks.py b/scripts/_tasks.py new file mode 100644 index 0000000..f28ef39 --- /dev/null +++ b/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.metrics import MetricTypes + +DPATH_TRAINING_PLANS = Path(__file__).resolve().parents[1] / "training_plans" + +COMMON_TAG = "brainflex" + + +@dataclass(frozen=True) +class TaskConfig: + training_plan_module: str + training_plan_class: str + target_column: str + test_metric: MetricTypes + # 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_module="tabular_regressor", + training_plan_class="TabularRegressorTrainingPlan", + target_column="nb:Age", + test_metric=MetricTypes.MEAN_SQUARE_ERROR, + with_target_mean=True, + secagg_clipping_range=100, + ), + "sex": TaskConfig( + training_plan_module="tabular_classifier", + training_plan_class="TabularClassifierTrainingPlan", + target_column="nb:Sex", + test_metric=MetricTypes.ACCURACY, + 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", + ] + + +def load_training_plan_class(module_name: str, class_name: str) -> type: + 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) diff --git a/scripts/request_training_plan_approval.py b/scripts/request_training_plan_approval.py index 5d3a9ad..b0509ea 100755 --- a/scripts/request_training_plan_approval.py +++ b/scripts/request_training_plan_approval.py @@ -5,6 +5,7 @@ import argparse import sys +from _tasks import COMMON_TAG, TASKS, load_training_plan_class from fedbiomed.common.constants import TrainingPlanApprovalStatus from fedbiomed.researcher.federated_workflows import Experiment from fedbiomed.researcher.federated_workflows.jobs import ( @@ -12,9 +13,6 @@ TrainingPlanCheckJob, ) -# importable because the script's directory is the first entry of sys.path -from run_experiments import COMMON_TAG, TASKS, _load_training_plan_class - # 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" @@ -43,7 +41,7 @@ def _request_approval(task: str) -> dict[str, str]: nodes = fds.node_ids() if fds is not None else [] if not nodes: raise ValueError(f"No node has a dataset for {task}") - training_plan = _load_training_plan_class( + training_plan = load_training_plan_class( config.training_plan_module, config.training_plan_class )() diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index a052b82..09e7f79 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -3,24 +3,26 @@ """Run the federated learning experiments.""" import argparse -import importlib.util import json -import logging import math -import re import shutil import sys from collections import defaultdict from collections.abc import Generator -from contextlib import contextmanager -from dataclasses import dataclass, field from datetime import date -from enum import Enum from pathlib import Path -import numpy as np +from _io_utils import ( + PREFIX_LOGS, + PREFIX_MODEL, + PREFIX_RESULTS, + PREFIX_TENSORBOARD, + PathHelper, + save_json, +) +from _log_utils import log_to_file, logger, print_logger_messages_at_exit +from _tasks import COMMON_TAG, TASKS, filter_feature_col, load_training_plan_class from declearn.optimizer.modules import ScaffoldServerModule -from fedbiomed.common.metrics import MetricTypes from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_args import TrainingArgs from fedbiomed.researcher.datasets import FederatedDataset @@ -28,14 +30,6 @@ from fedbiomed.researcher.federated_workflows.jobs import TrainingJob from fedbiomed.researcher.secagg import SecureAggregation -DPATH_TRAINING_PLANS = Path(__file__).resolve().parents[1] / "training_plans" -PREFIX_MODEL = "model" -PREFIX_RESULTS = "results" -PREFIX_TENSORBOARD = "tensorboard" -PREFIX_LOGS = "logs" - -COMMON_TAG = "brainflex" - AGGREGATORS = ["fedavg", "fedprox", "scaffold"] DEFAULT_TASKS = ["age", "sex"] @@ -52,92 +46,6 @@ DEFAULT_CV = True DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d") -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") - - -@dataclass(frozen=True) -class TaskConfig: - training_plan_module: str - training_plan_class: str - target_column: str - test_metric: MetricTypes - # 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 - - -@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}" - - -TASKS = { - "age": TaskConfig( - training_plan_module="tabular_regressor", - training_plan_class="TabularRegressorTrainingPlan", - target_column="nb:Age", - test_metric=MetricTypes.MEAN_SQUARE_ERROR, - with_target_mean=True, - secagg_clipping_range=100, - ), - "sex": TaskConfig( - training_plan_module="tabular_classifier", - training_plan_class="TabularClassifierTrainingPlan", - target_column="nb:Sex", - test_metric=MetricTypes.ACCURACY, - extra_model_args={"n_classes": 2}, - ), -} - def _get_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) @@ -239,139 +147,6 @@ def _get_parser() -> argparse.ArgumentParser: return parser -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", - ] - - -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" - ) - - -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)) - - -def _load_training_plan_class(module_name: str, class_name: str) -> type: - 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) - - def _get_node_data(task: str) -> dict[str, dict]: """Get the dataset metadata of each node that has data for a task.""" fds = Experiment(tags=[COMMON_TAG, task]).training_data() @@ -395,7 +170,7 @@ def _get_feature_stats( excluded_columns["not in the datasets of all nodes"] = missing_columns stats_columns: list[str] = [] for column in sorted(common_columns): - if column == target_column or _filter_feature_col(column): + if column == target_column or filter_feature_col(column): stats_columns.append(column) else: excluded_columns["not in feature column list"].append(column) @@ -616,7 +391,7 @@ def _run_task( model_args = _build_model_args(task, columns, mean, std, alpha, target_mean) exp.set_model_args(model_args) exp.set_training_plan_class( - _load_training_plan_class( + load_training_plan_class( config.training_plan_module, config.training_plan_class ) ) @@ -737,7 +512,7 @@ def run_experiments( " Use --overwrite to overwrite them" ) - with _print_logger_messages_at_exit(): + with print_logger_messages_at_exit(): for task, held_out_nodes in runs.items(): fpath_results = path_helper.get_fpath_out( task=task, prefix=PREFIX_RESULTS, extension=".json" @@ -747,7 +522,7 @@ def run_experiments( ) json_data = {"settings": {"runtime_args": settings}, "folds": {}} fpath_results.parent.mkdir(parents=True, exist_ok=True) - with _log_to_file(fpath_log): + with log_to_file(fpath_log): for held_out_node in held_out_nodes: if held_out_node is None: # use all nodes for training @@ -790,7 +565,7 @@ def run_experiments( secagg=secagg, ): json_data["folds"][fold_key] = fold - _save_json(fpath_results, json_data) + save_json(fpath_results, json_data) logger.info(f"Saved settings and results to {fpath_results}") logger.info(f"Saved log to {fpath_log}") diff --git a/tests/scripts/conftest.py b/tests/scripts/conftest.py new file mode 100644 index 0000000..d6705b5 --- /dev/null +++ b/tests/scripts/conftest.py @@ -0,0 +1,36 @@ +import importlib.util +import sys +from collections.abc import Callable +from pathlib import Path +from types import ModuleType + +import pytest + +DPATH_SCRIPTS = Path(__file__).parents[2] / "scripts" + + +@pytest.fixture +def load_script( + monkeypatch: pytest.MonkeyPatch, tmp_path_factory: pytest.TempPathFactory +) -> Callable[[str], ModuleType]: + """Load a module from the scripts directory.""" + # importing fedbiomed.researcher creates a component directory + monkeypatch.setenv( + "FBM_RESEARCHER_COMPONENT_ROOT", + str(tmp_path_factory.getbasetemp() / "fbm-researcher"), + ) + # copy so that modules imported by the script are not kept + monkeypatch.setattr(sys, "modules", dict(sys.modules)) + # have scripts directory be first in syspath + monkeypatch.syspath_prepend(str(DPATH_SCRIPTS)) + + def _load(name: str) -> ModuleType: + spec = importlib.util.spec_from_file_location( + name, DPATH_SCRIPTS / f"{name}.py" + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + return _load diff --git a/tests/scripts/test_io_utils.py b/tests/scripts/test_io_utils.py new file mode 100644 index 0000000..375bc9e --- /dev/null +++ b/tests/scripts/test_io_utils.py @@ -0,0 +1,84 @@ +import dataclasses +from collections.abc import Callable +from enum import Enum +from pathlib import Path +from types import ModuleType + +import numpy as np +import pytest + + +@pytest.fixture +def module(load_script: Callable[[str], ModuleType]) -> ModuleType: + return load_script("_io_utils") + + +@pytest.fixture +def path_helper(module: ModuleType, tmp_path: Path) -> object: + return module.PathHelper( + output_dir=tmp_path, + date_tag="2026_01_01", + n_rounds=10, + batch_size=32, + n_updates=50, + learning_rate=0.001, + test_ratio=0.2, + aggregator="scaffold", + fedprox_mu=2.0, + alpha=0.1, + secagg=True, + cv=True, + ) + + +class TestPathHelper: + def test_get_fpath_out(self, path_helper: object, tmp_path: Path): + assert path_helper.get_fpath_out("age", "results") == ( + tmp_path + / "2026_01_01" + / "age" + / ( + "results-scaffold-10rounds-50updates-32-0.001lr-0.2test-0.1alpha" + "-secagg-cv.json" + ) + ) + + def test_no_secagg_no_cv(self, path_helper: object): + path_helper = dataclasses.replace(path_helper, secagg=False, cv=False) + fpath = path_helper.get_fpath_out("age", "results") + assert fpath.stem.endswith("-0.1alpha-no_secagg-no_cv") + + def test_fedprox(self, path_helper: object): + path_helper = dataclasses.replace(path_helper, aggregator="fedprox") + fpath = path_helper.get_fpath_out("age", "results") + assert fpath.name.startswith("results-fedprox_2.0mu-") + + def test_extension(self, path_helper: object, tmp_path: Path): + fpath = path_helper.get_fpath_out("sex", "model", ".pkl") + assert fpath.parent == tmp_path / "2026_01_01" / "sex" + assert fpath.suffix == ".pkl" + + def test_held_out_node(self, path_helper: object): + fpath = path_helper.get_fpath_out( + "age", "model", ".pkl", held_out_node="NODE_1" + ) + assert fpath.name.endswith("-0.1alpha-secagg-cv-heldout_NODE_1.pkl") + + +class TestSerialize: + @pytest.mark.parametrize( + "obj,expected", + [ + (Enum("Color", "RED").RED, "RED"), + (Path("a/b"), "a/b"), + (np.array([1.0, 2.0]), [1.0, 2.0]), + (np.float64(1.5), 1.5), + ({"b", "a"}, ["a", "b"]), + ], + ) + def test_serialize(self, module: ModuleType, obj: object, expected: object): + assert module.serialize(obj) == expected + + def test_not_serializable(self, module: ModuleType): + with pytest.raises(TypeError, match="not serializable"): + module.serialize(object()) diff --git a/tests/scripts/test_log_utils.py b/tests/scripts/test_log_utils.py new file mode 100644 index 0000000..716cc87 --- /dev/null +++ b/tests/scripts/test_log_utils.py @@ -0,0 +1,98 @@ +import logging +from collections.abc import Callable +from pathlib import Path +from types import ModuleType + +import pytest + + +@pytest.fixture +def module(load_script: Callable[[str], ModuleType]) -> ModuleType: + return load_script("_log_utils") + + +@pytest.fixture(autouse=True) +def fbm_logger_level(caplog: pytest.LogCaptureFixture) -> None: + # set by Fed-BioMed's logger, which is not imported here + caplog.set_level(logging.INFO, logger="fedbiomed") + + +class TestLogToFile: + def test_log_to_file(self, module: ModuleType, tmp_path: Path): + fpath_log = tmp_path / "logs" / "logs.log" + fbm_logger = logging.getLogger("fedbiomed") + with module.log_to_file(fpath_log): + module.logger.info("script info") + module.logger.warning("script warning") + fbm_logger.info("fbm info") + fbm_logger.warning("fbm warning") + fbm_logger.error("fbm security", extra={"is_security": True}) + lines = fpath_log.read_text().splitlines() + assert len(lines) == 3 + assert lines[0].endswith("brainflex INFO - script info") + assert lines[1].endswith("brainflex WARNING - script warning") + assert lines[2].endswith("fedbiomed WARNING - fbm warning") + + @pytest.mark.parametrize( + ("node_level", "kept"), [("WARNING", True), ("ERROR", True), ("INFO", False)] + ) + def test_node_log_message( + self, module: ModuleType, tmp_path: Path, node_level: str, kept: bool + ): + fpath_log = tmp_path / "logs.log" + + # same name and message format as in Fed-BioMed's Requests class + def print_node_log_message() -> None: + logging.getLogger("fedbiomed").info( + f"\033[1m{node_level}\033[0m\n" + "\t\t\t\t\t\033[1m NODE\033[0m NODE_1\n" + "\t\t\t\t\t\033[1m MESSAGE:\033[0m node message\033[0m\n" + ) + + with module.log_to_file(fpath_log): + print_node_log_message() + text = fpath_log.read_text() + assert ("node message" in text) is kept + if kept: + assert f"fedbiomed INFO - {node_level}\n" in text + assert "\033" not in text + + def test_info_not_from_node(self, module: ModuleType, tmp_path: Path): + # the first line is only parsed for messages forwarded from the nodes + fpath_log = tmp_path / "logs.log" + with module.log_to_file(fpath_log): + logging.getLogger("fedbiomed").info("WARNING\nnot from a node") + assert fpath_log.read_text() == "" + + def test_handler_removed_on_error(self, module: ModuleType, tmp_path: Path): + fbm_logger = logging.getLogger("fedbiomed") + handlers_before = (list(module.logger.handlers), list(fbm_logger.handlers)) + with pytest.raises(RuntimeError), module.log_to_file(tmp_path / "logs.log"): + raise RuntimeError + assert (list(module.logger.handlers), list(fbm_logger.handlers)) == ( + handlers_before + ) + + +class TestPrintLoggerMessagesAtExit: + def test_print_logger_messages_at_exit( + self, module: ModuleType, capsys: pytest.CaptureFixture + ): + with module.print_logger_messages_at_exit(): + module.logger.info("[age] first") + logging.getLogger("fedbiomed").warning("not printed") + assert capsys.readouterr().out == "" + module.logger.warning("[age] second") + assert capsys.readouterr().out == ( + "INFO - [age] first\nWARNING - [age] second\n" + ) + assert not any( + isinstance(handler, module._ListHandler) + for handler in module.logger.handlers + ) + + def test_printed_on_error(self, module: ModuleType, capsys: pytest.CaptureFixture): + with pytest.raises(RuntimeError), module.print_logger_messages_at_exit(): + module.logger.info("[age] before error") + raise RuntimeError + assert capsys.readouterr().out == "INFO - [age] before error\n" diff --git a/tests/scripts/test_request_training_plan_approval.py b/tests/scripts/test_request_training_plan_approval.py index 7344a2e..f6a09f5 100644 --- a/tests/scripts/test_request_training_plan_approval.py +++ b/tests/scripts/test_request_training_plan_approval.py @@ -1,6 +1,4 @@ -import importlib.util -import sys -from pathlib import Path +from collections.abc import Callable from types import ModuleType from unittest.mock import MagicMock @@ -9,32 +7,12 @@ from fedbiomed.common.training_plans import BaseTrainingPlan from pytest_mock import MockerFixture -FPATH_SCRIPT = ( - Path(__file__).parents[2] / "scripts" / "request_training_plan_approval.py" -) NODES = ["NODE_1", "NODE_2", "NODE_3", "NODE_4"] @pytest.fixture -def module( - monkeypatch: pytest.MonkeyPatch, tmp_path_factory: pytest.TempPathFactory -) -> ModuleType: - # importing fedbiomed.researcher creates a component directory - monkeypatch.setenv( - "FBM_RESEARCHER_COMPONENT_ROOT", - str(tmp_path_factory.getbasetemp() / "fbm-researcher"), - ) - # copy so that modules imported by the script are not kept - monkeypatch.setattr(sys, "modules", dict(sys.modules)) - # as when the script is run - monkeypatch.syspath_prepend(str(FPATH_SCRIPT.parent)) - spec = importlib.util.spec_from_file_location( - "request_training_plan_approval", FPATH_SCRIPT - ) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module +def module(load_script: Callable[[str], ModuleType]) -> ModuleType: + return load_script("request_training_plan_approval") @pytest.fixture diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 98b01ca..6f6886f 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -1,10 +1,7 @@ -import dataclasses -import importlib.util import json import logging import sys -from collections.abc import Generator -from enum import Enum +from collections.abc import Callable, Generator from pathlib import Path from types import ModuleType from unittest.mock import MagicMock @@ -13,10 +10,8 @@ import pandas as pd import pytest from fedbiomed.common.training_args import TrainingArgs -from fedbiomed.common.training_plans import BaseTrainingPlan from pytest_mock import MockerFixture -FPATH_SCRIPT = Path(__file__).parents[2] / "scripts" / "run_experiments.py" GLOBAL_STATS = { "mean": { "nb:Age": np.float64(30.0), @@ -50,21 +45,8 @@ @pytest.fixture -def module( - monkeypatch: pytest.MonkeyPatch, tmp_path_factory: pytest.TempPathFactory -) -> ModuleType: - # importing fedbiomed.researcher creates a component directory - monkeypatch.setenv( - "FBM_RESEARCHER_COMPONENT_ROOT", - str(tmp_path_factory.getbasetemp() / "fbm-researcher"), - ) - # copy so that modules imported by the script are not kept - monkeypatch.setattr(sys, "modules", dict(sys.modules)) - spec = importlib.util.spec_from_file_location("run_experiments", FPATH_SCRIPT) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module +def module(load_script: Callable[[str], ModuleType]) -> ModuleType: + return load_script("run_experiments") @pytest.fixture @@ -153,69 +135,6 @@ def fpath_csv(tmp_path: Path) -> Path: return fpath -class TestLoadTrainingPlanClass: - @pytest.mark.parametrize("task", ["age", "sex"]) - def test_source(self, module: ModuleType, task: str): - config = module.TASKS[task] - plan_cls = module._load_training_plan_class( - config.training_plan_module, config.training_plan_class - ) - assert issubclass(plan_cls, BaseTrainingPlan) - assert f"class {config.training_plan_class}" in plan_cls().source() - - -@pytest.fixture -def path_helper(module: ModuleType, tmp_path: Path) -> object: - return module.PathHelper( - output_dir=tmp_path, - date_tag="2026_01_01", - n_rounds=10, - batch_size=32, - n_updates=50, - learning_rate=0.001, - test_ratio=0.2, - aggregator="scaffold", - fedprox_mu=2.0, - alpha=0.1, - secagg=True, - cv=True, - ) - - -class TestPathHelper: - def test_get_fpath_out(self, path_helper: object, tmp_path: Path): - assert path_helper.get_fpath_out("age", "results") == ( - tmp_path - / "2026_01_01" - / "age" - / ( - "results-scaffold-10rounds-50updates-32-0.001lr-0.2test-0.1alpha" - "-secagg-cv.json" - ) - ) - - def test_no_secagg_no_cv(self, path_helper: object): - path_helper = dataclasses.replace(path_helper, secagg=False, cv=False) - fpath = path_helper.get_fpath_out("age", "results") - assert fpath.stem.endswith("-0.1alpha-no_secagg-no_cv") - - def test_fedprox(self, path_helper: object): - path_helper = dataclasses.replace(path_helper, aggregator="fedprox") - fpath = path_helper.get_fpath_out("age", "results") - assert fpath.name.startswith("results-fedprox_2.0mu-") - - def test_extension(self, path_helper: object, tmp_path: Path): - fpath = path_helper.get_fpath_out("sex", "model", ".pkl") - assert fpath.parent == tmp_path / "2026_01_01" / "sex" - assert fpath.suffix == ".pkl" - - def test_held_out_node(self, path_helper: object): - fpath = path_helper.get_fpath_out( - "age", "model", ".pkl", held_out_node="NODE_1" - ) - assert fpath.name.endswith("-0.1alpha-secagg-cv-heldout_NODE_1.pkl") - - class TestGetNodeData: def test_get_node_data(self, module: ModuleType, mock_experiment: MagicMock): exp = mock_experiment.return_value @@ -322,35 +241,6 @@ def test_setup_held_out_secagg( assert secagg_arguments == secagg.train_arguments.return_value -class TestFilterFeatureCol: - @pytest.mark.parametrize( - "col", - [ - "nb:Age", - "nb:Sex", - "EstimatedTotalIntraCranialVol", - "Left-Hippocampus", - "Right-Accumbens-area", - ], - ) - def test_included(self, module: ModuleType, col: str): - assert module._filter_feature_col(col) - - @pytest.mark.parametrize( - "col", - [ - "Left-WM-hypointensities", - "non-WM-hypointensities", - "CSF", - "5th-Ventricle", - "lhSurfaceHoles", - "lh_insula", - ], - ) - def test_excluded(self, module: ModuleType, col: str): - assert not module._filter_feature_col(col) - - class TestGetFeatureStats: def test_target_dropped(self, module: ModuleType, analytics: MagicMock): columns, mean, std, excluded_columns, target_mean = module._get_feature_stats( @@ -449,7 +339,7 @@ def test_training_data( task, columns, mean, std, 0.1, target_mean ) - plan = module._load_training_plan_class( + plan = module.load_training_plan_class( config.training_plan_module, config.training_plan_class )() plan.post_init(model_args, TrainingArgs({}, only_required=False)) @@ -498,25 +388,6 @@ def test_other(self, module: ModuleType, aggregator: str): assert module._get_agg_optimizer(aggregator) is None -class TestSerialize: - @pytest.mark.parametrize( - "obj,expected", - [ - (Enum("Color", "RED").RED, "RED"), - (Path("a/b"), "a/b"), - (np.array([1.0, 2.0]), [1.0, 2.0]), - (np.float64(1.5), 1.5), - ({"b", "a"}, ["a", "b"]), - ], - ) - def test_serialize(self, module: ModuleType, obj: object, expected: object): - assert module._serialize(obj) == expected - - def test_not_serializable(self, module: ModuleType): - with pytest.raises(TypeError, match="not serializable"): - module._serialize(object()) - - class TestMoveTensorboardLogs: def test_move( self, module: ModuleType, dpath_fbm_tensorboard: Path, tmp_path: Path @@ -545,93 +416,12 @@ def test_overwrite_node_dir( ] -class TestLogToFile: - def test_log_to_file(self, module: ModuleType, tmp_path: Path): - fpath_log = tmp_path / "logs" / "logs.log" - fbm_logger = logging.getLogger("fedbiomed") - with module._log_to_file(fpath_log): - module.logger.info("script info") - module.logger.warning("script warning") - fbm_logger.info("fbm info") - fbm_logger.warning("fbm warning") - fbm_logger.error("fbm security", extra={"is_security": True}) - lines = fpath_log.read_text().splitlines() - assert len(lines) == 3 - assert lines[0].endswith("brainflex INFO - script info") - assert lines[1].endswith("brainflex WARNING - script warning") - assert lines[2].endswith("fedbiomed WARNING - fbm warning") - - @pytest.mark.parametrize( - ("node_level", "kept"), [("WARNING", True), ("ERROR", True), ("INFO", False)] - ) - def test_node_log_message( - self, module: ModuleType, tmp_path: Path, node_level: str, kept: bool - ): - fpath_log = tmp_path / "logs.log" - - # same name and message format as in Fed-BioMed's Requests class - def print_node_log_message() -> None: - logging.getLogger("fedbiomed").info( - f"\033[1m{node_level}\033[0m\n" - "\t\t\t\t\t\033[1m NODE\033[0m NODE_1\n" - "\t\t\t\t\t\033[1m MESSAGE:\033[0m node message\033[0m\n" - ) - - with module._log_to_file(fpath_log): - print_node_log_message() - text = fpath_log.read_text() - assert ("node message" in text) is kept - if kept: - assert f"fedbiomed INFO - {node_level}\n" in text - assert "\033" not in text - - def test_info_not_from_node(self, module: ModuleType, tmp_path: Path): - # the first line is only parsed for messages forwarded from the nodes - fpath_log = tmp_path / "logs.log" - with module._log_to_file(fpath_log): - logging.getLogger("fedbiomed").info("WARNING\nnot from a node") - assert fpath_log.read_text() == "" - - def test_handler_removed_on_error(self, module: ModuleType, tmp_path: Path): - fbm_logger = logging.getLogger("fedbiomed") - handlers_before = (list(module.logger.handlers), list(fbm_logger.handlers)) - with pytest.raises(RuntimeError), module._log_to_file(tmp_path / "logs.log"): - raise RuntimeError - assert (list(module.logger.handlers), list(fbm_logger.handlers)) == ( - handlers_before - ) - - -class TestPrintLoggerMessagesAtExit: - def test_print_logger_messages_at_exit( - self, module: ModuleType, capsys: pytest.CaptureFixture - ): - with module._print_logger_messages_at_exit(): - module.logger.info("[age] first") - logging.getLogger("fedbiomed").warning("not printed") - assert capsys.readouterr().out == "" - module.logger.warning("[age] second") - assert capsys.readouterr().out == ( - "INFO - [age] first\nWARNING - [age] second\n" - ) - assert not any( - isinstance(handler, module._ListHandler) - for handler in module.logger.handlers - ) - - def test_printed_on_error(self, module: ModuleType, capsys: pytest.CaptureFixture): - with pytest.raises(RuntimeError), module._print_logger_messages_at_exit(): - module.logger.info("[age] before error") - raise RuntimeError - assert capsys.readouterr().out == "INFO - [age] before error\n" - - class TestRunTask: def test_run_task( self, module: ModuleType, mock_experiment: MagicMock, tmp_path: Path ): fpath_model = tmp_path / "age" / "model.pkl" - with module._log_to_file(tmp_path / "logs.log"): + with module.log_to_file(tmp_path / "logs.log"): folds = list( module._run_task( "age", @@ -679,7 +469,9 @@ def test_run_task( # before training and after each round assert len(folds) == 4 assert {fold_key for fold_key, _ in folds} == {"NODE_1__NODE_2__NODE_3"} - fold = json.loads(json.dumps(folds[-1][1], default=module._serialize)) + fold = json.loads( + json.dumps(folds[-1][1], default=sys.modules["_io_utils"].serialize) + ) assert fold["settings"]["excluded_columns"] == {} model_args = fold["settings"]["model_args"] assert model_args["target_column"] == "nb:Age" @@ -792,7 +584,7 @@ def test_leakage_warnings( task: str, n_warnings: int, ): - with module._log_to_file(tmp_path / "logs.log"): + with module.log_to_file(tmp_path / "logs.log"): list( module._run_task( task, @@ -874,7 +666,7 @@ def test_held_out_node( assert fold["settings"]["held_out_node"] == node_data_held_out assert "node_metrics" in fold["results"] assert not any("data leakage" in r.getMessage() for r in caplog.records) - assert "===== AGE: held_out=NODE_3 =====" in caplog.messages + assert "===== AGE: held_out_node=NODE_3 =====" in caplog.messages def _read_results(dpath_task: Path) -> tuple[Path, dict]: @@ -943,7 +735,7 @@ def test_tasks_in_order( assert fpath_results.name.startswith("results-scaffold-2rounds-") assert kwargs["fpath_model"].name.startswith("model-scaffold-2rounds-") assert kwargs["fpath_model"].suffix == ".pkl" - # created by _log_to_file even though _run_task is mocked + # created by log_to_file even though _run_task is mocked fpath_log = dpath_date / "age" / fpath_results.name.replace("results-", "logs-") assert fpath_log.with_suffix(".log").exists() assert kwargs["dpath_tensorboard"].parent == dpath_date / "age" diff --git a/tests/scripts/test_tasks.py b/tests/scripts/test_tasks.py new file mode 100644 index 0000000..ec7a4a7 --- /dev/null +++ b/tests/scripts/test_tasks.py @@ -0,0 +1,50 @@ +from collections.abc import Callable +from types import ModuleType + +import pytest +from fedbiomed.common.training_plans import BaseTrainingPlan + + +@pytest.fixture +def module(load_script: Callable[[str], ModuleType]) -> ModuleType: + return load_script("_tasks") + + +class TestLoadTrainingPlanClass: + @pytest.mark.parametrize("task", ["age", "sex"]) + def test_source(self, module: ModuleType, task: str): + config = module.TASKS[task] + plan_cls = module.load_training_plan_class( + config.training_plan_module, config.training_plan_class + ) + assert issubclass(plan_cls, BaseTrainingPlan) + assert f"class {config.training_plan_class}" in plan_cls().source() + + +class TestFilterFeatureCol: + @pytest.mark.parametrize( + "col", + [ + "nb:Age", + "nb:Sex", + "EstimatedTotalIntraCranialVol", + "Left-Hippocampus", + "Right-Accumbens-area", + ], + ) + def test_included(self, module: ModuleType, col: str): + assert module.filter_feature_col(col) + + @pytest.mark.parametrize( + "col", + [ + "Left-WM-hypointensities", + "non-WM-hypointensities", + "CSF", + "5th-Ventricle", + "lhSurfaceHoles", + "lh_insula", + ], + ) + def test_excluded(self, module: ModuleType, col: str): + assert not module.filter_feature_col(col)