diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 1ce3998ae47a..5456ff7cf74e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -2267,12 +2267,24 @@ def _decode_video_rows( else: from concurrent.futures import wait + torch = sys.modules.get("torch") + if torch is not None: + grad_enabled = torch.is_grad_enabled() + inference_enabled = torch.is_inference_mode_enabled() + + def decode(collator, rows): + if torch is None: + return collator(rows) + with torch.inference_mode(inference_enabled), \ + torch.set_grad_enabled(grad_enabled): + return collator(rows) + pool = decode_pool if decode_pool is not None else _VideoDecodePool() futures = [] try: executor = pool.get(decode_parallelism) for collator, _, rows in tasks: - futures.append(executor.submit(collator, rows)) + futures.append(executor.submit(decode, collator, rows)) decoded_groups = [future.result() for future in futures] finally: for future in futures: diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index acd87cccb0c0..7d48264a3a88 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -273,6 +273,35 @@ def __call__(self, rows): self.assertEqual([8], worker_counts) + def test_video_row_fallback_preserves_torch_mode(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + from pypaimon.multimodal.video import _VideoDecodePool + + class Collator: + video_column = "camera" + output_column = "frame" + + def __call__(self, rows): + mode = (torch.is_grad_enabled(), + torch.is_inference_mode_enabled()) + return [dict(row, frame=mode) for row in rows] + + pool = _VideoDecodePool() + try: + for context, expected in ( + (torch.enable_grad(), (True, False)), + (torch.no_grad(), (False, False)), + (torch.inference_mode(), (False, True))): + with self.subTest(expected=expected), context: + rows = {0: {"camera": 1}} + _decode_video_rows([rows], [Collator()], decode_pool=pool) + self.assertEqual(expected, rows[0]["frame"]) + finally: + pool.close() + @unittest.skipUnless( av is not None and importlib.util.find_spec("torch") is not None, "PyAV and Torch are required for video decoding",