[ENH] Add random seed

scripts/_io_utils.py

@@ -27,6 +27,7 @@ class PathHelper:
     alpha: float
     secagg: bool
     cv: bool
+    random_seed: int
 
     def get_fpath_out(
         self,
@@ -49,6 +50,7 @@ def get_fpath_out(
                 f"{self.learning_rate}lr",
                 f"{self.test_ratio}test",
                 f"{self.alpha}alpha",
+                f"{self.random_seed}seed",
                 "secagg" if self.secagg else "no_secagg",
                 "cv" if self.cv else "no_cv",
             ]

scripts/run_experiments.py

@@ -44,6 +44,7 @@
 DEFAULT_TENSORBOARD = True
 DEFAULT_SECAGG = True
 DEFAULT_CV = True
+DEFAULT_RANDOM_SEED = 3791
 DEFAULT_DATE_TAG = date.today().strftime("%Y_%m_%d")
 
 
@@ -139,6 +140,15 @@ def _get_parser() -> argparse.ArgumentParser:
             f" evaluate on the held-out node. Default: {DEFAULT_CV}."
         ),
     )
+    parser.add_argument(
+        "--random-seed",
+        type=int,
+        default=DEFAULT_RANDOM_SEED,
+        help=(
+            "Random seed for the node data splits, batch shuffling and SGD."
+            f" Incremented at each round. Default: {DEFAULT_RANDOM_SEED}."
+        ),
+    )
     parser.add_argument(
         "--overwrite",
         action="store_true",
@@ -326,6 +336,7 @@ def _run_task(
     training_nodes: list[str] | None = None,
     node_data_held_out: dict[str, dict] | None = None,
     secagg: bool = True,
+    random_seed: int = DEFAULT_RANDOM_SEED,
 ) -> Generator[tuple[str, dict], None, None]:
     """Train a model for a task and save it in a task-specific output directory.
 
@@ -356,6 +367,7 @@ def _run_task(
             ),
             # resample the local train/validation split at each round
             "shuffle_testing_dataset": is_cv,
+            "random_seed": random_seed,
         },
         agg_optimizer=_get_agg_optimizer(aggregator),
         round_limit=n_rounds,
@@ -397,7 +409,7 @@ def _run_task(
     fold = {
         "settings": {
             "training_nodes": training_nodes,
-            "held_out_node": node_data_held_out,
+            "held_out_node": held_out_node,
             "excluded_columns": excluded_columns,
             "model_args": exp.model_args(),
             "training_args": exp.training_args(),
@@ -413,6 +425,11 @@ def _run_task(
 
     for round_ in range(1, n_rounds + 1):
         logger.info(f"Round {round_}/{n_rounds}")
+        # nodes reseed at the start of every round, so offset the seed to get a
+        # different (but deterministic) data split and batch order at each round
+        exp.set_training_args(
+            {**exp.training_args(), "random_seed": random_seed + round_}
+        )
         if is_cv:
             _evaluate_held_out(exp, node_data_held_out, held_out_secagg_arguments)
         exp.run_once()
@@ -453,6 +470,7 @@ def run_experiments(
     tensorboard: bool = True,
     secagg: bool = True,
     cv: bool = True,
+    random_seed: int = DEFAULT_RANDOM_SEED,
 ) -> None:
     """Train and save a global model for each task."""
     output_dir = output_dir.resolve()
@@ -463,10 +481,10 @@ def run_experiments(
     runs: dict[str, list[str | None]] = {}
     node_data: dict[str, dict[str, dict]] = {}
     for task in tasks:
+        node_data[task] = _get_node_data(task)
         if not cv:
             runs[task] = [None]
             continue
-        node_data[task] = _get_node_data(task)
         # secure aggregation needs at least 2 training nodes
         min_n_nodes = 3 if secagg else 2
         if len(node_data[task]) < min_n_nodes:
@@ -490,6 +508,7 @@ def run_experiments(
         alpha=alpha,
         secagg=secagg,
         cv=cv,
+        random_seed=random_seed,
     )
 
     # error out if one of the result files already exists
@@ -500,7 +519,7 @@ def run_experiments(
         )
         if fpath.exists() and any(
             "results" in fold
-            for fold in json.loads(fpath.read_text()).get("folds", {}).values()
+            for fold in json.loads(fpath.read_text()).get("results", {}).values()
         ):
             fpaths_results_existing.append(fpath)
 
@@ -518,7 +537,14 @@ def run_experiments(
             fpath_log = path_helper.get_fpath_out(
                 task=task, prefix=PREFIX_LOGS, extension=".log"
             )
-            json_data = {"settings": {"runtime_args": settings}, "folds": {}}
+            json_data = {
+                "settings": settings,
+                "node_names": {
+                    node: node_data[task][node]["name"]
+                    for node in sorted(node_data[task])
+                },
+                "results": {},
+            }
             fpath_results.parent.mkdir(parents=True, exist_ok=True)
             with log_to_file(fpath_log):
                 for held_out_node in held_out_nodes:
@@ -561,8 +587,9 @@ def run_experiments(
                         training_nodes=training_nodes,
                         node_data_held_out=node_data_held_out,
                         secagg=secagg,
+                        random_seed=random_seed,
                     ):
-                        json_data["folds"][fold_key] = fold
+                        json_data["results"][fold_key] = fold
                         save_json(fpath_results, json_data)
                     logger.info(f"Saved settings and results to {fpath_results}")
             logger.info(f"Saved log to {fpath_log}")
@@ -587,6 +614,7 @@ def main(argv: list[str] | None = None) -> None:
             tensorboard=args.tensorboard,
             secagg=args.secagg,
             cv=args.cv,
+            random_seed=args.random_seed,
         )
     except FileExistsError as e:
         sys.exit(f"Error: {e}")

tests/scripts/test_io_utils.py

@@ -28,6 +28,7 @@ def path_helper(module: ModuleType, tmp_path: Path) -> object:
         alpha=0.1,
         secagg=True,
         cv=True,
+        random_seed=3791,
     )
 
 
@@ -39,14 +40,14 @@ def test_get_fpath_out(self, path_helper: object, tmp_path: Path):
             / "age"
             / (
                 "results-scaffold-10rounds-50updates-32-0.001lr-0.2test-0.1alpha"
-                "-secagg-cv.json"
+                "-3791seed-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")
+        assert fpath.stem.endswith("-0.1alpha-3791seed-no_secagg-no_cv")
 
     def test_fedprox(self, path_helper: object):
         path_helper = dataclasses.replace(path_helper, aggregator="fedprox")
@@ -62,7 +63,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-secagg-cv-heldout_NODE_1.pkl")
+        assert fpath.name.endswith("-0.1alpha-3791seed-secagg-cv-heldout_NODE_1.pkl")
 
 
 class TestSerialize:

tests/scripts/test_run_experiments.py

@@ -464,6 +464,11 @@ def test_run_task(
         exp = mock_experiment.return_value
         assert exp.run_once.call_count == 3
         assert exp.save_training_replies.call_count == 3
+        assert training_args["random_seed"] == module.DEFAULT_RANDOM_SEED
+        # different seed at each round
+        assert [
+            call.args[0]["random_seed"] for call in exp.set_training_args.call_args_list
+        ] == [module.DEFAULT_RANDOM_SEED + round_ for round_ in (1, 2, 3)]
         exp.training_plan.return_value.export_model.assert_called_once_with(
             str(fpath_model)
         )
@@ -673,7 +678,7 @@ def test_held_out_node(
         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 fold["settings"]["held_out_node"] == "NODE_3"
         assert "node_metrics" in fold["results"]
         assert not any("data leakage" in r.getMessage() for r in caplog.records)
         assert "===== AGE: held_out_node=NODE_3 =====" in caplog.messages
@@ -686,8 +691,8 @@ def _read_results(dpath_task: Path) -> tuple[Path, dict]:
 
 
 RESULTS_WITH_FOLD = {
-    "settings": {"runtime_args": {}},
-    "folds": {"NODE_1__NODE_2": {"settings": {}, "results": {}}},
+    "settings": {},
+    "results": {"NODE_1__NODE_2": {"settings": {}, "results": {}}},
 }
 
 
@@ -713,12 +718,17 @@ def _write_results(
         alpha=module.DEFAULT_ALPHA,
         secagg=module.DEFAULT_SECAGG,
         cv=cv,
+        random_seed=module.DEFAULT_RANDOM_SEED,
     ).get_fpath_out(task, module.PREFIX_RESULTS)
     fpath.parent.mkdir(parents=True)
     fpath.write_text(json.dumps(data))
 
 
 class TestMain:
+    @pytest.fixture(autouse=True)
+    def _mock_get_node_data(self, mock_get_node_data: MagicMock) -> None:
+        """Avoid searching for node datasets."""
+
     def test_tasks_in_order(
         self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
     ):
@@ -752,13 +762,13 @@ def test_tasks_in_order(
         assert kwargs["dpath_tensorboard"].name.startswith(
             "tensorboard-scaffold-2rounds-"
         )
-        assert json_data["folds"] == {
+        assert json_data["results"] == {
             "NODE_1__NODE_2__NODE_3": {
                 "settings": {"training_nodes": None},
                 "results": {"node_metrics": {}},
             }
         }
-        settings = json_data["settings"]["runtime_args"]
+        settings = json_data["settings"]
         assert settings["output_dir"] == str(tmp_path.resolve())
         assert settings["tasks"] == ["sex", "age"]
         assert settings["date_tag"] == module.DEFAULT_DATE_TAG
@@ -798,7 +808,26 @@ def test_date_tag(
     ):
         module.main([str(tmp_path), "--task", "age", "--date-tag", "mytag", "--no-cv"])
         _, json_data = _read_results(tmp_path.resolve() / "mytag" / "age")
-        assert json_data["settings"]["runtime_args"]["date_tag"] == "mytag"
+        assert json_data["settings"]["date_tag"] == "mytag"
+
+    def test_random_seed(
+        self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
+    ):
+        module.main([str(tmp_path), "--task", "age", "--random-seed", "7", "--no-cv"])
+        kwargs = mock_run_task.call_args.kwargs
+        assert kwargs["random_seed"] == 7
+        assert "-7seed-" in kwargs["fpath_model"].name
+        fpath_results, json_data = _read_results(
+            tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age"
+        )
+        assert "-7seed-" in fpath_results.name
+        assert json_data["settings"]["random_seed"] == 7
+
+    def test_default_random_seed(
+        self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
+    ):
+        module.main([str(tmp_path), "--task", "age", "--no-cv"])
+        assert mock_run_task.call_args.kwargs["random_seed"] == 3791
 
     def test_no_secagg(
         self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
@@ -808,7 +837,7 @@ def test_no_secagg(
         _, json_data = _read_results(
             tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age"
         )
-        assert json_data["settings"]["runtime_args"]["secagg"] is False
+        assert json_data["settings"]["secagg"] is False
 
     def test_aggregator_options(
         self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
@@ -845,7 +874,7 @@ def test_no_tensorboard(
         _, json_data = _read_results(
             tmp_path.resolve() / module.DEFAULT_DATE_TAG / "age"
         )
-        assert json_data["settings"]["runtime_args"]["tensorboard"] is False
+        assert json_data["settings"]["tensorboard"] is False
 
     def test_existing_results(
         self, module: ModuleType, mock_run_task: MagicMock, tmp_path: Path
@@ -863,7 +892,7 @@ def test_existing_settings_only(
             module,
             tmp_path.resolve(),
             "age",
-            {"settings": {}, "folds": {"NODE_1__NODE_2": {"settings": {}}}},
+            {"settings": {}, "results": {"NODE_1__NODE_2": {"settings": {}}}},
         )
         module.main([str(tmp_path), "--task", "age", "--no-cv"])
         mock_run_task.assert_called_once()
@@ -907,13 +936,16 @@ def test_cv(
         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"]) == [
+            assert json_data["settings"]["cv"] is True
+            assert json_data["node_names"] == {
+                node: f"name_{node}" for node in NODE_DATA
+            }
+            assert list(json_data["results"]) == [
                 "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 all("results" in fold for fold in json_data["results"].values())
             assert [p.name for p in (dpath_date / task).glob("logs-*")] == [
                 fpath_results.name.replace("results-", "logs-").replace(".json", ".log")
             ]
@@ -926,10 +958,14 @@ def test_no_cv(
         tmp_path: Path,
     ):
         module.main([str(tmp_path), "--task", "age", "--no-cv"])
-        mock_get_node_data.assert_not_called()
+        mock_get_node_data.assert_called_once_with("age")
         kwargs = mock_run_task.call_args.kwargs
         assert kwargs["training_nodes"] is None
         assert kwargs["node_data_held_out"] is None
+        _, 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}
 
     @pytest.mark.parametrize(
         ("secagg", "n_nodes", "match"),