[BUG] Fix node names

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)

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