[RLlib] RLModule: `InferenceOnlyAPI`. #47572

sven1977 · 2024-09-09T19:26:42Z

Introduce the RLModule InferenceOnlyAPI to simplify and generalize the concept of defining an inference-only version of ones (custom) RLModule.

Users that want to save space on the EnvRunners (on which all RLModule are constructed with the inference_only flag set to True), can now override a single method in their RLModule code (get_non_inference_only_attributes) that returns those attribute names (possibly nested) that point to sub-components of the model NOT needed for action computations.

RLlib will then automatically:
a) remove those parts in inference-only mode.
b) adjust the state returned by get_state(inference_only=True) from an inference_only=False RLModule which does implement the InferenceOnlyAPI (capable of having an inference_only version). This way, when getting the weights from a Learner (inference_only=False) for an EnvRunner (inference_only=True), the weight matrix is reduced already on the Learner saving network traffic.

Why are these changes needed?

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: sven1977 <[email protected]>

simonsays1980

LGTM. Some questions here and there. Awesome PR that saves us a lot of code 👍

simonsays1980 · 2024-09-10T08:14:26Z

rllib/core/rl_module/apis/inference_only_api.py

+    """An API to be implemented by RLModules that have an inference-only mode.
+
+    Only the `get_non_inference_attributes` method needs to get implemented for
+    a RLModule to have the following functionality:


Small nit "a" -> "an RLModule" ;)

simonsays1980 · 2024-09-10T08:14:54Z

rllib/core/rl_module/apis/inference_only_api.py

+
+    Only the `get_non_inference_attributes` method needs to get implemented for
+    a RLModule to have the following functionality:
+    - On EnvRunners (or when self.config.inference_only=True), RLlib will remove


Can we indent these? It makes it easier to read.

I'm not sure it'll pass the doctests, but let's try ...

simonsays1980 · 2024-09-10T08:17:50Z

rllib/core/rl_module/apis/inference_only_api.py

+
+    @abc.abstractmethod
+    def get_non_inference_attributes(self) -> List[str]:
+        """Returns a list of names (str) of attributes of inference-only components.


This is confusing. The method is named get_NON_inference_attributes and the docstring tells that this returns attributes of inference-only attributes.

Maybe "Returns a list of names (str) of attributes not used in inference mode."

fixed and clarified

simonsays1980 · 2024-09-10T08:18:23Z

rllib/core/rl_module/apis/inference_only_api.py

+
+            from ray.rllib.core.rl_module.rl_module import RLModuleSpec
+
+            spec = RLModuleSpec(module_class=..., inference_only=True)


Very nice example!!

simonsays1980 · 2024-09-10T08:22:23Z

rllib/core/rl_module/torch/torch_rl_module.py

+        # If an inference-only class AND self.config.inference_only is True,
+        # remove all attributes that are returned by
+        # `self.get_non_inference_attributes()`.
+        if self.config.inference_only and isinstance(self, InferenceOnlyAPI):


simonsays1980 · 2024-09-10T08:28:00Z

rllib/core/rl_module/torch/torch_rl_module.py

+                    continue
+                target = getattr(self, parts[0])
+                # Traverse from the next part on (if nested).
+                for part in parts[1:]:


This stops after the first nesting of modules. Can't we just use a recursive implementation that calls itself again?

Also, this first builds all the components (could be many) and then deletes them again. This could induce some initialization cost which hurts performance in cases where workers need to be rebuild often.

This is true, BUT users can always decide to not even build these components in the first place inside their setup(). In this case, the above functionality will simply ignore the (already) missing attributes (which it would have deleted otherwise).

I think, the logic is actually fine. It does not stop after the first level, but goes all the way to the end of the specified attribute string, e.g:

"a.b.c.d.e".

If any sub-attribute above does not exist, break and del nothing (attribute already missing).

If "a.b.c.d.e" can be found -> del only the "e" sub-attribute.

simonsays1980 · 2024-09-10T08:28:35Z

rllib/core/rl_module/torch/torch_rl_module.py

