title: How to get rid of zombie processes using torch.multiprocessing.Pool (Python)
body: <p>I am using <code>torch.multiprocessing.Pool</code> to speed up my NN in inference, like this:</p>
  
  <pre><code>import torch.multiprocessing as mp
  mp = torch.multiprocessing.get_context('forkserver')
  
  def parallel_predict(predict_func, sequences, args):
      predicted_cluster_ids = []
      pool = mp.Pool(args.num_workers, maxtasksperchild=1)
      out = pool.imap(
          func=functools.partial(predict_func, args=args),
          iterable=sequences,
          chunksize=1)
      for item in tqdm(out, total=len(sequences), ncols=85):
          predicted_cluster_ids.append(item)
      pool.close()
      pool.terminate()
      pool.join()
      return predicted_cluster_ids
  </code></pre>
  
  <p>Note 1) I am using <code>imap</code> because I want to be able to show a progress bar with <code>tqdm</code>.<br>
  Note 2) I tried with both <code>forkserver</code> and <code>spawn</code> but no luck. I cannot use other methods because of how they interact (poorly) with CUDA.<br>
  Note 3) I am using <code>maxtasksperchild=1</code> and <code>chunksize=1</code> so for each sequence in <code>sequences</code> it spawns a new process.<br>
  Note 4) Adding or removing <code>pool.terminate()</code> and <code>pool.join()</code> makes no difference.<br>
  Note 5) <code>predict_func</code> is a method of a class I created. I could also pass the whole model to <code>parallel_predict</code> but it does not change anything.</p>
  
  <p>Everything works fine except the fact that after a while I run out of memory on the CPU (while on the GPU everything works as expected). Using <code>htop</code> to monitor memory usage I notice that, for every process I spawn with pool I get a zombie that uses 0.4% of the memory. They don't get cleared, so they keep using space. Still, <code>parallel_predict</code> does return the correct result and the computation goes on. My script is structured in a way that id does validation multiple times so next time <code>parallel_predict</code> is called the zombies add up. </p>
  
  <p>This is what I get in <code>htop</code>:
  <a href="https://i.stack.imgur.com/11u1i.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/11u1i.png" alt="enter image description here"></a></p>
  
  <p>Usually, these zombies get cleared after <code>ctrl-c</code> but in some rare cases I need to <code>killall</code>.</p>
  
  <p>Is there some way I can force the <code>Pool</code> to close them?</p>
  
  <p>UPDATE:
  I tried to kill the zombie processes using this:</p>
  
  <pre><code>def kill(pool):
      import multiprocessing
      import signal
      # stop repopulating new child
      pool._state = multiprocessing.pool.TERMINATE
      pool._worker_handler._state = multiprocessing.pool.TERMINATE
      for p in pool._pool:
          os.kill(p.pid, signal.SIGKILL)
      # .is_alive() will reap dead process
      while any(p.is_alive() for p in pool._pool):
          pass
      pool.terminate()
  </code></pre>
  
  <p>But it does not work. It gets stuck at <code>pool.terminate()</code></p>
  
  <p>UPDATE2:
  I tried to use the <code>initializer</code> arg in <code>imap</code> to catch signals like this:</p>
  
  <pre><code>def process_initializer():
      def handler(_signal, frame):
          print('exiting')
          exit(0)
      signal.signal(signal.SIGTERM, handler)
  
  
  def parallel_predict(predict_func, sequences, args):
      predicted_cluster_ids = []
      with mp.Pool(args.num_workers, initializer=process_initializer, maxtasksperchild=1) as pool:
          out = pool.imap(
              func=functools.partial(predict_func, args=args),
              iterable=sequences,
              chunksize=1)
          for item in tqdm(out, total=len(sequences), ncols=85):
              predicted_cluster_ids.append(item)
          for p in pool._pool:
              os.kill(p.pid, signal.SIGTERM)
          pool.close()
          pool.terminate()
          pool.join()
      return predicted_cluster_ids
  </code></pre>
  
  <p>but again it does not free memory.</p>