Tensor sharing across multiple venvs

username: guill
cooked: <h2><a name="p-472712-context-1" class="anchor" href="#p-472712-context-1"></a>Context</h2>
<p>Hello! I’m one of the developers working on <a href="https://github.com/comfyanonymous/ComfyUI" rel="noopener nofollow ugc">ComfyUI</a>. It’s a tool built for modular execution of AI workflows via a node graph system. One category of issues we encounter a lot is “custom nodes” with conflicting dependencies (e.g. one requiring numpy 1.x while another needs numpy 2.x).</p>
<p>One solution we’re exploring is allowing Python extensions to run in isolated virtual environments while sharing PyTorch tensors with the host process via PyTorch’s <code>torch.multiprocessing</code> support. This enables running extensions with conflicting dependencies while still allowing efficient tensor sharing for AI workloads.</p>
<p>To this end, we’ve developed a library named <a href="https://github.com/Comfy-Org/pyisolate" rel="noopener nofollow ugc">pyisolate</a></p>
<p>Our approach:</p>
<ol>
<li>Create separate venvs for different sets of dependencies</li>
<li>Use <code>multiprocessing.set_executable()</code> to spawn processes using each venv’s Python interpreter. (It’s a little more complicated for Windows.)</li>
<li>Use <code>torch.multiprocessing</code> with spawn method to enable tensor sharing between host and extension processes</li>
<li>All processes have the same PyTorch version installed</li>
</ol>
<h2><a name="p-472712-current-status-2" class="anchor" href="#p-472712-current-status-2"></a>Current Status</h2>
<p>This approach currently works well - tensors (including CUDA tensors) are successfully shared between the host process and extension processes running with different Python executables. The shared memory mechanisms and file descriptors work correctly across different Python executables.</p>
<h2><a name="p-472712-question-3" class="anchor" href="#p-472712-question-3"></a>Question</h2>
<p><strong>Is this use case (tensor sharing between processes spawned with different Python executables via <code>set_executable()</code>) something PyTorch is likely to stop supporting in the future?</strong></p>
<p>We want to ensure we’re not relying on behavior that is likely break in future versions. While the standard use case is sharing tensors between processes using the same Python executable, our use case extends this to processes using different executables from isolated environments.</p>
<p>Thanks in advance for any thoughts!</p>

username: ezyang
cooked: <p>Is this using our CUDA refcounted IPC stuff? We haven’t really touched it in forever so it’s “stable”, but if the setup there changes then yes it would break across versions. The most reliable thing to do is reimplement the sharing yourself (it’s not that complicated, and you might want to have a better refcounting scheme than what we do) and then use Python APIs to convert it into Tensors.</p>

username: Hamza_Javaid
cooked: <p>Maybe try to isolate and ensure the compatibility something like this</p>
<pre><code class="lang-auto">
def verify_pytorch_compatibility(venv_path):
    """Verify PyTorch versions match across environments"""
    import subprocess
    result = subprocess.run(
        [f"{venv_path}/bin/python", "-c", 
         "import torch; print(torch.__version__)"],
        capture_output=True, text=True
    )
    host_version = torch.__version__
    venv_version = result.stdout.strip()
    
    if host_version != venv_version:
        raise RuntimeError(f"Version mismatch: {host_version} != {venv_version}")
</code></pre>
<p>Afaik, PyTorch’s tensor sharing relies on low-level OS primitives. Other way could be to handle platform specific differences. Tricky problem</p>