SF-SKAT: keep parties up until all have finished

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 && \

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")
 
 

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:

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")

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,36 +248,53 @@ 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)
     sfskat_protocol.start_sfskat("1", "config_dir")
     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):