import sys | |
from lightning_utilities.core.imports import RequirementCache | |
__all__ = [] | |
if not RequirementCache("litdata"): | |
raise ModuleNotFoundError("Please, run `pip install litdata`") # E111 | |
else: | |
import litdata | |
# Enable resolution at least for lower data namespace | |
sys.modules["lightning.data"] = litdata | |
from litdata.processing.functions import map, optimize, walk | |
from litdata.streaming.combined import CombinedStreamingDataset | |
from litdata.streaming.dataloader import StreamingDataLoader | |
from litdata.streaming.dataset import StreamingDataset | |
__all__ += [ | |
"LightningDataset", | |
"StreamingDataset", | |
"CombinedStreamingDataset", | |
"StreamingDataLoader", | |
"LightningIterableDataset", | |
"map", | |
"optimize", | |
"walk", | |
] | |
# TODO: Move this to litdata | |
if RequirementCache("lightning_sdk"): | |
from lightning_sdk import Machine | |
__all__ += ["Machine"] | |