Repository: ray-project/ray

Issue:
[RLlib] Errors originating from environment workers cause training to stall
### What is the problem?

Ray: 0.8.0/0.8.1
OS: Ubuntu 16.04/18.04

Errors originating from environment `step()` or `reset()` functions cause training to stall when using  `AsyncSamplesOptimizer` (e.g. IMPALA) even if `ignore_worker_failures: True` is set in the experiment config.

If an environment request fails, the generator returned by `TaskPool.completed_prefetch()` never runs to completion and so `self._fetching` isn't cleared down properly. This results in stale object ids being used on the next training iteration causing #7105 to manifest itself, stalling the experiment.

### Reproduction

1. Set up an IMPALA experiment with external environments
2. After a few training iterations, kill one of the external environments

#### Result:

After `Trainer._try_recover()` has black listed the failed worker, the experiment will stall with the following call stack:
```
Thread 46562 (idle): "MainThread"
    get_objects (ray/worker.py:318)
    get (ray/worker.py:1450)
    ray_get_and_free (ray/rllib/utils/memory.py:33)
    _augment_with_replay (ray/rllib/optimizers/aso_aggregator.py:170)
    iter_train_batches (ray/rllib/optimizers/aso_aggregator.py:117)
    _step (ray/rllib/optimizers/async_samples_optimizer.py:178)
    step (ray/rllib/optimizers/async_samples_optimizer.py:136)
    _train (ray/rllib/agents/trainer_template.py:129)
    train (ray/tune/trainable.py:176)
    train (ray/rllib/agents/trainer.py:433)
    actor_method_executor (ray/function_manager.py:766)
    main_loop (ray/worker.py:433)
    <module> (ray/workers/default_worker.py:118)
```

### Potential Fix

Rather than using a list for `self._fetching` in `TaskPool`, you can use a queue so that items that are going to be freed are automatically removed from the list of tasks to be fetched. This also ensures that the fetching list is always up to date, even if the generator is stopped early.

