[RLlib] In the Json_writer convert all non string keys to keys #33896

avnishn · 2023-03-29T23:26:23Z

Currently ray data has 2 formats in which it can store python dictionaries as datasets internally.

it will try to use pyarrow to store the underlying data.
it will try to use pandas / dictionaries to store the underlying data in the dataset.
In case 1, when writing data out to a json file, the data will be written out in the expected schema. e.g. {YOUR DATA: ..., }.
IN case 2, when writing data out to json, the schema is {"value": {YOUR DATA: ....}.

Ray data tries to go for 1 always but falls back onto 2 in the case that the dictionary being converted to a dataset can't be processed by arrow.

One of the reasons for fall back from 1 to 2 is if any of the keys in this dictionary are not strings

IN RLLIB this happens whenever our sample batch includes environment infos, and those environment infos have non string keys. This happens when the env infos are mapping agent ids to infos and those agent ids are integers.

Env infos are included as a part of the batch with torch but not tensorflow which is why we saw this issue when we made the default framework torch last week.

The solution for now is to convert non string keys to string keys, but to also force ray data to always do 1 and use py arrow to store the underlying data. We do this with the output_arrow_format flag added to from_list in #33837

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: amogkam <[email protected]>

Signed-off-by: Amog Kamsetty <[email protected]>

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

… from-items-force-simple

…_writer_convert_keys_to_string

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

avnishn · 2023-03-29T23:33:59Z

python/ray/data/datasource/binary_datasource.py

all ray data changes in this pr are from #33837 and should be ignored during review.

avnishn · 2023-03-29T23:34:05Z

doc/source/ray-air/examples/rl_offline_example.ipynb

ptal at this file

avnishn · 2023-03-29T23:34:23Z

rllib/offline/dataset_writer.py

ptal at this file

avnishn · 2023-03-29T23:34:36Z

rllib/offline/json_writer.py

ptal at this file

gjoliver · 2023-03-30T04:42:34Z

rllib/offline/json_writer.py

 return v


+def _convert_keys_to_strings_json_dict(d: Dict[Any, Any]) -> Dict[Any, Any]:


I am actually not sure if you can do this.
like what should we do after we read this json data back in? how would we know which key to convert back into integers?
maybe just clear infos column here for now?

Ok I'm in agreement sure.

ollie-iterators · 2023-03-30T14:49:51Z

Strange, I would think the error about "can't pickle _thread.lock objects" would be related to #33660, but it appears that "Dataset tests (Arrow nightly)" is passing right now on the main branch.

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

…_writer_convert_keys_to_string

avnishn · 2023-03-30T21:02:17Z

we need to wait for #33837 to be merged but otherwise this is ready to go afaict.

gjoliver

ok, this looks reasonable.
can we wait for Amog to merge his PR, then rebase here, so we can have a final look before merge?
thanks for the fix.

…_writer_convert_keys_to_string

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

as it turns out there was a way to store infos in tensorflow for training. Its not publicly documented, but there is 1 test for it, so there were 2 occurences in the codebase. This whole pr pretty much disables you from being able to write env infos out in a sample batch, so we'll allow it to still be used for training, but not allow it for output writing. Signed-off-by: Avnish <[email protected]>

avnishn · 2023-04-10T22:48:07Z

rllib/tests/test_io.py

@@ -80,20 +80,6 @@ def test_agent_output_logdir(self):
 agent = self.write_outputs("logdir", fw)
 self.assertEqual(len(glob.glob(agent.logdir + "/output-*.json")), 1)

- def test_agent_output_infos(self):


as it turns out there was a way to store infos in tensorflow for training.
Its not publicly documented, but there is 1 test for it, so there were
2 occurences in the codebase. This whole pr pretty much disables you from being able
to write env infos out in a sample batch, so we'll allow it to still be used for training,
but not allow it for output writing.

I removed the test where we check for env infos in the output batch for the time being.

cc @gjoliver

…ay-project#33896) Signed-off-by: Avnish <[email protected]> Signed-off-by: elliottower <[email protected]>

…ay-project#33896) Signed-off-by: Avnish <[email protected]> Signed-off-by: Jack He <[email protected]>

amogkam and others added 8 commits March 28, 2023 18:07

update

376cc7e

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

update

4c76849

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

Update doc/requirements-doc.txt

f298199

Signed-off-by: Amog Kamsetty <[email protected]>

clarify

86f2293

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

Merge branch 'from-items-force-simple' of github.com:amogkam/ray into…

7046b4e

… from-items-force-simple

Merge branch 'master' of https://github.com/ray-project/ray into json…

53a8f12

…_writer_convert_keys_to_string

Have Json Writer convert all sample batch keys to strings

e6155aa

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

Have Json Writer convert all sample batch keys to strings

e5d12f0

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

avnishn requested review from richardliaw, gjoliver, krfricke, xwjiang2010, amogkam, matthewdeng, Yard1, maxpumperla, a team, sven1977, ArturNiederfahrenhorst, smorad, kouroshHakha, ericl, scv119, clarkzinzow, jjyao, jianoaix and c21 as code owners March 29, 2023 23:26

avnishn changed the title ~~[RLlib] Json_writer_convert_keys_to_string~~ [RLlib] In the Json_writer convert all non string keys to keys Mar 29, 2023

avnishn assigned gjoliver Mar 29, 2023

avnishn commented Mar 29, 2023

View reviewed changes

gjoliver reviewed Mar 30, 2023

View reviewed changes

avnishn added 2 commits March 30, 2023 10:03

Address feedback delete sample batch info key when writing out

feee3ac

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

Merge branch 'master' of https://github.com/ray-project/ray into json…

9ebcb01

…_writer_convert_keys_to_string

gjoliver reviewed Mar 30, 2023

View reviewed changes

Merge branch 'master' of https://github.com/ray-project/ray into json…

b5f2a96

…_writer_convert_keys_to_string

gjoliver approved these changes Apr 10, 2023

View reviewed changes

avnishn added 3 commits April 10, 2023 11:11

Merge branch 'master' of https://github.com/ray-project/ray into json…

5c18d70

…_writer_convert_keys_to_string

Check if infos in sample batch first before deleting

c283bc2

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

avnishn commented Apr 10, 2023

View reviewed changes

gjoliver merged commit a3544a2 into ray-project:master Apr 11, 2023

avnishn linked an issue Apr 11, 2023 that may be closed by this pull request

[RLlib] something broken about the way that torch ppo writes out to json #33789

Closed

elliottower pushed a commit to elliottower/ray that referenced this pull request Apr 22, 2023

[RLlib] Remove infos dict before Json_writer writes sample batches (r…

86ae35c

…ay-project#33896) Signed-off-by: Avnish <[email protected]> Signed-off-by: elliottower <[email protected]>

ProjectsByJackHe pushed a commit to ProjectsByJackHe/ray that referenced this pull request May 4, 2023

[RLlib] Remove infos dict before Json_writer writes sample batches (r…

bd1fe8b

…ay-project#33896) Signed-off-by: Avnish <[email protected]> Signed-off-by: Jack He <[email protected]>

Rohan138 mentioned this pull request Jun 20, 2023

[RLlib] RLlib does not write full info dict on terminated step #35440

Closed

sven1977 mentioned this pull request Sep 13, 2023

[RLlib] Issue 35440: JSON output writer should include INFOs. #39632

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RLlib] In the Json_writer convert all non string keys to keys #33896

[RLlib] In the Json_writer convert all non string keys to keys #33896

avnishn commented Mar 29, 2023 •

edited

Loading

avnishn Mar 29, 2023

avnishn Mar 29, 2023

avnishn Mar 29, 2023

avnishn Mar 29, 2023

gjoliver Mar 30, 2023

avnishn Mar 30, 2023

ollie-iterators commented Mar 30, 2023 •

edited

Loading

avnishn commented Mar 30, 2023

gjoliver left a comment

avnishn Apr 10, 2023

		return v


		def _convert_keys_to_strings_json_dict(d: Dict[Any, Any]) -> Dict[Any, Any]:

[RLlib] In the Json_writer convert all non string keys to keys #33896

[RLlib] In the Json_writer convert all non string keys to keys #33896

Conversation

avnishn commented Mar 29, 2023 • edited Loading

Why are these changes needed?

Related issue number

Checks

avnishn Mar 29, 2023

Choose a reason for hiding this comment

avnishn Mar 29, 2023

Choose a reason for hiding this comment

avnishn Mar 29, 2023

Choose a reason for hiding this comment

avnishn Mar 29, 2023

Choose a reason for hiding this comment

gjoliver Mar 30, 2023

Choose a reason for hiding this comment

avnishn Mar 30, 2023

Choose a reason for hiding this comment

ollie-iterators commented Mar 30, 2023 • edited Loading

avnishn commented Mar 30, 2023

gjoliver left a comment

Choose a reason for hiding this comment

avnishn Apr 10, 2023

Choose a reason for hiding this comment

avnishn commented Mar 29, 2023 •

edited

Loading

ollie-iterators commented Mar 30, 2023 •

edited

Loading