Testing PR #636 variant: streaming grandchild processes with os.fork model in self-test.py

source/src/python/PyRosetta/src/pyrosetta/__init__.py

@@ -190,7 +190,7 @@ def init(options='-ex1 -ex2aro', extra_options='', set_logging_handler=None, not
     v.extend(args)
 
     if not silent:
-        print( version() )
+        print( version() , flush=True)
         logger.info( "\n" + version() )
     else:
         logger.debug( "\n" + version() )

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/__init__.py

@@ -50,7 +50,7 @@
     "run",
     "update_scores",
 ]
-__version__: str = "4.1.0"
+__version__: str = "4.1.1"
 
 
 with warnings.catch_warnings() and suppress(NameError):

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/converter_tasks.py

@@ -473,7 +473,8 @@ def export_init_file(
                         f.write(init_file_json)
                 print(
                     f"Exported PyRosettaCluster decoy output file '{output_file}' to "
-                    + f"PyRosetta initialization file: '{output_init_file}'"
+                    + f"PyRosetta initialization file: '{output_init_file}'",
+                    flush=True,
                 )
             else:
                 raise ValueError(

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/serialization.py

@@ -221,50 +221,44 @@ def _cache_nonce(self, sealed: bytes) -> None:
         """
         package = self.unpack(sealed)
         if not isinstance(package, dict) or package.get("v", None) != 1:
+            _location = "Dask worker" if NonceCache._on_worker() else "head node"
             _err_msg = (
-                "Invalid sealed package or version on {0} nonce cache! "
+                f"Invalid sealed package or version on {_location} nonce cache! "
                 + f"Received: {type(package)!r}"
             )
-            if NonceCache._on_worker():
-                raise ValueError(_err_msg.format("worker"))
-            else:
-                raise ValueError(_err_msg.format("host"))
+            raise ValueError(_err_msg)
 
         _instance_id = package["a"] # `str`: PyRosettaCluster instance identifier/App
         _data = package["d"] # `bytes`: Data bytestring
         _mac = package["m"] # `bytes`: MAC (HMAC tag)
         _nonce = package["n"] # `bytes`: Nonce
         _version = package["v"] # `int`: Version
-
         if _instance_id != self.instance_id:
             raise ValueError("PyRosettaCluster instance identifier mismatch in sealed package.")
 
         msg = self.pack([_instance_id, _data, _nonce, _version])
         _expected_mac = hmac_digest(bytes(self.prk), msg)
         if not compare_digest(_expected_mac, _mac):
+            _location = "Dask worker" if NonceCache._on_worker() else "head node"
             _err_msg = (
-                "Task HMAC verification failed during nonce cache on {0}!\n"
+                f"Task HMAC verification failed during nonce cache on {_location}!\n"
                 + f"Expected: {_expected_mac!r}\n"
                 + f"Value:    {_mac!r}\n"
             )
-            if NonceCache._on_worker():
-                raise SystemExit(_err_msg.format("worker"))
-            else:
-                raise SystemExit(_err_msg.format("host"))
+            raise SystemExit(_err_msg)
 
         if _nonce is not None:
             if _nonce in self._seen:
                 # Replay protection
+                _location = "Dask worker" if NonceCache._on_worker() else "head node"
                 _err_msg = (
-                    "PyRosettaCluster detected a repeat nonce on the {0} for the instance identifier "
+                    f"PyRosettaCluster detected a repeat nonce on the {_location} for the instance identifier "
                     + f"'{self.instance_id}', which might indicate a replay attack is in progress! "
                     + "Exiting process for security. Please ensure that `PyRosettaCluster(security=True)` "
                     + f"is enabled in future PyRosettaCluster simulations. Received: '{_nonce}'."
                 )
-                if NonceCache._on_worker():
-                    raise SystemExit(_err_msg.format("worker"))
-                else:
-                    raise SystemExit(_err_msg.format("host"))
+                raise SystemExit(_err_msg)
+
             self._seen.add(_nonce)
             self._order.append(_nonce)
             while len(self._seen) > self._order.maxlen:
@@ -279,9 +273,9 @@ def _cache_nonce(self, sealed: bytes) -> None:
                     + f"Example: {sorted(self._seen)[0]}"
                 )
                 if NonceCache._on_worker():
-                    print(f"Remote worker ({get_worker().contact_address}) nonce cache: {_msg}")
+                    print(f"Remote Dask worker ({get_worker().contact_address}) nonce cache: {_msg}")
                 else:
-                    print(f"Local host nonce cache: {_msg}")
+                    print(f"Local head node nonce cache: {_msg}")
 
 
 @attr.s(kw_only=True, slots=False, frozen=False)

source/src/python/PyRosetta/src/pyrosetta/distributed/cluster/tools.py

@@ -574,73 +574,76 @@ def reproduce(
         )
 
     _tmp_dir = None
-    if isinstance(input_file, str):
-        if input_file.endswith((".init", ".init.bz2")):
-            _tmp_dir = tempfile.TemporaryDirectory(prefix="PyRosettaCluster_reproduce_")
-            default_init_from_file_kwargs = dict(
-                output_dir=os.path.join(_tmp_dir.name, "pyrosetta_init_input_files"),
-                skip_corrections=skip_corrections,
-                relative_paths=True,
-                dry_run=False,
-                max_decompressed_bytes=pow(2, 30), # 1 GiB
-                restore_rg_state=True,
-                database=None,
-                verbose=True,
-                set_logging_handler="logging",
-                notebook=None,
-                silent=False,
-            )
-            input_packed_pose, input_file = parse_init_file(
-                input_file,
-                input_packed_pose,
-                skip_corrections,
-                toolz.dicttoolz.merge(
-                    default_init_from_file_kwargs,
-                    init_from_file_kwargs if is_dict(init_from_file_kwargs) else {},
-                ),
-            )
-        elif input_file.endswith((".pkl_pose", ".pkl_pose.bz2", ".b64_pose", ".b64_pose.bz2")):
-            if not was_init_called():
-                raise PyRosettaIsNotInitializedError(
-                    "If providing a '.pkl_pose', '.pkl_pose.bz2', '.b64_pose', or '.b64_pose.bz2' file to the 'input_file' "
-                    + "keyword argument parameter, please ensure `pyrosetta.init()` or `pyrosetta.init_from_file()` has been "
-                    + "properly called (with the same residue type set as that used to generate the original '.pkl_pose', "
-                    + "'.pkl_pose.bz2', '.b64_pose', or '.b64_pose.bz2' file) before running the `reproduce()` function. "
-                    + "If an output '.init' file from the original simulation is available, it is recommended to run "
-                    + "`pyrosetta.init_from_file()` with that '.init' file before running the `reproduce()` function."
+    try:
+        if isinstance(input_file, str):
+            if input_file.endswith((".init", ".init.bz2")):
+                _tmp_dir = tempfile.TemporaryDirectory(prefix="PyRosettaCluster_reproduce_")
+                default_init_from_file_kwargs = dict(
+                    output_dir=os.path.join(_tmp_dir.name, "pyrosetta_init_input_files"),
+                    skip_corrections=skip_corrections,
+                    relative_paths=True,
+                    dry_run=False,
+                    max_decompressed_bytes=pow(2, 30), # 1 GiB
+                    restore_rg_state=True,
+                    database=None,
+                    verbose=True,
+                    set_logging_handler="logging",
+                    notebook=None,
+                    silent=False,
+                )
+                input_packed_pose, input_file = parse_init_file(
+                    input_file,
+                    input_packed_pose,
+                    skip_corrections,
+                    toolz.dicttoolz.merge(
+                        default_init_from_file_kwargs,
+                        init_from_file_kwargs if is_dict(init_from_file_kwargs) else {},
+                    ),
                 )
+            elif input_file.endswith((".pkl_pose", ".pkl_pose.bz2", ".b64_pose", ".b64_pose.bz2")):
+                if not was_init_called():
+                    raise PyRosettaIsNotInitializedError(
+                        "If providing a '.pkl_pose', '.pkl_pose.bz2', '.b64_pose', or '.b64_pose.bz2' file to the 'input_file' "
+                        + "keyword argument parameter, please ensure `pyrosetta.init()` or `pyrosetta.init_from_file()` has been "
+                        + "properly called (with the same residue type set as that used to generate the original '.pkl_pose', "
+                        + "'.pkl_pose.bz2', '.b64_pose', or '.b64_pose.bz2' file) before running the `reproduce()` function. "
+                        + "If an output '.init' file from the original simulation is available, it is recommended to run "
+                        + "`pyrosetta.init_from_file()` with that '.init' file before running the `reproduce()` function."
+                    )
 
-    PyRosettaCluster(
-        **toolz.dicttoolz.keyfilter(
-            lambda a: a not in ["client", "clients", "input_packed_pose"],
-            toolz.dicttoolz.merge(
-                get_instance_kwargs(
-                    input_file=input_file,
-                    scorefile=scorefile,
-                    decoy_name=decoy_name,
-                    skip_corrections=skip_corrections,
-                    with_metadata_kwargs=False,
+        PyRosettaCluster(
+            **toolz.dicttoolz.keyfilter(
+                lambda a: a not in ["client", "clients", "input_packed_pose"],
+                toolz.dicttoolz.merge(
+                    get_instance_kwargs(
+                        input_file=input_file,
+                        scorefile=scorefile,
+                        decoy_name=decoy_name,
+                        skip_corrections=skip_corrections,
+                        with_metadata_kwargs=False,
+                    ),
+                    parse_instance_kwargs(instance_kwargs),
                 ),
-                parse_instance_kwargs(instance_kwargs),
             ),
-        ),
-        client=parse_client(client),
-        clients=clients,
-        input_packed_pose=input_packed_pose,
-    ).distribute(
-        protocols=get_protocols(
-            protocols=protocols,
-            input_file=input_file,
-            scorefile=scorefile,
-            decoy_name=decoy_name,
-        ),
-        clients_indices=clients_indices,
-        resources=resources,
-        priorities=None,
-        retries=retries,
-    )
-    if isinstance(_tmp_dir, tempfile.TemporaryDirectory):
-        _tmp_dir.cleanup()
+            client=parse_client(client),
+            clients=clients,
+            input_packed_pose=input_packed_pose,
+        ).distribute(
+            protocols=get_protocols(
+                protocols=protocols,
+                input_file=input_file,
+                scorefile=scorefile,
+                decoy_name=decoy_name,
+            ),
+            clients_indices=clients_indices,
+            resources=resources,
+            priorities=None,
+            retries=retries,
+        )
+
+    finally:
+        if _tmp_dir:
+            _tmp_dir.cleanup()
 
 
 def produce(**kwargs: Any) -> Optional[NoReturn]:

source/src/python/PyRosetta/src/pyrosetta/tests/distributed/cluster/print_discovered_test_cases.py

@@ -10,6 +10,7 @@
 
 import argparse
 import os
+import sys
 import unittest
 
 def flatten(suite):
@@ -19,7 +20,8 @@ def flatten(suite):
 def main(root):
     suite = unittest.defaultTestLoader.discover(root)
     test_cases = (test_case.id() for test_case in flatten(suite))
-    print(*test_cases, sep="\n")
+    print(*test_cases, sep="\n", flush=True)
+    sys.stdout.flush()
 
 # if __name__ == "__main__":
 #     parser = argparse.ArgumentParser()