rename and move allow_child_to_ptrace to common.py to avoid circular deps; update tests

Author: Le

alf/algorithms/distributed_off_policy_algorithm.py

@@ -33,7 +33,7 @@
 from alf.environments.alf_environment import AlfEnvironment
 from alf.experience_replayers.replay_buffer import ReplayBuffer
 from alf.data_structures import Experience, make_experience, StepType
-from alf.trainers.evaluator import _allow_child_to_ptrace
+from alf.utils.common import allow_child_to_ptrace
 from alf.utils.per_process_context import PerProcessContext
 from alf.utils import dist_utils
 from alf.utils.summary_utils import record_time
@@ -523,7 +523,7 @@ def _create_data_receiver_subprocess(self):
                                     self._ddp_rank),
                               daemon=True)
         process.start()
-        _allow_child_to_ptrace(process.pid)
+        allow_child_to_ptrace(process.pid)
 
     def utd(self):
         total_exps = int(self._replay_buffer.get_current_position().sum())
@@ -713,7 +713,7 @@ def _create_pull_params_subprocess(self):
                                    self._params_socket_rank),
                              daemon=True)
         process.start()
-        _allow_child_to_ptrace(process.pid)
+        allow_child_to_ptrace(process.pid)
 
     def observe_for_replay(self, exp: Experience):
         """Send experience data to the trainer.

alf/experience_replayers/replay_buffer_test.py

@@ -328,11 +328,12 @@ def generate_step_types(self, num_envs, max_steps, end_prob):
         (False, True),
         (True, False),
     ])
-    def test_replay_buffer(self, allow_multiprocess, with_replacement):
+    def test_replay_buffer(self, use_mp_context, with_replacement):
+        mp_ctx = mp.get_context('spawn') if use_mp_context else None
         replay_buffer = ReplayBuffer(data_spec=self.data_spec,
                                      num_environments=self.num_envs,
                                      max_length=self.max_length,
-                                     allow_multiprocess=allow_multiprocess)
+                                     mp_context=mp_ctx)
 
         batch1 = get_exp_batch([0, 4, 7], self.dim, t=0, x=0.1)
         replay_buffer.add_batch(batch1, batch1.env_id)
@@ -702,17 +703,17 @@ class ReplayBufferSharingTest(parameterized.TestCase, alf.test.TestCase):
     def test_spawned_process_sharing(self, start_method):
         spec = alf.TensorSpec((10, ), dtype=torch.uint8)
 
+        ctx = mp.get_context(start_method)
+        buffer_ctx = ctx if start_method == 'spawn' else None
         buffer = ReplayBuffer(data_spec=spec,
                               num_environments=1,
                               max_length=10,
-                              device='cpu')
+                              device='cpu',
+                              mp_context=buffer_ctx)
 
-        original_start_method = mp.get_start_method()
-        mp.set_start_method(start_method, force=True)
-        p = mp.Process(target=_write_to_buffer, args=(buffer, ))
+        p = ctx.Process(target=_write_to_buffer, args=(buffer, ))
         p.start()
         p.join()
-        mp.set_start_method(original_start_method, force=True)
 
         if start_method == 'spawn':
             # The buffer in the main process is also changed

alf/trainers/evaluator.py

@@ -14,7 +14,6 @@
 
 from absl import logging
 from absl import flags
-import ctypes
 import math
 import torch.multiprocessing as mp
 import os
@@ -42,30 +41,6 @@
                      defaults=[None] * 4)
 
 
-def _allow_child_to_ptrace(child_pid: int) -> None:
-    """
-    In the *parent* process: allow a given child to ptrace us.
-    This relaxes Yama's ptrace restriction for this specific relationship.
-
-    Requires: kernel.yama.ptrace_scope <= 1 (default on many distros).
-    The value of kernel.yama.ptrace_scope can be checked with:
-        cat /proc/sys/kernel/yama/ptrace_scope
-    """
-    libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
-
-    # prctl constants from <linux/prctl.h>:
-    PR_SET_PTRACER = 0x59616D61
-
-    def _check(ret, what):
-        if ret != 0:
-            err = ctypes.get_errno()
-            raise OSError(err, f"{what} failed: {os.strerror(err)}")
-
-    # Whitelist the specific child PID
-    _check(libc.prctl(PR_SET_PTRACER, ctypes.c_ulong(int(child_pid)), 0, 0, 0),
-           "PR_SET_PTRACER")
-
-
 class Evaluator(object):
     """Evaluator for performing evaluation on the current algorithm.
 
