Add support for filtering results in PyRosettaCluster (#459)

Author: klimaj

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

@@ -23,6 +23,7 @@
     produce,
     recreate_environment,
     reproduce,
+    requires_packed_pose,
     reserve_scores,
     run,
     update_scores,
@@ -43,11 +44,12 @@
     "produce",
     "recreate_environment",
     "reproduce",
+    "requires_packed_pose",
     "reserve_scores",
     "run",
     "update_scores",
 ]
-__version__: str = "2.0.0"
+__version__: str = "2.1.0"
 
 _print_conda_warnings()
 

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

@@ -422,9 +422,17 @@ def _parse_dict(obj: Dict[Any, Any]) -> Dict[Any, Any]:
     for k in obj.keys():
         if k in ["client", "clients", "input_packed_pose"]:
             raise NotImplementedError(
-                f"The parameter '{k}' must be passed directly to reproduce(), "
+                f"The parameter '{k}' must be passed directly to `reproduce()`, "
                 + "not as a member of the 'instance_kwargs' dictionary."
             )
+        elif k == "filter_results":
+            raise ValueError(
+                f"The parameter '{k}' cannot be set as a PyRosettaCluster attribute "
+                + "in `reproduce()` because the saved 'decoy_ids' attribute from the "
+                + "original simulation depends on the original decoy output order from "
+                + "each protocol, so results must be filtered identically. Please remove "
+                + "this keyword argument to run `reproduce()`."
+            )
     return obj
 
 
@@ -433,6 +441,21 @@ def _default_none(obj: None) -> Dict[Any, Any]:
     return {}
 
 
+@singledispatch
+def is_empty(obj: Any) -> NoReturn:
+    """Test whether a `PackedPose` object is empty."""
+    raise NotImplementedError(type(obj))
+
+@is_empty.register(type(None))
+def _from_none(obj: None) -> bool:
+    # Protocol results return a `None` object when a segmentation fault occurs with `ignore_errors=True`
+    return False
+
+@is_empty.register(PackedPose)
+def _from_packed(obj: PackedPose) -> bool:
+    return obj.empty()
+
+
 def is_bytes(obj: Any) -> bool:
     return isinstance(obj, bytes)
 

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

@@ -30,13 +30,15 @@
 import pyrosetta
 import sys
 import types
