

TITLE
End2End RAG training hangs

COMMENTS
2

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
### System Info

```shell
Python 3.9.0 (default, Nov 15 2020, 14:28:56) 
...
>>> transformers.__version__
'4.17.0'
>>> torch.__version__
'1.7.1'
>>> ray.__version__
'1.11.0'
```


### Who can help?

@shamanez, thanks for a great example! I'm having trouble training a model end2end and wonder if you could help out. I have two GPUs available so I've set the arguments relevant for the end2end training to:
```
    --gpus 1  \
    --end2end \
    --distributed_retriever ray \
    --num_retrieval_workers 4  \
    --index_gpus 1 \
    --gpu_order [3,6]
```

This configuration doesn't seem to work since it results in `len(self.retrieval_workers)=0` which means that the call to `re_load()` in `RagRayDistributedRetriever` just hangs eternally. 

I also tried setting `--gpu 2` above, but that breaks with error 

```*** ValueError: Failed to look up actor with name 'retrieval_worker_0'. This could because 1. You are trying to look up a named actor you didn't create. 2. The named actor died. 3. You did not use a namespace matching the namespace of the actor.```

I don't quite know how ray works so I don't know if ranging `re_load()' to this would break it:

```
        if len(self.retrieval_workers) > 0:
          ray.get([worker.clear_object.remote() for worker in self.retrieval_workers])

          # build the index object again
          index = self._build_index(self.config)

          ray.get(
              [
                  worker.create_rag_retriever.remote(
                      self.config, self.question_encoder_tokenizer, self.generator_tokenizer, index
                  )
                  for worker in self.retrieval_workers
              ]
          )
        else:
          self.index = self._build_index(self.config)
```

What are your thoughts? Thanks!

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Set end2end arguments as shown above.

### Expected behavior

```shell
Continuous training without hanging at function re_load()
```
