diff --git a/Dockerfile b/Dockerfile index d1b074e..488d20d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -65,7 +65,7 @@ RUN git clone https://github.com/froelich/sf-relate . && \ FROM go AS sf-skat RUN git clone https://github.com/swanhong/secure-skat . && \ - git checkout 1d94279 && \ + git checkout e2b5186 && \ # dial peers via the SOCKS proxy from the environment (as sfgwas does), # which sfkit-proxy relies on; fail if the patch no longer applies go get golang.org/x/net@v0.20.0 && \ diff --git a/sfkit/utils/sfgwas_protocol.py b/sfkit/utils/sfgwas_protocol.py index d63bc69..7cbce7e 100644 --- a/sfkit/utils/sfgwas_protocol.py +++ b/sfkit/utils/sfgwas_protocol.py @@ -363,7 +363,7 @@ def build_sfgwas() -> None: print("Finished building sfgwas code") -def sync_with_other_vms(role: str, demo: bool, skip_cp0: bool = False) -> None: +def sync_with_other_vms(role: str, demo: bool, skip_cp0: bool = False, stagger: int = 15) -> None: update_firestore("update_firestore::status=syncing up") update_firestore("update_firestore::task=Syncing up machines") print("Begin syncing up") @@ -381,7 +381,7 @@ def sync_with_other_vms(role: str, demo: bool, skip_cp0: bool = False) -> None: break print("Waiting for all participants to sync up...") time.sleep(5) - time.sleep(15 * int(role)) + time.sleep(stagger * int(role)) print("Finished syncing up") diff --git a/sfkit/utils/sfskat_protocol.py b/sfkit/utils/sfskat_protocol.py index c4ec0cf..e8a29c1 100644 --- a/sfkit/utils/sfskat_protocol.py +++ b/sfkit/utils/sfskat_protocol.py @@ -30,6 +30,11 @@ from sfkit.utils.sfgwas_protocol import sync_with_other_vms GLOBAL_CONFIG_FILENAME = "configGlobal.toml" +LOCAL_CONFIG_FILENAME = "configLocal.toml" +# Set once this party's secure-rvas has exited. It must not contain "Finished protocol", +# which the website takes as the cue to stop the auxiliary party's VM. +COMPUTATION_DONE_STATUS = "finishing SF-SKAT protocol" +FINISHED_STATUS = "Finished protocol!" def run_sfskat_protocol(role: str, demo: bool = False) -> None: @@ -46,7 +51,11 @@ def run_sfskat_protocol(role: str, demo: bool = False) -> None: print("Begin updating config files") config_dir = update_config(role) - sync_with_other_vms(role, demo) + # Start all parties together instead of staggering them by role: sfkit-proxy sets up its + # NAT bindings as soon as it boots, and they expire (after 30s on Google Cloud NAT) while + # an early party waits for the later ones. secure-rvas retries its connections, so the + # parties do not need to come up in order. + sync_with_other_vms(role, demo, stagger=0) start_sfskat(role, config_dir) @@ -169,11 +178,11 @@ def generate_shared_keys(role: int, ancestries: list) -> None: def update_config(role: str) -> str: """ Write the config directory that secure-rvas's ``party`` subcommand expects: - configGlobal.toml (study parameters and networking, the same for every party - apart from run_dir) and configLocal.Party{role}.toml (this party's keys, - thread count and prepared input paths). Settings that only matter to - ``secure-rvas prepare``, which sfkit does not invoke (participants must - already have run it out-of-band), are left out. + configGlobal.toml (study parameters and networking, the same for every party) + and configLocal.toml (this party's run directory, keys, thread count and + prepared input paths). Settings that only matter to ``secure-rvas prepare``, + which sfkit does not invoke (participants must already have run it + out-of-band), are left out. """ doc_ref_dict: dict = get_doc_ref_dict() pars = {**doc_ref_dict["parameters"], **doc_ref_dict["advanced_parameters"]} @@ -185,7 +194,6 @@ def scalar(key: str, default): run_dir = get_run_dir(role) global_config = { - "run_dir": run_dir, "chromosomes": [int(c) for c in parse_list_param(pars.get("chromosomes", {}).get("value"), ["21", "22"])], "phenotype_columns": parse_list_param(pars.get("phenotype_columns", {}).get("value"), ["phenotype1"]), "ancestries": parse_ancestries(pars.get("ancestries", {}).get("value")), @@ -214,6 +222,7 @@ def scalar(key: str, default): global_config["servers"][f"party{i}"] = server local_config = { + "run_dir": run_dir, "shared_keys_path": shared_keys_root(), "local_num_threads": os.cpu_count() or 1, } @@ -233,7 +242,7 @@ def scalar(key: str, default): os.makedirs(config_dir, exist_ok=True) for filename, data in ( (GLOBAL_CONFIG_FILENAME, global_config), - (f"configLocal.Party{role}.toml", local_config), + (LOCAL_CONFIG_FILENAME, local_config), ): with open(os.path.join(config_dir, filename), "w") as f: f.write(tomlkit.dumps(data)) @@ -266,15 +275,35 @@ def start_sfskat(role: str, config_dir: str) -> None: role=role, ) + # secure-rvas exits as soon as it has sent its last message, without waiting for the + # other parties to receive it; the auxiliary party in particular is done within seconds + # of connecting. So keep sfkit-proxy up, and hold off reporting completion (which makes + # the website stop the auxiliary party's VM), until every party's secure-rvas has exited. + wait_for_other_parties() + if sfkit_proxy: sfkit_proxy.terminate() print("\n\n Finished SF-SKAT \n\n") - if role == "1": + # both cohorts receive the results; the auxiliary party has none + if role != "0": process_output_files(role) - update_firestore("update_firestore::status=Finished protocol!") + update_firestore(f"update_firestore::status={FINISHED_STATUS}") + + +def wait_for_other_parties() -> None: + update_firestore(f"update_firestore::status={COMPUTATION_DONE_STATUS}") + while True: + statuses = get_doc_ref_dict()["status"].values() + if all( + status in (COMPUTATION_DONE_STATUS, FINISHED_STATUS) or "FAILED" in status + for status in statuses + ): + break + print("Waiting for the other parties to finish...") + time.sleep(5) def process_output_files(role: str) -> None: diff --git a/tests/utils/test_sfgwas_protocol.py b/tests/utils/test_sfgwas_protocol.py index ba4361e..79d5746 100644 --- a/tests/utils/test_sfgwas_protocol.py +++ b/tests/utils/test_sfgwas_protocol.py @@ -117,6 +117,26 @@ def test_build_sfgwas(mocker): sfgwas_protocol.build_sfgwas() +def test_sync_with_other_vms(mocker): + mocker.patch("sfkit.utils.sfgwas_protocol.update_firestore") + mocker.patch( + "sfkit.utils.sfgwas_protocol.get_doc_ref_dict", + side_effect=lambda: copy.deepcopy(mock_doc_ref_dict_syncing_up), + ) + sleep = mocker.patch("sfkit.utils.sfgwas_protocol.time.sleep") + + sfgwas_protocol.sync_with_other_vms("2", True) + sleep.assert_not_called() + + # by default, each role starts 15 seconds after the previous one + sfgwas_protocol.sync_with_other_vms("2", False) + sleep.assert_called_once_with(30) + + sleep.reset_mock() + sfgwas_protocol.sync_with_other_vms("2", False, skip_cp0=True, stagger=0) + sleep.assert_called_once_with(0) + + def test_start_sfgwas(mocker): mocker.patch("sfkit.utils.sfgwas_protocol.boot_sfkit_proxy") mocker.patch("sfkit.utils.sfgwas_protocol.update_firestore") diff --git a/tests/utils/test_sfskat_protocol.py b/tests/utils/test_sfskat_protocol.py index 954c991..b5f2bfd 100644 --- a/tests/utils/test_sfskat_protocol.py +++ b/tests/utils/test_sfskat_protocol.py @@ -66,6 +66,8 @@ def test_run_sfskat_protocol(mocker): sfskat_protocol.run_sfskat_protocol("1") sfskat_protocol.generate_shared_keys.assert_called_once_with(1, ["EUR", "AFR"]) + # the parties must not be staggered, or the early ones' NAT bindings go stale + sfskat_protocol.sync_with_other_vms.assert_called_with("1", False, stagger=0) sfskat_protocol.start_sfskat.assert_called_with("1", "config_dir") mocker.patch("sfkit.utils.sfskat_protocol.constants.IS_DOCKER", True) @@ -193,7 +195,6 @@ def test_update_config(mocker, tmp_path): config_dir = sfskat_protocol.update_config("1") assert config_dir == str(tmp_path / "skat_config") assert load_toml(os.path.join(config_dir, "configGlobal.toml")) == { - "run_dir": str(data_path), "chromosomes": [21, 22], "phenotype_columns": ["phenotype1", "phenotype2"], "ancestries": ["EUR", "AFR"], @@ -220,7 +221,8 @@ def test_update_config(mocker, tmp_path): }, }, } - assert load_toml(os.path.join(config_dir, "configLocal.Party1.toml")) == { + assert load_toml(os.path.join(config_dir, "configLocal.toml")) == { + "run_dir": str(data_path), "shared_keys_path": str(tmp_path / "skat_shared_keys"), "local_num_threads": 8, "genotype_dir": f"{chr_dir}/A/geno", @@ -231,7 +233,8 @@ def test_update_config(mocker, tmp_path): } sfskat_protocol.update_config("2") - assert load_toml(os.path.join(config_dir, "configLocal.Party2.toml")) == { + assert load_toml(os.path.join(config_dir, "configLocal.toml")) == { + "run_dir": str(data_path), "shared_keys_path": str(tmp_path / "skat_shared_keys"), "local_num_threads": 8, "genotype_dir": f"{chr_dir}/B/geno", @@ -245,22 +248,27 @@ def test_update_config(mocker, tmp_path): # the auxiliary party has no data, and the proxy makes every peer local mocker.patch("sfkit.utils.sfskat_protocol.constants.SFKIT_PROXY_ON", True) sfskat_protocol.update_config("0") - assert load_toml(os.path.join(config_dir, "configLocal.Party0.toml")) == { + assert load_toml(os.path.join(config_dir, "configLocal.toml")) == { + "run_dir": str(tmp_path / "skat_run"), "shared_keys_path": str(tmp_path / "skat_shared_keys"), "local_num_threads": 8, } global_config = load_toml(os.path.join(config_dir, "configGlobal.toml")) - assert global_config["run_dir"] == str(tmp_path / "skat_run") assert {server["ipaddr"] for server in global_config["servers"].values()} == { "127.0.0.1" } def test_start_sfskat(mocker): - mocker.patch("sfkit.utils.sfskat_protocol.update_firestore") - mocker.patch("sfkit.utils.sfskat_protocol.boot_sfkit_proxy") - mocker.patch("sfkit.utils.sfskat_protocol.run_command") - mocker.patch("sfkit.utils.sfskat_protocol.process_output_files") + calls = mocker.Mock() + for name in [ + "update_firestore", + "boot_sfkit_proxy", + "run_command", + "wait_for_other_parties", + "process_output_files", + ]: + calls.attach_mock(mocker.patch(f"sfkit.utils.sfskat_protocol.{name}"), name) mocker.patch("sfkit.utils.sfskat_protocol.constants.IS_DOCKER", True) mocker.patch("sfkit.utils.sfskat_protocol.constants.SFKIT_PROXY_ON", True) @@ -268,13 +276,25 @@ def test_start_sfskat(mocker): sfskat_protocol.boot_sfkit_proxy.assert_called_once_with( "1", "config_dir/configGlobal.toml" ) - sfskat_protocol.boot_sfkit_proxy.return_value.terminate.assert_called_once() sfskat_protocol.run_command.assert_called_once_with( ["secure-rvas", "party", "--config", "config_dir", "--party", "1"], fail_message="Failed SF-SKAT protocol", role="1", ) sfskat_protocol.process_output_files.assert_called_once_with("1") + # the proxy stays up, and completion goes unreported, until the other parties are done + assert [call[0] for call in calls.mock_calls if "__" not in call[0]] == [ + "update_firestore", + "boot_sfkit_proxy", + "run_command", + "wait_for_other_parties", + "boot_sfkit_proxy().terminate", + "process_output_files", + "update_firestore", + ] + sfskat_protocol.update_firestore.assert_called_with( + "update_firestore::status=Finished protocol!" + ) mocker.patch("sfkit.utils.sfskat_protocol.constants.SFKIT_PROXY_ON", False) mocker.patch("sfkit.utils.sfskat_protocol.constants.IS_DOCKER", False) @@ -284,7 +304,43 @@ def test_start_sfskat(mocker): "sf-skat", "secure-rvas" ) sfskat_protocol.boot_sfkit_proxy.assert_called_once() - sfskat_protocol.process_output_files.assert_called_once() + # both cohorts receive results, the auxiliary party does not + sfskat_protocol.process_output_files.assert_called_with("2") + sfskat_protocol.start_sfskat("0", "config_dir") + assert sfskat_protocol.process_output_files.call_count == 2 + + +def test_wait_for_other_parties(mocker): + update_firestore = mocker.patch("sfkit.utils.sfskat_protocol.update_firestore") + sleep = mocker.patch("sfkit.utils.sfskat_protocol.time.sleep") + done = sfskat_protocol.COMPUTATION_DONE_STATUS + # the website stops the auxiliary party's VM on any status containing this + assert "Finished protocol" not in done + + def statuses(*values): + return {"status": dict(zip(PARTICIPANTS, values))} + + # the auxiliary party is done first, and has to outlast both cohorts + mocker.patch( + "sfkit.utils.sfskat_protocol.get_doc_ref_dict", + side_effect=[ + statuses(done, "running SF-SKAT protocol", "running SF-SKAT protocol"), + statuses(done, done, "running SF-SKAT protocol"), + statuses(done, "Finished protocol!", done), + ], + ) + sfskat_protocol.wait_for_other_parties() + update_firestore.assert_called_once_with(f"update_firestore::status={done}") + assert sleep.call_count == 2 + + # a failed party will never finish, so it must not be waited for + sleep.reset_mock() + mocker.patch( + "sfkit.utils.sfskat_protocol.get_doc_ref_dict", + return_value=statuses(done, "FAILED - Failed SF-SKAT protocol", done), + ) + sfskat_protocol.wait_for_other_parties() + sleep.assert_not_called() def test_process_output_files(mocker, tmp_path):