ModelCheckpoint with custom directory path not working on TPU

🐛 Bug

Saving a checkpoint to a custom directory on Colab with a TPU runtime fails using ModelCheckpoint(dirpath="/tmp/checkpoints"). When using ModelCheckpoint(dirpath="/tmp") or simply ModelCheckpoint() the checkpoint is saved correctly

To Reproduce

The code is based on the TPU tutorial:

Package installation:

!pip install torch==1.9.1 torchtext==0.10.1 torchvision==0.10.1 pytorch-lightning==1.5.8 cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl

Code:

import os
from typing import Any, List, Optional

import torch
import torch.nn.functional as F
from pytorch_lightning import LightningDataModule, LightningModule, Trainer
from pytorch_lightning.callbacks import BasePredictionWriter, ModelCheckpoint
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchmetrics.functional import accuracy
from torchvision import transforms
from torchvision.datasets import MNIST

BATCH_SIZE = 1024


class MNISTDataModule(LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

        # self.dims is returned when you call dm.size()
        # Setting default dims here because we know them.
        # Could optionally be assigned dynamically in dm.setup()
        self.dims = (1, 28, 28)
        self.num_classes = 10

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
        self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=BATCH_SIZE)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=BATCH_SIZE)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=BATCH_SIZE)


class LitModel(LightningModule):
    def __init__(self, channels, width, height, num_classes, hidden_size=64, learning_rate=2e-4):
        super().__init__()

        self.save_hyperparameters()

        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(channels * width * height, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, num_classes),
        )

    def forward(self, x):
        x = self.model(x)
        return F.log_softmax(x, dim=1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        return optimizer

dm = MNISTDataModule()
model = LitModel(*dm.size(), dm.num_classes)
ckpt_callback = ModelCheckpoint(dirpath="/tmp/checkpoints")
trainer = Trainer(max_epochs=3, progress_bar_refresh_rate=20, tpu_cores=8, callbacks=[ckpt_callback])
trainer.fit(model, dm)

Error log extract:

Exception in device=TPU:0: [Errno 2] No such file or directory: '/tmp/checkpoints/epoch=0-step=6.ckpt'
...
FileNotFoundError: [Errno 2] No such file or directory: '/tmp/checkpoints/epoch=0-step=6.ckpt'
Exception in device=TPU:2: tensorflow/compiler/xla/xla_client/mesh_service.cc:364 : Failed to meet rendezvous 'torch_xla.core.xla_model.save': Socket closed (14)
Exception in device=TPU:4: tensorflow/compiler/xla/xla_client/mesh_service.cc:364 : Failed to meet rendezvous 'torch_xla.core.xla_model.save': Socket closed (14)
...

Expected behavior

The checkpoint should be saved correctly to the specified directory.

Environment

Colab with TPU runtime.

• CUDA:
  • GPU:
  • available: False
  • version: 10.2
• Packages:
  • numpy: 1.19.5
  • pyTorch_debug: False
  • pyTorch_version: 1.9.1+cu102
  • pytorch-lightning: 1.5.8
  • tqdm: 4.62.3
• System:
  • OS: Linux
  • architecture:
    • 64bit
  • processor: x86_64
  • python: 3.7.12
  • version: #1 SMP Tue Dec 7 09:58:10 PST 2021

Start with the ModelCheckpoint dirpath handling and the Trainer TPU setup shown in the reproduction, then run the provided Colab example with /tmp/checkpoints. Compare it with /tmp and the default path; done means checkpoints save successfully in the specified custom directory without the reported FileNotFoundError or TPU rendezvous failures.