"shared_cache" in rebuild_cuda_tensor seems to have trouble

92hyungjun: <p>Hello, I’m trying to share a same Tensor between two Processes.<br>
In this case, the Tensor will be allocated in shared cache pool to be able to get accessed from both.</p>
<pre><code class="lang-auto">import torch
import torch.multiprocessing as mp

def sub_process(queue):
    tensor = queue.get()
    # do something...
    tensor = None

    tensor = queue.get()
    # do something...
    tensor = None

if __name__ == '__main__':
    mp.set_start_method('spawn')
    recv_queue = mp.Queue()
    t1 = torch.rand(2, 2, device="cuda")
    p = mp.Process(name="sub_process",target=sub_process, args=(recv_queue, ))
    p.start()

    recv_queue.put(t1)
    # update t1 ... 
    recv_queue.put(t1)
    p.join()

</code></pre>
<p>What I expected was, when I call the second <code>queue.get()</code> function in sub_process, it would take far less time than the first <code>queue.get()</code> since it will allocate the tensor in cached memory pool using storage_from_cache.</p>
<p>However, I found there exists a cached storage in shared_cache (line 298 in reductions.py), but when _new_with_weak_ptr is called to use the storage, it returns None.</p>
<p>The storage cache works as expected only when I control the object manually without calling <code>tensor = None</code>.</p>
<p>Should I handle all objects one by one to use the shared cache pool, or there is a better way?</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L298">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L298" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L298" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L298</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="288" style="counter-reset: li-counter 287 ;">
          <li>
          </li>
<li>def fd_id(fd):</li>
          <li>    # Returns a tuple which uniquely identifies a file descriptor. In Mac OS,</li>
          <li>    # this doesn't work with shared memory handles, which is why we don't</li>
          <li>    # support the "file_descriptor" sharing method on that platform.</li>
          <li>    stat = os.fstat(fd)</li>
          <li>    return (stat.st_ino, stat.st_dev)</li>
          <li>
          </li>
<li>
          </li>
<li>def storage_from_cache(cls, key):</li>
          <li class="selected">    storage_ref = shared_cache.get(key)</li>
          <li>    if storage_ref is None:</li>
          <li>        return None</li>
          <li>    return torch.UntypedStorage._new_with_weak_ptr(storage_ref.cdata)</li>
          <li>
          </li>
<li>
          </li>
<li>def rebuild_storage_fd(cls, df, size):</li>
          <li>    fd = df.detach()</li>
          <li>    try:</li>
          <li>        storage = storage_from_cache(cls, fd_id(fd))</li>
          <li>        if storage is not None:</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>–<br>
While debugging, I came up with another question.</p>
<p>It seems like Senders also make the Storage Cache,<br>
I don’t think I’m understanding the purpose of that code, since I could not find where it actually uses that Storage Cache; all uses Storage Cache from Receivers.</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L261">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L261" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L261" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/7df0878b9936961cc1bde9d20c834ac4331d140a/torch/multiprocessing/reductions.py#L261</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="251" style="counter-reset: li-counter 250 ;">
          <li>if storage.is_cuda:</li>
          <li>    (device,</li>
          <li>     handle,</li>
          <li>     storage_size_bytes,</li>
          <li>     storage_offset_bytes,</li>
          <li>     ref_counter_handle,</li>
          <li>     ref_counter_offset,</li>
          <li>     event_handle,</li>
          <li>     event_sync_required) = storage._share_cuda_()</li>
          <li>    tensor_offset = tensor.storage_offset()</li>
          <li class="selected">    shared_cache[handle] = StorageWeakRef(storage)</li>
          <li>    # _backward_hooks purposely omitted here, see</li>
          <li>    # Note [Don't serialize hooks]</li>
          <li>    return (rebuild_cuda_tensor,</li>
          <li>            (type(tensor),</li>
          <li>             tensor.size(),</li>
          <li>             tensor.stride(),</li>
          <li>             tensor_offset,  # tensor offset in its storage</li>
          <li>             type(storage),</li>
          <li>             tensor.dtype,</li>
          <li>             device,</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>