schedulers: set /dev/shm size for docker based schedulers to bypass 64M default

[d4l3k]

PyTorch dataloaders use /dev/shm to transfer data between processes. The default in docker containers is 64MB so we need to increase this for more complex models.

• DockerScheduler: sets shm_size to mem request
• AWSBatchScheduler: sets sharedMemorySize to mem request
• KubernetesScheduler: mounts an unlimited tmpfs onto /dev/shm

Fixes #428

Test plan:

import torch
from torch.utils.data import Dataset, DataLoader

class BigDataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return 20

    def __getitem__(self, idx):
        return torch.zeros((1,self.size))

dataset = BigDataset(100_000_000)
dataloader = DataLoader(dataset, batch_size=1, num_workers=1)

for i, x in enumerate(dataloader):
    print(i, x.shape)

torchx run --scheduler {local_docker,kubernetes,aws_batch} --wait --log dist.ddp --memMB 2000 -j 1x1 --script large-shm.py

[codecov]

Codecov Report

Merging #429 (090fc11) into main (90b05b0) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #429   +/-   ##
=======================================
  Coverage   94.41%   94.41%           
=======================================
  Files          67       67           
  Lines        3829     3830    +1     
=======================================
+ Hits         3615     3616    +1     
  Misses        214      214           

Impacted Files	Coverage Δ
torchx/schedulers/aws_batch_scheduler.py	84.18% <100.00%> (ø)
torchx/schedulers/docker_scheduler.py	95.74% <100.00%> (ø)
torchx/schedulers/kubernetes_scheduler.py	92.68% <100.00%> (+0.03%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 90b05b0...090fc11. Read the comment docs.

[facebook-github-bot]

@d4l3k has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.