Memory management in PyTorch implementation of multi-processing queues

mortazavi · August 24, 2020, 7:23pm

I have a question for the PyTorch development team.

How is the memory consumed by queues in PyTorch implementation of multi-processing libraries managed?

If you can point me to the relevant piece of code (if available) and/or provide a textual description, I would appreciate it.

pritamdamania87 · August 25, 2020, 2:25am

@VitalyFedyunin Could you help out here since its a torch.multiprocessing question?

VitalyFedyunin · August 27, 2020, 10:47pm

Please check

github.com

pytorch/pytorch/blob/e75fb4356b752097d093c7013ba85c9eb82961ef/torch/multiprocessing/reductions.py

import torch
import torch.utils.hooks
from torch._namedtensor_internals import check_serializing_named_tensor
import os
import threading
import multiprocessing
from multiprocessing.util import register_after_fork
from multiprocessing.reduction import ForkingPickler
try:
    # Early load resource_sharer to prevent a partially initialized instance
    # from being inherited in a forked child process. The reduce_storage method
    # requires this module indirectly through DupFd(). The built-in mp.Queue
    # class pickles arguments in a background thread which may overlap with the
    # fork.
    import multiprocessing.resource_sharer
except ImportError:
    pass


class StorageWeakRef(object):

This file has been truncated. show original

and

github.com

pytorch/pytorch/blob/cca247635c6edb323176eeac7a18d3e9ab71c558/torch/multiprocessing/queue.py

import io
import multiprocessing
import multiprocessing.queues
from multiprocessing.reduction import ForkingPickler
import pickle


class ConnectionWrapper(object):
    """Proxy class for _multiprocessing.Connection which uses ForkingPickler to
    serialize objects"""

    def __init__(self, conn):
        self.conn = conn

    def send(self, obj):
        buf = io.BytesIO()
        ForkingPickler(buf, pickle.HIGHEST_PROTOCOL).dump(obj)
        self.send_bytes(buf.getvalue())

    def recv(self):

This file has been truncated. show original

as methods are different for CPU/GPU

generally speaking we are passing storage descriptors and do usage ref counting.