From 12112d139ea8e7d79213c7f04a89a4e91412a161 Mon Sep 17 00:00:00 2001 From: Michelle Wang Date: Thu, 5 Nov 2026 21:15:56 -0500 Subject: [PATCH] fix --- scripts/run_experiments.py | 36 +++++++++++++----- tests/scripts/test_run_experiments.py | 55 ++++++++++++++++++++++----- 2 files changed, 73 insertions(+), 18 deletions(-) diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index 3fd6194..d815b5a 100755 --- a/scripts/run_experiments.py +++ b/scripts/run_experiments.py @@ -23,11 +23,13 @@ from _log_utils import log_to_file, logger, print_logger_messages_at_exit from _tasks import COMMON_TAG, TASKS, filter_feature_col from declearn.optimizer.modules import ScaffoldServerModule +from fedbiomed.common.message import PingRequest from fedbiomed.common.optimizers.optimizer import Optimizer from fedbiomed.common.training_args import TrainingArgs from fedbiomed.researcher.datasets import FederatedDataset from fedbiomed.researcher.federated_workflows import Experiment, FederatedAnalytics from fedbiomed.researcher.federated_workflows.jobs import TrainingJob +from fedbiomed.researcher.requests import DiscardOnTimeout from fedbiomed.researcher.secagg import SecureAggregation AGGREGATORS = ["fedavg", "fedprox", "scaffold"] @@ -157,10 +159,28 @@ def _get_parser() -> argparse.ArgumentParser: return parser -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() - return fds.data() if fds is not None else {} +def _get_node_names(exp: Experiment, nodes: list[str]) -> dict[str, str | None]: + """Get the name of each node.""" + message = PingRequest(researcher_id=exp.researcher_id) + with exp.requests.send(message, nodes, policies=[DiscardOnTimeout(5)]) as request: + replies = request.replies() + node_names = { + node: replies[node].node_name if node in replies else None + for node in sorted(nodes) + } + if missing := [node for node, name in node_names.items() if name is None]: + logger.warning(f"Could not get the name of nodes: {missing}") + return node_names + + +def _get_node_data(task: str) -> tuple[dict[str, dict], dict[str, str | None]]: + """Get the dataset metadata and the name of each node that has data for a task.""" + exp = Experiment(tags=[COMMON_TAG, task]) + fds = exp.training_data() + if fds is None: + return {}, {} + node_data = fds.data() + return node_data, _get_node_names(exp, list(node_data)) def _get_column_stats( @@ -480,8 +500,9 @@ def run_experiments( # values: held-out node (None if not doing CV) runs: dict[str, list[str | None]] = {} node_data: dict[str, dict[str, dict]] = {} + node_names: dict[str, dict[str, str | None]] = {} for task in tasks: - node_data[task] = _get_node_data(task) + node_data[task], node_names[task] = _get_node_data(task) if not cv: runs[task] = [None] continue @@ -539,10 +560,7 @@ def run_experiments( ) json_data = { "settings": settings, - "node_names": { - node: node_data[task][node]["name"] - for node in sorted(node_data[task]) - }, + "node_names": node_names[task], "results": {}, } fpath_results.parent.mkdir(parents=True, exist_ok=True) diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 67cee4a..ec65ae2 100644 --- a/tests/scripts/test_run_experiments.py +++ b/tests/scripts/test_run_experiments.py @@ -42,6 +42,7 @@ } for node in ["NODE_1", "NODE_2", "NODE_3"] } +NODE_NAMES = {node: f"node_name_{node}" for node in NODE_DATA} @pytest.fixture @@ -118,7 +119,9 @@ def mock_setup_held_out_secagg(module: ModuleType, mocker: MockerFixture) -> Mag @pytest.fixture def mock_get_node_data(module: ModuleType, mocker: MockerFixture) -> MagicMock: - return mocker.patch.object(module, "_get_node_data", return_value=NODE_DATA) + return mocker.patch.object( + module, "_get_node_data", return_value=(NODE_DATA, NODE_NAMES) + ) @pytest.fixture @@ -135,16 +138,46 @@ def fpath_csv(tmp_path: Path) -> Path: return fpath +def _mock_ping_replies(exp: MagicMock, node_names: dict[str, str]) -> None: + exp.researcher_id = "RESEARCHER_ID" + request = exp.requests.send.return_value.__enter__.return_value + request.replies.return_value = { + node: MagicMock(node_name=name) for node, name in node_names.items() + } + + +class TestGetNodeNames: + def test_get_node_names(self, module: ModuleType): + exp = MagicMock() + _mock_ping_replies(exp, NODE_NAMES) + assert module._get_node_names(exp, list(NODE_DATA)) == NODE_NAMES + message, nodes = exp.requests.send.call_args.args + assert message.researcher_id == "RESEARCHER_ID" + assert nodes == list(NODE_DATA) + + def test_missing_reply(self, module: ModuleType, caplog: pytest.LogCaptureFixture): + exp = MagicMock() + _mock_ping_replies( + exp, {node: NODE_NAMES[node] for node in ["NODE_1", "NODE_3"]} + ) + with caplog.at_level(logging.WARNING): + node_names = module._get_node_names(exp, list(NODE_DATA)) + assert node_names == {**NODE_NAMES, "NODE_2": None} + assert "Could not get the name of nodes: ['NODE_2']" in caplog.text + + class TestGetNodeData: def test_get_node_data(self, module: ModuleType, mock_experiment: MagicMock): exp = mock_experiment.return_value exp.training_data.return_value.data.return_value = NODE_DATA - assert module._get_node_data("age") == NODE_DATA + _mock_ping_replies(exp, NODE_NAMES) + assert module._get_node_data("age") == (NODE_DATA, NODE_NAMES) assert mock_experiment.call_args.kwargs["tags"] == ["brainflex", "age"] def test_no_nodes(self, module: ModuleType, mock_experiment: MagicMock): mock_experiment.return_value.training_data.return_value = None - assert module._get_node_data("age") == {} + assert module._get_node_data("age") == ({}, {}) + mock_experiment.return_value.requests.send.assert_not_called() class TestEvaluateHeldOut: @@ -937,9 +970,7 @@ def test_cv( fpath_results, json_data = _read_results(dpath_date / task) assert fpath_results.name.endswith("-secagg-cv.json") assert json_data["settings"]["cv"] is True - assert json_data["node_names"] == { - node: f"name_{node}" for node in NODE_DATA - } + assert json_data["node_names"] == NODE_NAMES assert list(json_data["results"]) == [ "NODE_2__NODE_3", "NODE_1__NODE_3", @@ -965,7 +996,7 @@ def test_no_cv( _, json_data = _read_results( tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age" ) - assert json_data["node_names"] == {node: f"name_{node}" for node in NODE_DATA} + assert json_data["node_names"] == NODE_NAMES @pytest.mark.parametrize( ("secagg", "n_nodes", "match"), @@ -984,7 +1015,10 @@ def test_cv_not_enough_nodes( n_nodes: int, match: str, ): - mock_get_node_data.return_value = dict(list(NODE_DATA.items())[:n_nodes]) + mock_get_node_data.return_value = ( + dict(list(NODE_DATA.items())[:n_nodes]), + NODE_NAMES, + ) with pytest.raises(ValueError, match=match): module.run_experiments(tmp_path, ["age"], secagg=secagg) mock_run_task.assert_not_called() @@ -996,7 +1030,10 @@ def test_cv_two_nodes_no_secagg( mock_get_node_data: MagicMock, tmp_path: Path, ): - mock_get_node_data.return_value = dict(list(NODE_DATA.items())[:2]) + mock_get_node_data.return_value = ( + dict(list(NODE_DATA.items())[:2]), + NODE_NAMES, + ) module.run_experiments(tmp_path, ["age"], secagg=False) assert mock_run_task.call_count == 2