[RLlib] created check_specs decorator, RLModule PR 1/N #29599

kouroshHakha · 2022-10-24T06:45:00Z

Why are these changes needed?

submitting this PR in pieces:
check_specs decorator can be added to any module method to enforce input/output struct types. This is useful for imposing a certain input/output behavior in RLModule without taking away the flexibility of implementation details from the user. User would also be efficiently informed about what needs to be implemented.

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 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 :(

…he names a little bit for generality Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

kouroshHakha · 2022-10-24T06:45:54Z

rllib/models/specs/specs_base.py

@@ -28,7 +28,7 @@ def validate(self, data: Any) -> None:


 @DeveloperAPI
-class TensorSpecs(SpecsAbstract):


Realized they can represent only one spec, hence renaming :)

sven1977 · 2022-10-24T09:33:37Z

rllib/utils/nested_dict.py

@@ -85,7 +85,7 @@ class NestedDict(Generic[T], MutableMapping[str, Union[T, "NestedDict"]]):
            >>>                             #            'b': {'c': 200, 'd': 300}}
            >>> # Getting elements, possibly nested:
            >>> print(foo_dict['b', 'c'])   # 200
-            >>> print(foo_dict['b']) # IndexError("Use get for partial indexing.")
+            >>> print(foo_dict['b']) # IndexError


Dumb question: Is there a reason why we make this seemingly arbitrary distinction between accessing by get (no error) vs direct (error)? Why should both not return the sub-dict? How would the user know this difference?

good point @sven1977, In hindsight I don't see any reason for not returning the sub-nested dict if __getitem__ is used. It actually confused myself at some point during a later pr.

sven1977 · 2022-10-24T09:34:42Z

rllib/utils/nested_dict.py

+            raise IndexError(
+                f"Key `{k}` is not a complete key in the given "
+                f"{self.__class__.__name__}. It results in a container "
+                f"with subkeys {set(output.keys())}. To get partial indexing, "


-> "To use partial indexing and thus retrieve a sub-structure ..."

I removed partial indexing error all together due the valid comment above. Also update the examples in the docstring to show that index error is not raised anymore.

sven1977 · 2022-10-24T09:40:28Z

rllib/models/specs/tests/test_check_specs.py

+            lambda: correct_module.check_input_and_output_wo_filter(input_dict),
+        )
+
+    def test_cache(self):


Very nice test, going the extra mile!

sven1977 · 2022-10-24T09:43:38Z

rllib/models/specs/tests/test_check_specs.py

+        )
+
+        # this should not raise an error because output is not forced to be checked
+        incorrect_module.check_only_input({"input": 2})


Dumb question, why would the decorator itself not already complain when it's being instantiated b/c of the missing output check? In other words, why is it allowed to have an implementation that doesn't check both, in- and output?

incorrect_module's run implementation does not have the correct output type. Therefore those functions that enforce output type checking should raise an error and those that don't should just ignore the output spec enforcement. This is only determined when the function is actually executed and not during the function decoration itself. Implementation if the decorated function is only visible when the function is invoked.

Ah, yes, that makes sense! Thanks for clarifying.

sven1977 · 2022-10-24T09:44:52Z

rllib/models/specs/specs_jax.py

@@ -11,20 +11,20 @@


 @DeveloperAPI
-class JAXSpecs(TensorSpecs):
-    @override(TensorSpecs)
+class JAXTensorSpec(TensorSpec):


Why call this JAXTensorSpec, but the others only XYZSpec (w/o the "tensor", e.g. "TFSpec")?

changed everything to XXXTensorSpec to be more precise, we may have XXXDistributionSpecs down the line too.

sven1977 · 2022-10-24T09:45:38Z

rllib/models/specs/specs_dict.py

+        return f"ModelSpec({repr(self._data)})"
+
+
+def check_specs(


Super nice! This is going to be dope for early-catching user errors.

sven1977

Awesome PR! Thanks for being very meticulous about designing these new APIs from the ground up. These will comprise ground-breaking advances for RLlib toward more user-friendliness and transparency.

Just a bunch of questions and nits.

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

kouroshHakha · 2022-10-24T17:48:37Z

@sven1977 Please re-review.

sven1977

LGTM.

sven1977 · 2022-10-24T18:42:55Z

test_nested_dict failing
cc: @kouroshHakha

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

gjoliver

looks pretty good. just some minor issues.

gjoliver · 2022-10-24T19:33:26Z

rllib/models/specs/specs_dict.py

        ...     },
-        ...     "action": TensorSpecs("b, d_a", h=12),
+        ...     "action": TensorSpec("b, d_a", h=12),
        ...     "action_dist": torch.distributions.Categorical


I have a random question. do you intend for these specs to be checkpointed and restored?
if so, maybe a registry of distribution is a good idea.

I see. This is a good point. This should be incorporated into the RLModule PR then. I'll look into this there.

We should be able to write this directly to the state_dict, but this might break logic that assumes state_dict only contains tensors.

gjoliver · 2022-10-24T19:35:50Z

rllib/models/specs/specs_dict.py

        if exact_match:
            data_spec_missing_keys = data_keys_set.difference(self._keys_set)
            if data_spec_missing_keys:
                raise ValueError(_MISSING_KEYS_FROM_SPEC.format(data_spec_missing_keys))

        for spec_name, spec in self.items():
            data_to_validate = data[spec_name]
-            if isinstance(spec, TensorSpecs):
-                spec.validate(data_to_validate)
+            if isinstance(spec, TensorSpec):


for safety, this needs an "else" clause that just raises error if spec is some random data?

gjoliver · 2022-10-24T19:52:12Z

rllib/models/specs/specs_dict.py

+                        )
+                    data = NestedDict(data)
+
+                if should_validate():


