diff --git a/comet/models/base.py b/comet/models/base.py index 9872d2b7..e3aed878 100644 --- a/comet/models/base.py +++ b/comet/models/base.py @@ -622,7 +622,11 @@ def predict( sampler=sampler, collate_fn=self.prepare_for_inference, num_workers=num_workers, - multiprocessing_context="fork" if torch.backends.mps.is_available() else None, + multiprocessing_context=( + "fork" + if num_workers > 0 and torch.backends.mps.is_available() + else None + ), ) if gpus > 1: pred_writer = CustomWriter() diff --git a/tests/unit/test_predict_dataloader_context.py b/tests/unit/test_predict_dataloader_context.py new file mode 100644 index 00000000..be07261a --- /dev/null +++ b/tests/unit/test_predict_dataloader_context.py @@ -0,0 +1,75 @@ +"""Exercise predict's real DataLoader without models or accelerator hardware.""" + +import os +import unittest +from contextlib import ExitStack +from unittest.mock import Mock, patch + +import torch +from comet.models.base import CometModel +from torch.utils.data import DataLoader + + +@unittest.skipUnless(os.name == "posix", "MPS worker setup uses POSIX fork") +class TestPredictDataLoaderContext(unittest.TestCase): + def test_context_matches_worker_configuration(self): + cases = [ + # MPS availability, GPUs, requested workers, resolved workers, context + (True, 0, None, 0, None), + (True, 1, 0, 0, None), + (True, 0, 2, 2, "fork"), + (True, 1, None, 2, "fork"), + (False, 0, None, 0, None), + (False, 1, 0, 0, None), + (False, 0, 2, 2, None), + (False, 1, None, 2, None), + ] + samples = [{"src": "source", "mt": "translation", "ref": "reference"}] + for mps, gpus, workers, expected_workers, expected_context in cases: + with self.subTest(mps=mps, gpus=gpus, workers=workers): + model = Mock() + with ExitStack() as stack: + stack.enter_context( + patch( + "comet.models.base.torch.backends.mps.is_available", + return_value=mps, + ) + ) + loader_factory = stack.enter_context( + patch("comet.models.base.DataLoader", wraps=DataLoader) + ) + trainer_factory = stack.enter_context( + patch("comet.models.base.ptl.Trainer") + ) + trainer_factory.return_value.predict.return_value = [ + {"scores": torch.tensor([0.75])} + ] + output = CometModel.predict( + model, + samples, + gpus=gpus, + num_workers=workers, + progress_bar=False, + length_batching=False, + ) + + kwargs = loader_factory.call_args.kwargs + self.assertEqual(kwargs["num_workers"], expected_workers) + self.assertEqual( + kwargs["multiprocessing_context"], expected_context + ) + loader = trainer_factory.return_value.predict.call_args.kwargs[ + "dataloaders" + ] + self.assertIsInstance(loader, DataLoader) + context = loader.multiprocessing_context + self.assertEqual( + context.get_start_method() if context else None, + expected_context, + ) + self.assertEqual(output.scores, [0.75]) + self.assertEqual(output.system_score, 0.75) + + +if __name__ == "__main__": + unittest.main()