@@ -102,7 +77,7 @@ def __init__(self, config: TrainerConfig, conf_file: str):
             self._worker.start()
 
             # Need this to avoid "pidfd_getfd: Operation not permitted"
-            _allow_child_to_ptrace(self._worker.pid)
+            common.allow_child_to_ptrace(self._worker.pid)
         else:
             self._env = create_environment(for_evaluation=True,
                                            num_parallel_environments=num_envs,

alf/utils/common.py

@@ -17,6 +17,8 @@
 from absl import logging
 import contextlib
 import copy
+import ctypes
+import ctypes.util
 from fasteners.process_lock import InterProcessLock
 from filelock import FileLock
 from functools import wraps
@@ -95,6 +97,32 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def allow_child_to_ptrace(child_pid: int) -> None:
+    """In the *parent* process: allow a given child to ptrace us.
+
+    This relaxes Yama's ptrace restriction for this specific relationship.
+
+    Requires: ``kernel.yama.ptrace_scope <= 1`` (default on many distros).
+    The value of ``kernel.yama.ptrace_scope`` can be checked with::
+
+        cat /proc/sys/kernel/yama/ptrace_scope
+    """
+
+    libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+
+    # prctl constants from <linux/prctl.h>
+    PR_SET_PTRACER = 0x59616D61
+
+    def _check(ret, what):
+        if ret != 0:
+            err = ctypes.get_errno()
+            raise OSError(err, f"{what} failed: {os.strerror(err)}")
+
+    _check(
+        libc.prctl(PR_SET_PTRACER, ctypes.c_ulong(int(child_pid)), 0, 0, 0),
+        "PR_SET_PTRACER")
+
+
 def as_list(x):
     """Convert ``x`` to a list.
 

alf/utils/common_test.py

@@ -21,7 +21,6 @@
 import torch.nn as nn
 
 import alf
-from alf.trainers.evaluator import _allow_child_to_ptrace
 import alf.utils.common as common
 
 
@@ -97,6 +96,13 @@ def _test_worker(m_):
         m_.z[:] = 1.
 
 
+def _launch_worker_with_ctx(ctx, module):
+    """Launch a child process to mutate ``module`` using the provided context."""
+    process = ctx.Process(target=_test_worker, args=(module, ))
+    process.start()
+    process.join()
+
+
 def _test_tensor_sharing():
     """This function is for testing whether tensors are automatically moved
     to shared memory, even when ``Module.share_memory()`` or ``Tensor.share_memory_()``
@@ -116,26 +122,40 @@ def _test_tensor_sharing():
 
     ctx = mp.get_context('spawn')
     # Change ``m`` in the child process
-    process = ctx.Process(target=_test_worker, args=(m, ))
-    process.start()
-    _allow_child_to_ptrace(process.pid)
-    process.join()
+    _launch_worker_with_ctx(ctx, m)
 
     # numpy array should not be modified
     assert np.all(m.y == np.zeros([2]))
     # cuda tensor should be modified
     if torch.cuda.is_available():
         assert torch.all(m.z.cpu() == torch.ones([2]).cpu())
-    # check that ``m``'s tensor also been modified in the parent process
-    assert m.x.is_shared() and torch.all(m.x == torch.ones([2]).cpu()), (
-        "Your pytorch version has a different behavior of sharing CPU tensors "
-        "between processes. Please report the version to the ALF team.")
+    # check whether ``m``'s tensor has been modified in the parent process
+    auto_shared = m.x.is_shared() and torch.all(m.x == torch.ones([2]).cpu())
+    if auto_shared:
+        return True
+
+    logging.fatal(
+        "CPU tensors are not automatically shared between processes on this "
+        "PyTorch build. Need to fall back to explicit share_memory checks.")
+
+    # Explicitly share the module and verify the behaviour still works.
+    m_explicit = _TestModule()
+    m_explicit.share_memory()
+    _launch_worker_with_ctx(ctx, m_explicit)
+    assert torch.all(m_explicit.x == torch.ones([2]).cpu()), (
+        "Explicit share_memory() failed to propagate tensor updates between "
+        "processes.")
+    return False
 
 
 class TensorSharingTest(alf.test.TestCase):
 
     def test_tensor_sharing(self):
-        _test_tensor_sharing()
+        auto_shared = _test_tensor_sharing()
+        if not auto_shared:
+            self.skipTest(
+                "Automatic CPU tensor sharing is disabled on this PyTorch "
+                f"version ({torch.__version__}).")
 
 
 if __name__ == '__main__':

alf/utils/data_buffer.py

@@ -58,9 +58,9 @@ def atomic_wrapper(self, *args, **kwargs):
 class RingBuffer(nn.Module):
     """Batched Ring Buffer.
 
-    Multiprocessing safe, optionally via: ``allow_multiprocess`` flag, blocking
-    modes to ``enqueue`` and ``dequeue``, a stop event to terminate blocked
-    processes, and putting buffer into shared memory.
+    Multiprocessing safe, optionally via providing ``mp_context`` to enable
+    blocking modes to ``enqueue`` and ``dequeue``, a stop event to terminate
+    blocked processes, and putting buffer into shared memory.
 
     This is the underlying implementation of ``ReplayBuffer`` and ``Queue``.
 
@@ -221,7 +221,7 @@ def enqueue(self, batch, env_ids=None, blocking=False):
         """
         if blocking:
             assert self._allow_multiprocess, (
-                "Set allow_multiprocess to enable blocking mode.")
+                "Pass a multiprocessing context to enable blocking mode.")
             env_ids = self.check_convert_env_ids(env_ids)
             while not self._stop.is_set():
                 with self._lock:
