fix(test): fix AWS_IO_DNS_INVALID_NAME in multiprocess tests

Add multiprocessing_context=mp.get_context() to DataLoader calls to ensure
spawn method is used instead of fork in Darwin, preventing S3 client fork
handlers from corrupting AWS CRT DNS resolver threads on macOS github runners.

Author: jet-tong

s3torchconnector/tst/e2e/test_distributed_training.py

@@ -79,7 +79,11 @@ def dataloader_for_map(
     )
     sampler = DistributedSampler(dataset)
     dataloader = DataLoader(
-        dataset, batch_size=batch_size, num_workers=num_workers, sampler=sampler
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        sampler=sampler,
+        multiprocessing_context=mp.get_context(),
     )
     return dataloader
 
@@ -93,7 +97,12 @@ def dataloader_for_iterable(
         enable_sharding=True,
         reader_constructor=reader_constructor,
     )
-    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        multiprocessing_context=mp.get_context(),
+    )
     return dataloader
 
 

s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py

@@ -12,6 +12,7 @@
 from lightning.pytorch.demos import WikiText2
 from lightning.pytorch.plugins import AsyncCheckpointIO
 from torch.utils.data import DataLoader
+import torch.multiprocessing as mp
 
 from s3torchconnector import S3Checkpoint
 from s3torchconnector._s3client import S3Client
@@ -79,7 +80,9 @@ def test_delete_checkpoint(checkpoint_directory):
 def test_load_trained_checkpoint(checkpoint_directory):
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
-    dataloader = DataLoader(dataset, num_workers=3)
+    dataloader = DataLoader(
+        dataset, num_workers=3, multiprocessing_context=mp.get_context()
+    )
     model = LightningTransformer(vocab_size=dataset.vocab_size)
     trainer = L.Trainer(accelerator=LIGHTNING_ACCELERATOR, fast_dev_run=2)
     trainer.fit(model=model, train_dataloaders=dataloader)
@@ -95,7 +98,9 @@ def test_load_trained_checkpoint(checkpoint_directory):
 def test_compatibility_with_trainer_plugins(checkpoint_directory):
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
-    dataloader = DataLoader(dataset, num_workers=3)
+    dataloader = DataLoader(
+        dataset, num_workers=3, multiprocessing_context=mp.get_context()
+    )
     model = LightningTransformer(vocab_size=dataset.vocab_size)
     s3_lightning_checkpoint = S3LightningCheckpoint(region=checkpoint_directory.region)
     _verify_user_agent(s3_lightning_checkpoint)
@@ -121,7 +126,9 @@ def test_compatibility_with_trainer_plugins(checkpoint_directory):
 def test_compatibility_with_checkpoint_callback(checkpoint_directory):
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
-    dataloader = DataLoader(dataset, num_workers=3)
+    dataloader = DataLoader(
+        dataset, num_workers=3, multiprocessing_context=mp.get_context()
+    )
 
     model = LightningTransformer(vocab_size=dataset.vocab_size)
     s3_lightning_checkpoint = S3LightningCheckpoint(checkpoint_directory.region)
@@ -161,7 +168,9 @@ def test_compatibility_with_checkpoint_callback(checkpoint_directory):
 def test_compatibility_with_async_checkpoint_io(checkpoint_directory):
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
-    dataloader = DataLoader(dataset, num_workers=3)
+    dataloader = DataLoader(
+        dataset, num_workers=3, multiprocessing_context=mp.get_context()
+    )
 
     model = LightningTransformer(vocab_size=dataset.vocab_size)
     s3_lightning_checkpoint = S3LightningCheckpoint(checkpoint_directory.region)
@@ -192,7 +201,9 @@ def test_compatibility_with_async_checkpoint_io(checkpoint_directory):
 def test_compatibility_with_lightning_checkpoint_load(checkpoint_directory):
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
-    dataloader = DataLoader(dataset, num_workers=3)
+    dataloader = DataLoader(
+        dataset, num_workers=3, multiprocessing_context=mp.get_context()
+    )
     model = LightningTransformer(vocab_size=dataset.vocab_size)
     s3_lightning_checkpoint = S3LightningCheckpoint(region=checkpoint_directory.region)
     trainer = L.Trainer(

s3torchconnector/tst/e2e/test_multiprocess_dataloading.py

@@ -10,6 +10,7 @@
 import pytest
 from torch.utils.data import DataLoader, get_worker_info
 from torchdata.datapipes.iter import IterableWrapper
+import torch.multiprocessing as mp
 
 from s3torchconnector import (
     S3IterableDataset,
@@ -85,7 +86,12 @@ def test_s3iterable_dataset_multiprocess_torchdata(
     batch_size = 2
     num_workers = 3
 
-    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        multiprocessing_context=mp.get_context(),
+    )
 
     total_objects = 0
     uris_seen = Counter()
@@ -123,7 +129,9 @@ def test_s3iterable_dataset_multiprocess(
     num_epochs = 2
     num_images = len(image_directory.contents)
 
-    dataloader = DataLoader(dataset, num_workers=num_workers)
+    dataloader = DataLoader(
+        dataset, num_workers=num_workers, multiprocessing_context=mp.get_context()
+    )
     counter = 0
     for epoch in range(num_epochs):
         s3keys = Counter()
@@ -160,7 +168,9 @@ def test_s3mapdataset_multiprocess(
     num_epochs = 2
     num_images = len(image_directory.contents)
 
-    dataloader = DataLoader(dataset, num_workers=num_workers)
+    dataloader = DataLoader(
+        dataset, num_workers=num_workers, multiprocessing_context=mp.get_context()
+    )
 
     for epoch in range(num_epochs):
         s3keys = Counter()