should you check should_validate() first thing, so we don't waste time building NestedDict, etc?

I don't think so, You may still need to filter the data if it's a mapping regardless of whether you should validate or not.

gjoliver · 2022-10-24T19:59:41Z

rllib/models/specs/specs_dict.py

+
+            input_data_ = input_data
+            if input_spec:
+                input_spec_ = getattr(self, input_spec)()


can input_spec_ just be a member variable on self?
or maybe we can simply assume that the specs will be provided by RLModule under some hardcoded key names.
so by default, folks can simply use the decorator without specifying any parameters.

That will assume some hard-coded names on the base class and is not the intention of this general-purpose decorator. The decorator let's the user choose their own spec names. so it can be applied to any base class essentially. It may become handy in defining Pi base class that looks different than RLModule base class. You can see the use-case in the RLModule PR.

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

kouroshHakha · 2022-10-25T01:30:06Z

@gjoliver Can we merge this? The failed tests are again not related.

gjoliver · 2022-10-25T02:36:12Z

mem leak test failures are not related.

…9599) * 1. created check_specs decorator 2. updated unittests 3. refactored the names a little bit for generality Signed-off-by: Kourosh Hakhamaneshi <[email protected]> Signed-off-by: Weichen Xu <[email protected]>

kouroshHakha added 2 commits October 23, 2022 23:39

1. created check_specs decorator 2. updated unittests 3. refactored t…

b22a6d8

…he names a little bit for generality Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

updated errors in nested dict

422d5ad

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

kouroshHakha assigned gjoliver Oct 24, 2022

kouroshHakha requested review from sven1977, gjoliver, avnishn, ArturNiederfahrenhorst, smorad, maxpumperla and krfricke as code owners October 24, 2022 06:45

kouroshHakha commented Oct 24, 2022

View reviewed changes

kouroshHakha assigned sven1977 Oct 24, 2022

kouroshHakha changed the title ~~[RLlib] created check_specs decorator~~ [RLlib] created check_specs decorator, RLModule PR 1/N Oct 24, 2022

Merge branch 'master' into update-specs

8098120

sven1977 reviewed Oct 24, 2022

View reviewed changes

kouroshHakha added 3 commits October 24, 2022 10:25

updated bazel

5be7182

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

made the names consistent

30fe44f

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

removed the partial indexing error in __getitem__

9c74871

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

sven1977 approved these changes Oct 24, 2022

View reviewed changes

kouroshHakha added 3 commits October 24, 2022 11:53

nested dict test update

f7dcb90

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

bazel update

a019dd5

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

making cache test unflakey

11fe4f4

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

gjoliver reviewed Oct 24, 2022

View reviewed changes

made the condition checking safer

f0ec25f

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

gjoliver approved these changes Oct 24, 2022

View reviewed changes

kouroshHakha added 2 commits October 24, 2022 16:09

fixed a hidden bug with tuple's getting skipped in ModelSpecs

e9ffeb1

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

attempt to deflake

9cfb161

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

kouroshHakha added the tests-ok The tagger certifies test failures are unrelated and assumes personal liability. label Oct 25, 2022

gjoliver merged commit 45420f5 into ray-project:master Oct 25, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RLlib] created check_specs decorator, RLModule PR 1/N #29599

[RLlib] created check_specs decorator, RLModule PR 1/N #29599

kouroshHakha commented Oct 24, 2022 •

edited

Loading

kouroshHakha Oct 24, 2022

sven1977 Oct 24, 2022

kouroshHakha Oct 24, 2022

kouroshHakha Oct 24, 2022

sven1977 Oct 24, 2022

kouroshHakha Oct 24, 2022

sven1977 Oct 24, 2022

sven1977 Oct 24, 2022

kouroshHakha Oct 24, 2022 •

edited

Loading

sven1977 Oct 24, 2022

sven1977 Oct 24, 2022

kouroshHakha Oct 24, 2022

sven1977 Oct 24, 2022

sven1977 left a comment

kouroshHakha commented Oct 24, 2022

sven1977 left a comment

sven1977 commented Oct 24, 2022

gjoliver left a comment

gjoliver Oct 24, 2022

kouroshHakha Oct 24, 2022

smorad Oct 24, 2022

gjoliver Oct 24, 2022

kouroshHakha Oct 24, 2022

gjoliver Oct 24, 2022

kouroshHakha Oct 24, 2022

gjoliver Oct 24, 2022

gjoliver Oct 24, 2022

kouroshHakha Oct 24, 2022 •

edited

Loading

kouroshHakha commented Oct 25, 2022

gjoliver commented Oct 25, 2022

		@@ -28,7 +28,7 @@ def validate(self, data: Any) -> None:


		@DeveloperAPI
		class TensorSpecs(SpecsAbstract):

[RLlib] created check_specs decorator, RLModule PR 1/N #29599

[RLlib] created check_specs decorator, RLModule PR 1/N #29599

Conversation

kouroshHakha commented Oct 24, 2022 • edited Loading

Why are these changes needed?

Related issue number

Checks

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

kouroshHakha Oct 24, 2022 • edited Loading

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

sven1977 left a comment

Choose a reason for hiding this comment

kouroshHakha commented Oct 24, 2022

sven1977 left a comment

Choose a reason for hiding this comment

sven1977 commented Oct 24, 2022

gjoliver 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

kouroshHakha Oct 24, 2022 • edited Loading

Choose a reason for hiding this comment

kouroshHakha commented Oct 25, 2022

gjoliver commented Oct 25, 2022

kouroshHakha commented Oct 24, 2022 •

edited

Loading

kouroshHakha Oct 24, 2022 •

edited

Loading

kouroshHakha Oct 24, 2022 •

edited

Loading