feat: multi-turn search R1 example #914

soodoshll · 2025-08-13T17:29:47Z

What does this PR do ?

Address #657

Issues

List issues that this PR closes (syntax):

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

...

wangshangsam

I have a feeling that the retrieval server needs to be moved into nemo_rl/environments/search/, cuz SearchEnv seems to be coupled with this particular retrival server with its particular (HTTP) API design.

wangshangsam · 2025-08-15T20:01:55Z

examples/searchR1/readme.md

+
+```bash
+uv pip install -U cmake
+git clone https://github.com/facebookresearch/faiss.git thirdparty/faiss


Following the existing convention, it might make sense to make faiss a submodule of this repo.

wangshangsam · 2025-08-15T20:05:48Z

examples/searchR1/readme.md

+
+The following instructions install Faiss-GPU from source. Please refer to [this](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md) for more details.
+
+```bash


To me, it seems there are enough shell commands in here to automate them into a script.

wangshangsam · 2025-08-15T20:25:12Z

examples/searchR1/readme.md

+```
+local_dir=./data/searchR1
+uv run --active python searchr1_download.py --local_dir $local_dir
+cat $local_dir/part_* > $local_dir/e5_Flat.index
+gzip -d $local_dir/wiki-18.jsonl.gz
+uv run --active python searchr1_dataset.py --local_dir $local_dir
+```


Did you intend to replace this with bash prepare.sh?

yes, fixed.

wangshangsam · 2025-08-15T20:31:12Z

examples/searchR1/requirements.txt

+numpy==1.26.4
+transformers
+datasets
+pyserini
+huggingface_hub
+# faiss-gpu-cu12
+uvicorn
+fastapi
+torch==2.6.0


Suggested change

numpy==1.26.4

transformers

datasets

pyserini

huggingface_hub

# faiss-gpu-cu12

uvicorn

fastapi

torch==2.6.0

pyserini

huggingface_hub

# faiss-gpu-cu12

uvicorn

fastapi

Only include dependencies that are not already presented in pyproject.toml

wangshangsam · 2025-08-15T20:38:45Z

examples/searchR1/retrieval_server.py

+    return corpus
+
+
+def read_jsonl(file_path):


Could you annonate the types and add docstrings for all the functions in this file?

wangshangsam · 2025-08-15T20:41:54Z

examples/searchR1/run_grpo_searchr1.py

+    tokenizer: TokenizerType,
+    max_seq_length: int,
+    idx: int,
+) -> DatumSpec:


Add docstrings to the functions in this file.

wangshangsam · 2025-08-15T20:46:24Z

examples/searchR1/searchr1_dataset.py

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Adapted from https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl-train/examples/search/searchr1_dataset.py


(This is just a side note; no need to address anything)

Both files are from NovaSky-AI/SkyRL, but somehow this one has noticeably better coding style than examples/searchR1/retrieval_server.py

wangshangsam · 2025-08-15T20:48:40Z

nemo_rl/environments/search/search.py

+INITIAL_RETRY_DELAY = 1
+
+
+class SearchEnvConfig(TypedDict):


Add docstring for all classes and functions/methods in this file.

wangshangsam · 2025-08-15T21:00:04Z

nemo_rl/environments/search/search.py

+    answer: Optional[str]
+
+
+def call_search_api(


Suggested change

def call_search_api(

def _call_search_api(

This seems more like an internal helper function to me.

wangshangsam · 2025-08-15T21:07:44Z

nemo_rl/environments/search/search.py

+
+
+@ray.remote
+class SearchEnv(EnvironmentInterface[SearchEnvMetadata]):


We should probably add unit test for this. I'm suspecting that it's not gonna be easy to actually bring up a 2-GPU retrival server for unit tests, hence you might need to mock requests.session for unit testing this.

euronymous-aithal · 2025-08-20T00:24:11Z

@soodoshll will you be able to address comments ASAP

wangshangsam · 2025-08-20T01:03:12Z

@soodoshll will you be able to address comments ASAP

@soodoshll is currently working on #883 which, as I understand, has a much higher priority over this Search-R1 example (i.e., a nice-to-have). If the priority should be the reverse, we can adapt accordingly.

soodoshll · 2025-08-20T14:36:50Z

I have a feeling that the retrieval server needs to be moved into nemo_rl/environments/search/, cuz SearchEnv seems to be coupled with this particular retrival server with its particular (HTTP) API design.

A general issue is that the retrieval server might run in an environment (quite) different from the nemo-rl env. For example, faiss depends on numpy 1.x

Signed-off-by: Qidong Su <[email protected]> update Signed-off-by: Qidong Su <[email protected]> upd Signed-off-by: Qidong Su <[email protected]> stash Signed-off-by: Qidong Su <[email protected]> fix many things Signed-off-by: Qidong Su <[email protected]> stash Signed-off-by: Qidong Su <[email protected]> clean Signed-off-by: Qidong Su <[email protected]> clean Signed-off-by: Qidong Su <[email protected]> update copyright Signed-off-by: Qidong Su <[email protected]> update copyright Signed-off-by: Qidong Su <[email protected]> clean fix Signed-off-by: Qidong Su <[email protected]> fix Signed-off-by: Qidong Su <[email protected]> fix Signed-off-by: Qidong Su <[email protected]>

Signed-off-by: Qidong Su <[email protected]>

soodoshll force-pushed the searchR1 branch from e672072 to 636a141 Compare August 13, 2025 17:36

github-actions bot added documentation Improvements or additions to documentation CI Relating to CI labels Aug 13, 2025

soodoshll marked this pull request as draft August 13, 2025 17:38

soodoshll force-pushed the searchR1 branch from 636a141 to ea25edb Compare August 13, 2025 18:12

github-actions bot removed documentation Improvements or additions to documentation CI Relating to CI labels Aug 13, 2025

soodoshll force-pushed the searchR1 branch from ebd840d to 5684764 Compare August 13, 2025 18:16

soodoshll marked this pull request as ready for review August 13, 2025 18:32

soodoshll requested a review from wangshangsam August 13, 2025 18:32

wangshangsam requested changes Aug 15, 2025

View reviewed changes

wangshangsam assigned soodoshll Aug 20, 2025

soodoshll and others added 3 commits August 20, 2025 07:43

clean

531b243

Signed-off-by: Qidong Su <[email protected]>

fix

7ad1988

Signed-off-by: Qidong Su <[email protected]>

soodoshll force-pushed the searchR1 branch from e104244 to 7ad1988 Compare August 20, 2025 14:43


		The following instructions install Faiss-GPU from source. Please refer to [this](https://github.com/facebookresearch/faiss/blob/main/INSTALL.md) for more details.

		```bash



		@ray.remote
		class SearchEnv(EnvironmentInterface[SearchEnvMetadata]):

feat: multi-turn search R1 example #914

Are you sure you want to change the base?

feat: multi-turn search R1 example #914

Uh oh!

Conversation

soodoshll commented Aug 13, 2025

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

wangshangsam left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

euronymous-aithal commented Aug 20, 2025

Uh oh!

wangshangsam commented Aug 20, 2025

Uh oh!

soodoshll commented Aug 20, 2025

Uh oh!

Uh oh!