diff --git a/scripts/_io_utils.py b/scripts/_io_utils.py index cc06f49..a92a5fc 100644 --- a/scripts/_io_utils.py +++ b/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", ] diff --git a/scripts/run_experiments.py b/scripts/run_experiments.py index cae84c5..3fd6194 100755 --- a/scripts/run_experiments.py +++ b/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}") diff --git a/tests/scripts/test_io_utils.py b/tests/scripts/test_io_utils.py index 375bc9e..10a1fb5 100644 --- a/tests/scripts/test_io_utils.py +++ b/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: diff --git a/tests/scripts/test_run_experiments.py b/tests/scripts/test_run_experiments.py index 5764d5c..67cee4a 100644 --- a/tests/scripts/test_run_experiments.py +++ b/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"),