+        if (
+            inference_only
+            and not self.config.inference_only
+            and isinstance(self, InferenceOnlyAPI)


I like this form of testing.

Yeah, me, too. This is what Learners should do: Check, whether the given RLModules are actually compatible with them (or expose certain known APIs to work with).
This way, we can eliminate the need to have algo-specific RLModule classes, freeing the user to think about these beforehand and making it easier to switch between algos later.

simonsays1980 · 2024-09-10T08:43:31Z

rllib/algorithms/dqn/dqn_rainbow_rl_module.py

@@ -85,6 +79,10 @@ def make_target_networks(self) -> None:
        if self.uses_dueling:
            self._target_vf = make_target_network(self.vf)

+    @override(InferenceOnlyAPI)
+    def get_non_inference_attributes(self) -> List[str]:
+        return ["_target_encoder", "_target_af", "_target_vf"]


return ["_target_encoder", "_target_af"] + ["_target_vf"] if self.uses_dueling else []

simonsays1980 · 2024-09-10T08:50:57Z

rllib/algorithms/ppo/ppo_rl_module.py

+        ret = ["vf"]
+        if hasattr(self.encoder, "critic_encoder"):
+            ret += ["encoder.critic_encoder"]
+        return ret


return ["vf"] + [] if self.config.model_config_dict["vf_share_layers"] else ["encoder.critic_encoder"]

simonsays1980 · 2024-09-10T08:52:57Z

rllib/algorithms/sac/sac_rl_module.py

+    @override(InferenceOnlyAPI)
+    def get_non_inference_attributes(self) -> List[str]:
+        ret = ["qf", "target_qf", "qf_encoder", "target_qf_encoder"]
+        if self.twin_q:


return ["qf", "target_qf", "qf_encoder", "target_qf_encoder"] + [ "qf_twin", "target_qf_twin", "qf_twin_encoder", "target_qf_twin_encoder", ] if self.twin_q else []

Signed-off-by: sven1977 <[email protected]>

Signed-off-by: ujjawal-khare <[email protected]>

wip

8c2f471

Signed-off-by: sven1977 <[email protected]>

sven1977 requested review from ArturNiederfahrenhorst and simonsays1980 as code owners September 9, 2024 19:26

sven1977 assigned simonsays1980 Sep 9, 2024

wip

119d6df

Signed-off-by: sven1977 <[email protected]>

simonsays1980 approved these changes Sep 10, 2024

View reviewed changes

sven1977 enabled auto-merge (squash) September 10, 2024 12:27

github-actions bot added the go add ONLY when ready to merge, run all tests label Sep 10, 2024

wip

fb5e7dd

Signed-off-by: sven1977 <[email protected]>

github-actions bot disabled auto-merge September 10, 2024 12:27

sven1977 enabled auto-merge (squash) September 10, 2024 12:31

wip

9247ba8

Signed-off-by: sven1977 <[email protected]>

github-actions bot disabled auto-merge September 10, 2024 14:14

sven1977 added 2 commits September 10, 2024 16:37

wip

7d0da38

Signed-off-by: sven1977 <[email protected]>

wip

35afb88

Signed-off-by: sven1977 <[email protected]>

sven1977 requested review from maxpumperla and a team as code owners September 10, 2024 16:56

sven1977 merged commit 7cb95ce into ray-project:master Sep 11, 2024
5 checks passed

sven1977 deleted the rl_module_api_inference_only branch September 11, 2024 07:55

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

4dcf4c2

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

cb728a5

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

69959f9

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

2bff109

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

db63e7b

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

6de1367

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

54b9084

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

47d9b0d

Signed-off-by: ujjawal-khare <[email protected]>

ujjawal-khare pushed a commit to ujjawal-khare-27/ray that referenced this pull request Oct 15, 2024

[RLlib] RLModule: InferenceOnlyAPI. (ray-project#47572)

6e90110

Signed-off-by: ujjawal-khare <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RLlib] RLModule: `InferenceOnlyAPI`. #47572

[RLlib] RLModule: `InferenceOnlyAPI`. #47572

sven1977 commented Sep 9, 2024 •

edited

Loading

simonsays1980 left a comment

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

simonsays1980 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024

sven1977 Sep 10, 2024

simonsays1980 Sep 10, 2024


		from ray.rllib.core.rl_module.rl_module import RLModuleSpec

		spec = RLModuleSpec(module_class=..., inference_only=True)

[RLlib] RLModule: InferenceOnlyAPI. #47572

[RLlib] RLModule: InferenceOnlyAPI. #47572

Conversation

sven1977 commented Sep 9, 2024 • edited Loading

Why are these changes needed?

Related issue number

Checks

simonsays1980 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[RLlib] RLModule: `InferenceOnlyAPI`. #47572

[RLlib] RLModule: `InferenceOnlyAPI`. #47572

sven1977 commented Sep 9, 2024 •

edited

Loading