@@ -333,7 +333,8 @@ def dequeue(self, env_ids=None, n=1, blocking=False):
         assert n <= self._max_length
         if blocking:
             assert self._allow_multiprocess, [
-                "Set allow_multiprocess", "to enable blocking mode."
+                "Pass a multiprocessing context",
+                "to enable blocking mode."
             ]
             env_ids = self.check_convert_env_ids(env_ids)
             while not self._stop.is_set():
@@ -487,7 +488,6 @@ def __init__(self,
                          num_environments=1,
                          max_length=capacity,
                          device=device,
-                         allow_multiprocess=False,
                          name=name)
         self._capacity = torch.as_tensor(self._max_length,
                                          dtype=torch.int64,

alf/utils/data_buffer_test.py

@@ -64,6 +64,25 @@ def get_batch(env_ids, dim, t, x):
                     reward=r)
 
 
+def _enqueue_after_delay(ring_buffer, batch, delay=0.04):
+    alf.set_default_device("cpu")
+    sleep(delay)
+    ring_buffer.enqueue(batch, batch.env_id)
+
+
+def _dequeue_after_delay(ring_buffer):
+    # cpu tensor on subprocess. Otherwise, spawn method is needed.
+    alf.set_default_device("cpu")
+    sleep(0.04)
+    ring_buffer.dequeue()  # 6(deleted), 7, 8, 9
+    sleep(0.04)  # 10, 7, 8, 9
+    ring_buffer.dequeue()  # 10, 7(deleted), 8, 9
+
+
+def _blocking_dequeue(ring_buffer):
+    ring_buffer.dequeue(blocking=True)
+
+
 class RingBufferTest(parameterized.TestCase, alf.test.TestCase):
     dim = 20
     max_length = 4
@@ -86,15 +105,16 @@ def __init__(self, *args):
         ('test_sync', False),
         ('test_async', True),
     ])