Here's a example fix that I'm currently testing:
[actors.py.diff](https://gist.github.com/elpollouk/1d0cb83bd98c7a9fa4c9226c66ea07ee)

I haven't dug into the access pattern for `TaskPool` beyond my current scenario, so I'm not sure if a lock should be added around modifications to `self._fetching`. I'm also unaware if there are any subtle side effects of removing items from the fetching list that the original authors were considering when implementing the current logic, so any thoughts here would be greatly appreciated.


Discussion:
@elpollouk could you try out the latest wheel and see if this is still an issue? The hang fix is merged in master as of https://github.com/ray-project/ray/pull/7117
Trying the latest nightly and the experiment does continue to run after knocking out workers.

However, it does appear that you're still re-using object ids for sample batches you've already processed after an error. Is this a concern?

The issue I was attempting to fix with my patch was that you are relying on the `completed_prefetch` generator to run to completion in order to maintain correct state in the `TaskPool`. However, generator completion is never guaranteed as iteration can be stopped at any time either by an exception or other coded stopping condition. Although fixing the symptom of #7105 is working now, the fact that it possible for `TaskPool` to get itself into an inconsistent state will likely manifest itself again in the future via other non-obvious bugs.
The patch makes sense. Do you want to make a PR for it?

Resolving patch:
diff --git a/rllib/utils/actors.py b/rllib/utils/actors.py
--- a/rllib/utils/actors.py
+++ b/rllib/utils/actors.py
@@ -1,6 +1,7 @@
 import logging
 import os
 import ray
+from collections import deque
 
 logger = logging.getLogger(__name__)
 
@@ -11,7 +12,7 @@ class TaskPool:
     def __init__(self):
         self._tasks = {}
         self._objects = {}
-        self._fetching = []
+        self._fetching = deque()
 
     def add(self, worker, all_obj_ids):
         if isinstance(all_obj_ids, list):
@@ -38,15 +39,11 @@ def completed_prefetch(self, blocking_wait=False, max_yield=999):
         for worker, obj_id in self.completed(blocking_wait=blocking_wait):
             self._fetching.append((worker, obj_id))
 
-        remaining = []
-        num_yielded = 0
-        for worker, obj_id in self._fetching:
-            if num_yielded < max_yield:
-                yield (worker, obj_id)
-                num_yielded += 1
-            else:
-                remaining.append((worker, obj_id))
-        self._fetching = remaining
+        for _ in range(max_yield):
+            if not self._fetching:
+                break
+
+            yield self._fetching.popleft()
 
     def reset_workers(self, workers):
         """Notify that some workers may be removed."""
@@ -54,11 +51,14 @@ def reset_workers(self, workers):
             if ev not in workers:
                 del self._tasks[obj_id]
                 del self._objects[obj_id]
-        ok = []
-        for ev, obj_id in self._fetching:
+
+        # We want to keep the same deque reference so that we don't suffer from
+        # stale references in generators that are still in flight
+        for _ in range(len(self._fetching)):
+            ev, obj_id = self._fetching.popleft()
             if ev in workers:
-                ok.append((ev, obj_id))
-        self._fetching = ok
+                # Re-queue items that are still valid
+                self._fetching.append((ev, obj_id))
 
     @property
     def count(self):


Test patch:
diff --git a/rllib/utils/tests/test_taskpool.py b/rllib/utils/tests/test_taskpool.py
new file mode 100644
--- /dev/null
+++ b/rllib/utils/tests/test_taskpool.py
@@ -0,0 +1,138 @@
+import unittest
+from unittest.mock import patch
+
+import ray
+from ray.rllib.utils.actors import TaskPool
+
+
+def createMockWorkerAndObjectId(obj_id):
+    return ({obj_id: 1}, obj_id)
+
+
+class TaskPoolTest(unittest.TestCase):
+    @patch("ray.wait")
+    def test_completed_prefetch_yieldsAllComplete(self, rayWaitMock):
+        task1 = createMockWorkerAndObjectId(1)
+        task2 = createMockWorkerAndObjectId(2)
+        # Return the second task as complete and the first as pending
+        rayWaitMock.return_value = ([2], [1])
+
+        pool = TaskPool()
+        pool.add(*task1)
+        pool.add(*task2)
+
+        fetched = list(pool.completed_prefetch())
+        self.assertListEqual(fetched, [task2])
+
+    @patch("ray.wait")
+    def test_completed_prefetch_yieldsAllCompleteUpToDefaultLimit(
+            self, rayWaitMock):
+        # Load the pool with 1000 tasks, mock them all as complete and then
+        # check that the first call to completed_prefetch only yields 999
+        # items and the second call yields the final one
+        pool = TaskPool()
+        for i in range(1000):
+            task = createMockWorkerAndObjectId(i)
+            pool.add(*task)
+
+        rayWaitMock.return_value = (list(range(1000)), [])
+
+        # For this test, we're only checking the object ids
+        fetched = [pair[1] for pair in pool.completed_prefetch()]
+        self.assertListEqual(fetched, list(range(999)))
+
+        # Finally, check the next iteration returns the final taks
+        fetched = [pair[1] for pair in pool.completed_prefetch()]
+        self.assertListEqual(fetched, [999])
+
+    @patch("ray.wait")
+    def test_completed_prefetch_yieldsAllCompleteUpToSpecifiedLimit(
+            self, rayWaitMock):
+        # Load the pool with 1000 tasks, mock them all as complete and then
+        # check that the first call to completed_prefetch only yield 999 items
+        # and the second call yields the final one
+        pool = TaskPool()
+        for i in range(1000):
+            task = createMockWorkerAndObjectId(i)
+            pool.add(*task)
+
+        rayWaitMock.return_value = (list(range(1000)), [])
+
+        # Verify that only the first 500 tasks are returned, this should leave
+        # some tasks in the _fetching deque for later
+        fetched = [pair[1] for pair in pool.completed_prefetch(max_yield=500)]
+        self.assertListEqual(fetched, list(range(500)))
+
+        # Finally, check the next iteration returns the remaining tasks
+        fetched = [pair[1] for pair in pool.completed_prefetch()]
+        self.assertListEqual(fetched, list(range(500, 1000)))
+
+    @patch("ray.wait")
+    def test_completed_prefetch_yieldsRemainingIfIterationStops(
+            self, rayWaitMock):
+        # Test for issue #7106
+        # In versions of Ray up to 0.8.1, if the pre-fetch generator failed to
+        # run to completion, then the TaskPool would fail to clear up already
+        # fetched tasks resulting in stale object ids being returned
+        pool = TaskPool()
+        for i in range(10):
+            task = createMockWorkerAndObjectId(i)
+            pool.add(*task)
+
+        rayWaitMock.return_value = (list(range(10)), [])
+
+        # This should fetch just the first item in the list
+        try:
+            for _ in pool.completed_prefetch():
+                # Simulate a worker failure returned by ray.get()
+                raise ray.exceptions.RayError
+        except ray.exceptions.RayError:
+            pass
+
+        # This fetch should return the remaining pre-fetched tasks
+        fetched = [pair[1] for pair in pool.completed_prefetch()]
+        self.assertListEqual(fetched, list(range(1, 10)))
+
+    @patch("ray.wait")
+    def test_reset_workers_pendingFetchesFromFailedWorkersRemoved(
+            self, rayWaitMock):
+        pool = TaskPool()
+        # We need to hold onto the tasks for this test so that we can fail a
+        # specific worker
+        tasks = []
+
+        for i in range(10):
+            task = createMockWorkerAndObjectId(i)
+            pool.add(*task)
+            tasks.append(task)
+
+        # Simulate only some of the work being complete and fetch a couple of
+        # tasks in order to fill the fetching queue
+        rayWaitMock.return_value = ([0, 1, 2, 3, 4, 5], [6, 7, 8, 9])
+        fetched = [pair[1] for pair in pool.completed_prefetch(max_yield=2)]
+
+        # As we still have some pending tasks, we need to update the
+        # completion states to remove the completed tasks
+        rayWaitMock.return_value = ([], [6, 7, 8, 9])
+
+        pool.reset_workers([
+            tasks[0][0],
+            tasks[1][0],
+            tasks[2][0],
+            tasks[3][0],
+            # OH NO! WORKER 4 HAS CRASHED!
+            tasks[5][0],
+            tasks[6][0],
+            tasks[7][0],
+            tasks[8][0],
+            tasks[9][0]
+        ])
+
+        # Fetch the remaining tasks which should already be in the _fetching
+        # queue
+        fetched = [pair[1] for pair in pool.completed_prefetch()]
+        self.assertListEqual(fetched, [2, 3, 5])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)