[BUG] Data collection error on M1 Macs

[itwasabhi]

Describe the bug

SyncDataCollector RuntimeError on Mac-M1 when using mps device.

To Reproduce

from pettingzoo.butterfly import pistonball_v6
from torchrl.collectors import SyncDataCollector
from torchrl.envs import ParallelEnv
from torchrl.envs.libs.pettingzoo import PettingZooWrapper


DEVICE = "mps"

kwargs = {"n_pistons": 21, "continuous": True}
inner_env_fn = lambda: PettingZooWrapper(
    env=pistonball_v6.parallel_env(**kwargs),
    return_state=True,
    group_map=None,
    device=DEVICE,
)
collector_env = lambda: ParallelEnv(num_workers=2, create_env_fn=inner_env_fn)
collector = SyncDataCollector(collector_env, frames_per_batch=100)

collector.reset()
print(next(iter(collector)))
collector.shutdown()

  File ".../.venv/lib/python3.11/site-packages/torch/storage.py", line 447, in wrapper
    return fn(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../.venv/lib/python3.11/site-packages/torch/storage.py", line 526, in _share_filename_cpu_
    return super()._share_filename_cpu_(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: _share_filename_: only available on CPU

Expected behavior

Should be able to create env on mps and collect data.

Screenshots

If applicable, add screenshots to help explain your problem.

System info

>>> print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)
0.8.0+83cb41d 2.0.2 3.11.8 (main, Feb 25 2024, 03:55:37) [Clang 17.0.6 ] darwin

>>> platform.platform()
'macOS-15.5-arm64-arm-64bit'

Running via uv

Additional context

• No issue when collecting directly on cpu or cuda
• No issue when using SerialEnv over ParallelEnv

Reason and Possible fixes

Seems related to pytorch/pytorch#87688

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[Kushagra481]

Hey can I work on this

[vmoens]

Hello

That's kind of expected - even though I agree it's not thrilling.
You don't have shared memory with MPS tensors. The solution is to not use a shared buffer (use_buffers=False).

Here's what I would suggest as a way forward:

• if mps is asked for the storage device and you don't specify if you want to use buffers or not, we warn you that we won't be using buffers and the whole thing can work (we might need to dump the data on cpu mem to make it transition from process to process)
• if you ask for mps storage and ask for buffers we error out.

If @Kushagra481 you want to give it a stab be my guest!

[Kushagra481]

I’ll fork the repo, implement the MPS check in SyncDataCollector, and open a PR with tests and a warning message as discussed. Thanks for the guidance.

[Kushagra481]

import warnings
import torch
from typing import Optional, Union, Callable, Any

class SyncDataCollector:
def init(
self,
create_env_fn: Callable[[], Any],
policy: Optional[Any] = None,
*,
frames_per_batch: int = 200,
total_frames: int = -1,
device: Optional[Union[torch.device, str]] = None,
storing_device: Optional[Union[torch.device, str]] = None,
use_buffers: Optional[bool] = None,
max_frames_per_traj: int = -1,
init_random_frames: int = -1,
reset_at_each_iter: bool = False,
postproc: Optional[Any] = None,
split_trajs: bool = False,
exploration_type: str = "random",
collector_class: Optional[type] = None,
collector_kwargs: Optional[dict] = None,
num_workers: int = 1,
sync: bool = True,
slicing: int = 1,
update_at_each_batch: bool = False,
preemptive_threshold: float = 0.0,
norm_stats: bool = False,
**kwargs
):
if device is not None:
device = torch.device(device)
if storing_device is not None:
storing_device = torch.device(storing_device)

    self.use_buffers = self._handle_mps_device_compatibility(device, storing_device, use_buffers)

    self.create_env_fn = create_env_fn
    self.policy = policy
    self.frames_per_batch = frames_per_batch
    self.total_frames = total_frames
    self.device = device
    self.storing_device = storing_device
    self.max_frames_per_traj = max_frames_per_traj
    self.init_random_frames = init_random_frames
    self.reset_at_each_iter = reset_at_each_iter
    self.postproc = postproc
    self.split_trajs = split_trajs
    self.exploration_type = exploration_type
    self.collector_class = collector_class
    self.collector_kwargs = collector_kwargs or {}
    self.num_workers = num_workers
    self.sync = sync
    self.slicing = slicing
    self.update_at_each_batch = update_at_each_batch
    self.preemptive_threshold = preemptive_threshold
    self.norm_stats = norm_stats

    self._init_collector()

def _handle_mps_device_compatibility(
    self,
    device: Optional[torch.device],
    storing_device: Optional[torch.device],
    use_buffers: Optional[bool]
) -> bool:
    is_mps_device = self._is_mps_device(device)
    is_mps_storing = self._is_mps_device(storing_device)

    if is_mps_device or is_mps_storing:
        device_name = "device" if is_mps_device else "storing_device"

        if use_buffers is True:
            raise RuntimeError(
                f"Cannot use shared memory buffers with MPS {device_name}. "
                f"MPS tensors do not support shared memory operations."
            )
        elif use_buffers is None:
            warnings.warn(
                f"MPS {device_name} detected. Shared memory buffers are not "
                f"supported with MPS tensors. Automatically setting use_buffers=False.",
                UserWarning,
                stacklevel=3
            )
            use_buffers = False

    return use_buffers

def _is_mps_device(self, device: Optional[torch.device]) -> bool:
    if device is None:
        return False
    if isinstance(device, str):
        return device.startswith('mps')
    return device.type == 'mps'

def _init_collector(self):
    pass

def _create_shared_buffer(self, *args, **kwargs):
    if not self.use_buffers:
        return None
    if self._is_mps_device(self.device) or self._is_mps_device(self.storing_device):
        raise RuntimeError(
            "Internal error: Attempted to create shared buffer with MPS device."
        )
    return None

def collect(self, *args, **kwargs):
    pass

def reset(self):
    pass

def shutdown(self):
    pass

def detect_mps_in_env(env_fn: Callable) -> bool:
try:
temp_env = env_fn()
if hasattr(temp_env, 'device'):
return temp_env.device.type == 'mps'
return False
except Exception:
return False
finally:
if 'temp_env' in locals():
try:
temp_env.close()
except:
pass
Hey bro if you could check this out before I generate a pull request @vmoens

[vmoens]

I think we should go over the env spec leaves and check their device instead of looking at the env device

[Kushagra481]

import warnings
import torch
from typing import Optional, Union, Callable, Any

class SyncDataCollector:
def init(
self,
create_env_fn: Callable[[], Any],
policy: Optional[Any] = None,
*,
frames_per_batch: int = 200,
total_frames: int = -1,
device: Optional[Union[torch.device, str]] = None,
storing_device: Optional[Union[torch.device, str]] = None,
use_buffers: Optional[bool] = None,
max_frames_per_traj: int = -1,
init_random_frames: int = -1,
reset_at_each_iter: bool = False,
postproc: Optional[Any] = None,
split_trajs: bool = False,
exploration_type: str = "random",
collector_class: Optional[type] = None,
collector_kwargs: Optional[dict] = None,
num_workers: int = 1,
sync: bool = True,
slicing: int = 1,
update_at_each_batch: bool = False,
preemptive_threshold: float = 0.0,
norm_stats: bool = False,
**kwargs
):
if device is not None:
device = torch.device(device)
if storing_device is not None:
storing_device = torch.device(storing_device)
self.use_buffers = self._handle_mps_device_compatibility(
device, storing_device, use_buffers, create_env_fn
)
self.create_env_fn = create_env_fn
self.policy = policy
self.frames_per_batch = frames_per_batch
self.total_frames = total_frames
self.device = device
self.storing_device = storing_device
self.max_frames_per_traj = max_frames_per_traj
self.init_random_frames = init_random_frames
self.reset_at_each_iter = reset_at_each_iter
self.postproc = postproc
self.split_trajs = split_trajs
self.exploration_type = exploration_type
self.collector_class = collector_class
self.collector_kwargs = collector_kwargs or {}
self.num_workers = num_workers
self.sync = sync
self.slicing = slicing
self.update_at_each_batch = update_at_each_batch
self.preemptive_threshold = preemptive_threshold
self.norm_stats = norm_stats
self._init_collector()

def _detect_mps_in_env_specs(self, env_fn: Callable) -> bool:
    temp_env = None
    try:
        temp_env = env_fn()
        spec_attrs = ['observation_spec', 'action_spec', 'reward_spec', 'done_spec']
        for attr_name in spec_attrs:
            if hasattr(temp_env, attr_name):
                spec = getattr(temp_env, attr_name)
                if callable(spec):
                    spec = spec()
                if self._check_spec_for_mps(spec):
                    return True
        return False
    except Exception:
        return False
    finally:
        if temp_env is not None:
            try:
                temp_env.close()
            except:
                pass

def _check_spec_for_mps(self, spec) -> bool:
    if spec is None:
        return False
    if hasattr(spec, 'device'):
        if self._is_mps_device(spec.device):
            return True
    if hasattr(spec, 'values') and callable(spec.values):
        try:
            for sub_spec in spec.values():
                if self._check_spec_for_mps(sub_spec):
                    return True
        except:
            pass
    if hasattr(spec, 'items') and callable(spec.items):
        try:
            for key, sub_spec in spec.items():
                if self._check_spec_for_mps(sub_spec):
                    return True
        except:
            pass
    if hasattr(spec, '_specs'):
        if self._check_spec_for_mps(spec._specs):
            return True
    try:
        if hasattr(spec, '__iter__') and not isinstance(spec, (str, bytes)):
            for sub_spec in spec:
                if self._check_spec_for_mps(sub_spec):
                    return True
    except (TypeError, AttributeError):
        pass
    return False

def _handle_mps_device_compatibility(
    self,
    device: Optional[torch.device],
    storing_device: Optional[torch.device],
    use_buffers: Optional[bool],
    create_env_fn: Callable[[], Any]
) -> bool:
    is_mps_device = self._is_mps_device(device)
    is_mps_storing = self._is_mps_device(storing_device)
    env_uses_mps = self._detect_mps_in_env_specs(create_env_fn)
    if is_mps_device or is_mps_storing or env_uses_mps:
        mps_sources = []
        if is_mps_device:
            mps_sources.append("device")
        if is_mps_storing:
            mps_sources.append("storing_device")
        if env_uses_mps:
            mps_sources.append("environment specs")
        device_name = " and ".join(mps_sources)
        if use_buffers is True:
            raise RuntimeError(
                f"Cannot use shared memory buffers with MPS {device_name}. MPS tensors do not support shared memory operations."
            )
        elif use_buffers is None:
            warnings.warn(
                f"MPS {device_name} detected. Shared memory buffers are not supported with MPS tensors. Automatically setting use_buffers=False.",
                UserWarning,
                stacklevel=3
            )
            use_buffers = False
    return use_buffers

def _is_mps_device(self, device: Optional[Union[torch.device, str]]) -> bool:
    if device is None:
        return False
    if isinstance(device, str):
        return device.startswith('mps')
    return device.type == 'mps'

def _init_collector(self):
    pass

def _create_shared_buffer(self, *args, **kwargs):
    if not self.use_buffers:
        return None
    if (
        self._is_mps_device(self.device)
        or self._is_mps_device(self.storing_device)
        or self._detect_mps_in_env_specs(self.create_env_fn)
    ):
        raise RuntimeError(
            "Internal error: Attempted to create shared buffer with MPS device."
        )
    return None

def collect(self, *args, **kwargs):
    pass

def reset(self):
    pass

def shutdown(self):
    pass

@vmoens I implemented the changes you suggested please check so i can proceed and thanks for the guidance once again

[vmoens]

I would rather not copy paste code in gh issues :)
Feel free to open a PR!

[discobot]

Dug into this: use_buffers=False on its own is not enough -- every pipe send of an
MPS tensordict crashes in the same mp reduction (reduce_storage ->
_share_filename_cpu_()), and consolidate(device="cpu") doesn't help because the
consolidation metadata pickles the original device, so the receiver skips the cast
back. The data has to be staged on CPU before consolidation on both sides of the pipe.
Fix in #3867: the spec-leaf device check with the warn/raise behavior described above,
plus CPU staging for the no-buffers pipe traffic.