-    def test_ring_buffer(self, allow_multiprocess):
+    def test_ring_buffer(self, use_mp_context):
+        mp_ctx = mp.get_context('spawn') if use_mp_context else None
         ring_buffer = RingBuffer(data_spec=self.data_spec,
                                  num_environments=self.num_envs,
                                  max_length=self.max_length,
-                                 allow_multiprocess=allow_multiprocess)
+                                 mp_context=mp_ctx)
 
         batch1 = get_batch([1, 2, 3, 5, 6], self.dim, t=1, x=0.4)
-        if not allow_multiprocess:
-            # enqueue: blocking mode only available under allow_multiprocess
+        if not use_mp_context:
+            # enqueue: blocking mode only available when multiprocessing is enabled
             self.assertRaises(AssertionError,
                               ring_buffer.enqueue,
                               batch1,
@@ -107,8 +127,8 @@ def test_ring_buffer(self, allow_multiprocess):
             # test that the created batch has gradients
             self.assertTrue(batch1.x.requires_grad)
             ring_buffer.enqueue(batch1, batch1.env_id)
-        if not allow_multiprocess:
-            # dequeue: blocking mode only available under allow_multiprocess
+        if not use_mp_context:
+            # dequeue: blocking mode only available when multiprocessing is enabled
             self.assertRaises(AssertionError,
                               ring_buffer.dequeue,
                               env_ids=batch1.env_id,
@@ -159,17 +179,13 @@ def test_ring_buffer(self, allow_multiprocess):
         ring_buffer.remove_up_to(3)
         self.assertEqual(ring_buffer._current_size, torch.tensor([0] * 8))
 
-        if allow_multiprocess:
+        if use_mp_context:
             # Test block on dequeue without enough data
-            def delayed_enqueue(ring_buffer, batch):
-                alf.set_default_device("cpu")
-                sleep(0.04)
-                ring_buffer.enqueue(batch, batch.env_id)
-
-            p = mp.Process(target=delayed_enqueue,
-                           args=(ring_buffer,
-                                 alf.nest.map_structure(
-                                     lambda x: x.cpu(), batch1)))
+            batch1_cpu = alf.nest.map_structure(
+                lambda tensor: tensor.detach().cpu()
+                if isinstance(tensor, torch.Tensor) else tensor, batch1)
+            p = mp_ctx.Process(target=_enqueue_after_delay,
+                               args=(ring_buffer, batch1_cpu))
             p.start()
             batch = ring_buffer.dequeue(env_ids=batch1.env_id, blocking=True)
             self.assertEqual(batch.step_type, torch.tensor([[9]] * 5))
@@ -180,26 +196,16 @@ def delayed_enqueue(ring_buffer, batch):
                 batch2 = get_batch(range(0, 8), self.dim, t=t, x=0.4)
                 ring_buffer.enqueue(batch2)
 
-            def delayed_dequeue():
-                # cpu tensor on subprocess.  Otherwise, spawn method is needed.
-                alf.set_default_device("cpu")
-                sleep(0.04)
-                ring_buffer.dequeue()  # 6(deleted), 7, 8, 9
-                sleep(0.04)  # 10, 7, 8, 9
-                ring_buffer.dequeue()  # 10, 7(deleted), 8, 9
-
-            p = mp.Process(target=delayed_dequeue)
+            p = mp_ctx.Process(target=_dequeue_after_delay,
+                               args=(ring_buffer, ))
             p.start()
             batch2 = get_batch(range(0, 8), self.dim, t=10, x=0.4)
             ring_buffer.enqueue(batch2, blocking=True)
             p.join()
             self.assertEqual(ring_buffer._current_size[0], torch.tensor(3))
 
             # Test stop queue event
-            def blocking_dequeue(ring_buffer):
-                ring_buffer.dequeue(blocking=True)
-
-            p = mp.Process(target=blocking_dequeue, args=(ring_buffer, ))
+            p = mp_ctx.Process(target=_blocking_dequeue, args=(ring_buffer, ))
             ring_buffer.clear()
             p.start()
             sleep(0.02)  # for subprocess to enter while loop