+import warnings
 
 from functools import singledispatch
 from pyrosetta.distributed.cluster.converter_tasks import (
     environment_cmd,
     get_yml,
     is_bytes,
     is_dict,
+    is_empty,
     is_packed,
     parse_input_packed_pose as _parse_input_packed_pose,
     to_int,
@@ -63,6 +65,38 @@
 S = TypeVar("S", bound=Serialization)
 
 
+def _parse_filter_results(obj: Any) -> Union[bool, NoReturn]:
+    """Parse the input `filter_results` attribute of PyRosettaCluster."""
+    _issue_future_warning = True
+
+    @singledispatch
+    def converter(obj: Any) -> NoReturn:
+        raise ValueError("'filter_results' must be of type `bool` or `NoneType`!")
+
+    @converter.register(type(None))
+    def _parse_none(obj: None) -> bool:
+        if _issue_future_warning:
+            warnings.warn(
+                (
+                    "As of PyRosettaCluster version 2.1.0, the 'filter_results' instance attribute "
+                    "is enabled by default, which automatically filters empty `PackedPose` objects between "
+                    "user-provided PyRosetta protocols to help reduce compute overhead. Please explicitly set "
+                    "either `filter_results=True` (the currently enabled, new setting) or `filter_results=False` "
+                    "(to revert to legacy behavior before version 2.1.0) to silence this notice. This notice "
+                    "will disappear in a future version of PyRosettaCluster."
+                ),
+                FutureWarning,
+                stacklevel=5,
+            )
+        return True
+
+    @converter.register(bool)
+    def _parse_bool(obj: bool) -> bool:
+        return obj
+
+    return converter(obj)
+
+
 def _parse_decoy_ids(objs: Any) -> List[int]:
     """
     Normalize user-provided PyRosetta 'decoy_ids' to a `list` object containing `int` objects.

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

@@ -206,6 +206,16 @@
         It's recommended to set this option to at least 1 second, but longer times may
         be used as a safety throttle in cases of overwhelmed dask scheduler processes.
         Default: 3.0
+    filter_results: A `bool` object specifying whether or not to filter out empty
+        `PackedPose` objects between user-provided PyRosetta protocols. When a protocol
+        returns or yields `NoneType`, PyRosettaCluster converts it to an empty `PackedPose`
+        object that gets passed to the next protocol. If `True`, then filter out any empty
+        `PackedPose` objects where there are no residues in the conformation as given by
+        `Pose.empty()`, otherwise if `False` then continue to pass empty `PackedPose` objects
+        to the next protocol. This is used for filtering out decoys mid-trajectory through
+        user-provided PyRosetta protocols if protocols return or yield any `None`, empty
+        `Pose`, or empty `PackedPose` objects.
+        Default: True
     save_all: A `bool` object specifying whether or not to save all of the returned
         or yielded `Pose` and `PackedPose` objects from all user-provided
         PyRosetta protocols. This option may be used for checkpointing trajectories.
@@ -257,7 +267,9 @@
 from datetime import datetime
 from pyrosetta.distributed.cluster.base import TaskBase, _get_residue_type_set
 from pyrosetta.distributed.cluster.converters import (
+    is_empty as _is_empty,
     _parse_decoy_ids,
+    _parse_filter_results,
     _parse_environment,
     _parse_input_packed_pose,
     _parse_logging_address,
@@ -567,6 +579,12 @@ class PyRosettaCluster(IO[G], LoggingSupport[G], SchedulerManager[G], TaskBase[G
         validator=[_validate_float, attr.validators.instance_of((float, int))],
         converter=attr.converters.default_if_none(default=3.0),
     )
+    filter_results = attr.ib(
+        type=bool,
+        default=None,
+        validator=attr.validators.instance_of(bool),
+        converter=_parse_filter_results,
+    )
     save_all = attr.ib(
         type=bool,
         default=False,
@@ -822,6 +840,8 @@ def _run(
                             compressed_packed_pose,
                             self.serializer.deepcopy_kwargs(kwargs),
                         )
+                    if self.filter_results and _is_empty(self.serializer.decompress_packed_pose(compressed_packed_pose)):
+                        continue
                     compressed_kwargs, pyrosetta_init_kwargs, protocol, clients_index, resource = self._setup_kwargs(
                         kwargs, clients_indices, resources
                     )

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

@@ -133,6 +133,8 @@ def wrapper(
                 + WorkerError._get_msg(protocol_name, ignore_errors)
             )
             if ignore_errors:
+                # Return a `NoneType` object to be converted to an empty `PackedPose` object
+                # when a non-system-exiting Python exception is raised
                 result = None
             else:
                 raise

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

@@ -21,6 +21,7 @@
     produce,
     recreate_environment,
     reproduce,
+    requires_packed_pose,
     reserve_scores,
     run,
     update_scores,

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

@@ -34,6 +34,7 @@
 from functools import wraps
 from pyrosetta.distributed.cluster.converters import _parse_protocols, _parse_yield_results
 from pyrosetta.distributed.cluster.converter_tasks import (
+    is_empty,
     get_protocols_list_of_str,
     get_yml,
     parse_client,
@@ -427,6 +428,53 @@ def wrapper(packed_pose, **kwargs):
     return cast(P, wrapper)
 
 
+def requires_packed_pose(func: P) -> Union[PackedPose, None, P]:
+    """
+    Use this as a Python decorator of any user-provided PyRosetta protocol.
+    If a user-provided PyRosetta protocol requires that the first argument
+    parameter be a non-empty `PackedPose` object, then return any received empty
+    `PackedPose` objects or `NoneType` objects and skip the decorated protocol,
+    otherwise run the decorated protocol.
+
+    If using `PyRosettaCluster(filter_results=False)` and the preceding protocol
+    returns or yields either `None`, an empty `Pose` object, or an empty `PackedPose`
+    object, then an empty `PackedPose` object is distributed to the next user-provided
+    PyRosetta protocol, in which case the next protocol and/or any downstream
+    protocols are skipped if they are decorated with this decorator. If using
+    `PyRosettaCluster(ignore_errors=True)` and an error is raised in the preceding
+    protocol, then a `NoneType` object is distributed to the next user-provided
+    PyRosetta protocol, in which case the next protocol and/or any downstream
+    protocols are skipped if they are decorated with this decorator.
+
+    For example:
+
+    @requires_packed_pose
+    def my_pyrosetta_protocol(packed_pose, **kwargs):
+        assert packed_pose.pose.size() > 0
+        return packed_pose
+
+    Args:
+        A user-provided PyRosetta function.
+
+    Returns:
+        The input `packed_pose` argument parameter if it is an empty `PackedPose` object
+        or a `NoneType` object, otherwise the results from the decorated protocol.
+    """
+    @wraps(func)
+    def wrapper(packed_pose, **kwargs):
+        _msg = "User-provided PyRosetta protocol '{0}' received and is duly returning {1} object."
+        if is_empty(packed_pose):
+            logging.info(_msg.format(func.__name__, "an empty `PackedPose`"))
+            return packed_pose
+        elif packed_pose is None:
+            logging.info(_msg.format(func.__name__, "a `NoneType`"))
+            return packed_pose
+        else:
+            return func(packed_pose, **kwargs)
+
+    return cast(P, wrapper)
+
+
 def reproduce(
     input_file: Optional[str] = None,
     scorefile: Optional[str] = None,

source/src/python/PyRosetta/src/pyrosetta/distributed/packed_pose/core.py

@@ -74,6 +74,9 @@ def clone(self):
         result.scores = pickle.loads(pickle.dumps(self.scores))
         return result
 
+    def empty(self):
+        return self.pose.empty()
+
 
 def pack_result(func):
     @functools.wraps(func)

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

@@ -147,6 +147,7 @@ def my_pyrosetta_protocol_2(packed_pose, **kwargs):
                 save_all=False,
                 system_info=None,
                 pyrosetta_build=None,
+                filter_results=None,
                 max_delay_time=3.0,
             )
             cluster.distribute(my_pyrosetta_protocol_1, my_pyrosetta_protocol_2)