fkat 0.1.2__py3-none-any.whl

fkat/__init__.py

@@ -0,0 +1,147 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import sys
+import asyncio
+from collections.abc import Callable
+
+import hydra
+import omegaconf as oc
+import lightning as L
+import torch
+import torch.multiprocessing as mp
+from torch.distributed.elastic.multiprocessing.errors import record
+
+from fkat.utils import config, pdb
+from fkat.utils.logging import rank0_logger
+from fkat.utils.config import SingletonResolver
+
+log = rank0_logger(__name__)
+
+
+def run_main(main: Callable[[], None]) -> None:
+    patch_args()
+
+    @record
+    async def async_main() -> None:
+        try:
+            main()
+        except Exception as e:
+            import traceback
+
+            traceback.print_tb(e.__traceback__)
+            raise e
+
+    asyncio.run(async_main())  # type: ignore[arg-type]
+
+
+def patch_args() -> None:
+    """
+    In case we need to pass wildcard arguments (e.g. overrides) as expected by Hydra,
+    but the runtime only allows named arguments we pass them using a bogus "--overrides" flag.
+    This function will take care of removing this flag by the time we call Hydra.
+    """
+    overrides_pos = -1
+    for i, a in enumerate(sys.argv):
+        if a == "--overrides":
+            overrides_pos = i
+            break
+    if overrides_pos >= 0:
+        overrides = sys.argv[overrides_pos + 1] if overrides_pos + 1 < len(sys.argv) else ""
+        # skipping overrides when constructing new args, there could be more args ahead
+        sys.argv = sys.argv[:overrides_pos] + (
+            sys.argv[overrides_pos + 2 :] if overrides_pos + 2 < len(sys.argv) else []
+        )
+        if overrides:
+            # adding overrides at the end
+            sys.argv.extend(overrides.split(" "))
+
+
+def setup(
+    cfg: oc.DictConfig | None = None,
+    print_config: bool = False,
+    multiprocessing: str = "spawn",
+    seed: int | None = None,
+    post_mortem: bool = False,
+    determinism: bool = False,
+    resolvers: dict[str, "oc.Resolver"] | None = None,
+) -> SingletonResolver:
+    """Setup the training environment.
+
+    Args:
+        cfg (oc.OmegaConf | None): Full configuration
+        print_config (bool): Whether to print configuration to output. Defaults to ``False``
+        multiprocessing (str): Multiprocessing mode. Defaults to ``spawn``
+        seed (int | None): Random number generator seed to start off when set
+        post_mortem (bool): Whether to start pdb debugger when an uncaught exception encoutered. Defaults to ``False``
+        determinism (bool): Whether to enforce deterministric algorithms. Defaults to ``False```
+        resolvers (dict[str, oc.Resover] | None): Custom resolvers to register for configuration processing
+
+    Returns:
+        :class:`SingletonResolver` object that holds initialized data, trainer, model, etc.
+    """
+    if print_config:
+        log.info(config.to_str(cfg))
+
+    mp.set_start_method(multiprocessing, force=True)
+
+    if seed:
+        L.seed_everything(seed)
+
+    if post_mortem:
+        pdb.post_mortem()
+
+    if determinism:  # Enable deterministic algorithms globally
+        assert seed is not None, "seed has to be set for deterministic runs"
+        torch.use_deterministic_algorithms(True)
+        if torch.cuda.is_available():
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+            deterministic_env_vars = {
+                "CUBLAS_WORKSPACE_CONFIG": [":16:8", ":4096:8"],
+                "NCCL_ALGO": ["^NVLS"],
+                "NVTE_ALLOW_NONDETERMINISTIC_ALGO": ["0"],
+            }
+            for var, vals in deterministic_env_vars.items():
+                if (val := os.environ.get(var, vals[-1])) not in vals:
+                    raise ValueError(f"{var} has to be set to one of {vals} for deterministic runs, got: {val}")
+                os.environ[var] = val
+    for rn, fn in (resolvers or {}).items():
+        oc.OmegaConf.register_new_resolver(rn, fn, replace=True)
+
+    s = config.register_singleton_resolver()
+    return s
+
+
+def initialize(cfg: oc.DictConfig) -> SingletonResolver:
+    """Initialize data, model and trainer with supplied configurations.
+
+    Args:
+        cfg (oc.DictConfig): Configurations supplied by user through yaml file.
+
+    Returns:
+        :class:`SingletonResolver` object that holds initialized data, trainer, model, etc.
+    """
+    # 0. setup the training environment
+    s = setup(cfg, **(hydra.utils.instantiate(cfg["setup"]) if "setup" in cfg else {}))
+
+    # 1. instantiate `trainer`
+    s.trainer = hydra.utils.instantiate(cfg.trainer)
+
+    # 2. instantiate optional `data`
+    s.data = hydra.utils.instantiate(cfg.get("data"))
+
+    # 3. instantiate `model` after `trainer`
+    s.model = hydra.utils.instantiate(cfg.model)
+
+    # 4. obtain optional `ckpt_path`, `return_predictions` after `model`
+    s.ckpt_path = hydra.utils.call(cfg.get("ckpt_path"))
+    s.return_predictions = hydra.utils.call(cfg.get("return_predictions"))
+
+    # 5. save and upload the config
+    config.save(cfg, s.trainer)
+
+    # 6. run tuners
+    s.tuners = hydra.utils.instantiate(cfg.get("tuners"))
+
+    return s

fkat/data/__init__.py

@@ -0,0 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .data_module import DataModule, PersistStates, RestoreStates
+from .shm import ShmDataLoader
+from .sharded import ShardedDataLoader
+from .dict import DictDataLoader
+
+__all__ = [
+    "DataModule",
+    "ShmDataLoader",
+    "ShardedDataLoader",
+    "DictDataLoader",
+    "PersistStates",
+    "RestoreStates",
+]

fkat/data/data_module.py

@@ -0,0 +1,198 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import logging
+from functools import partial
+from typing import Any
+from collections.abc import Callable, Iterable
+
+import lightning as L
+from lightning.pytorch.profilers import Profiler
+from lightning.pytorch.core.hooks import CheckpointHooks
+from torch.utils.data import DataLoader
+from typing_extensions import Protocol, override, runtime_checkable
+
+from fkat.utils.rng import get_rng_states, set_rng_states
+from fkat.utils.profiler import profile_until_exit
+
+logger = logging.getLogger(__name__)
+
+
+def _call(dataloader: Iterable[Any] | None, attr: str, *args: Any, **kwargs: Any) -> None:
+    if dataloader is None:
+        return
+    for obj in (
+        dataloader,
+        getattr(dataloader, "dataset", None),
+        getattr(dataloader, "sampler", None),
+        getattr(dataloader, "batch_sampler", None),
+    ):
+        if not obj:
+            continue
+        if impl := getattr(obj, attr, None):
+            impl(*args, **kwargs)
+
+
+def worker_init_fn(
+    profiler: Profiler,
+    stage: str,
+    init_fn: Callable[[int], Any] | None,
+    worker_id: int,
+) -> None:
+    action = f"DataWorker[{stage}][{worker_id}]"
+    profile_until_exit(profiler, action=action, filename_suffix=f"_{stage}_{worker_id}")
+
+    if init_fn is not None:
+        init_fn(worker_id)
+
+    # TODO: Add Dataloader worker have consistent seed based on worker_id
+    # Reference: https://github.com/pytorch/pytorch/issues/5059#issuecomment-404232359
+
+
+def instrument(
+    cfg: dict[str, Any],
+    profiler: Profiler | None,
+    stage: str,
+) -> dict[str, Any]:
+    if not profiler:
+        return cfg
+    cfg["worker_init_fn"] = partial(worker_init_fn, profiler, stage, cfg.get("worker_init_fn"))
+    return cfg
+
+
+@runtime_checkable
+class PersistStates(Protocol):
+    def state_dict(self) -> dict[str, Any]: ...
+
+
+@runtime_checkable
+class RestoreStates(Protocol):
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None: ...
+
+
+class DataModule(L.LightningDataModule, CheckpointHooks):
+    """A :class:`LightningDataModule` that manages multiple :class:`DataLoader`\\s for different stages.
+
+    Args:
+        dataloaders (dict[str, dict[str, Any] | Callable[[], Iterable[Any]]]): Dataloaders for different stages.
+        profiler (Profiler | None): Profiler instance for worker initialization.
+    """
+
+    SUPPORTED_STAGES = ("train", "test", "val", "predict")
+
+    def __init__(
+        self,
+        dataloaders: dict[str, dict[str, Any] | Callable[[], Iterable[Any]]],
+        profiler: Profiler | None = None,
+    ) -> None:
+        super().__init__()
+        self.profiler = profiler
+        self.dataloader_factory: dict[str, Callable[[], Iterable[Any]]] = {}
+        for stage, cfg in dataloaders.items():
+            if stage not in DataModule.SUPPORTED_STAGES:
+                raise ValueError(f"Unsupported stage {stage}, use one of {DataModule.SUPPORTED_STAGES}")
+            dataloader_factory: Callable[[], Iterable[Any]]
+            if isinstance(cfg, dict):
+                cfg = instrument(cfg, profiler, stage)  # type: ignore[arg-type]
+                dataloader_factory = partial(DataLoader, **cfg)
+            else:
+                dataloader_factory = cfg
+            self.dataloader_factory[stage] = dataloader_factory
+        self.dataloaders: dict[str, Iterable[Any] | None] = {}
+
+    def _new_dataloader(self, stage: str) -> Iterable[Any] | None:
+        dataloader_factory = self.dataloader_factory.get(stage, lambda: None)
+        self.dataloaders[stage] = (dataloader := dataloader_factory())
+        return dataloader
+
+    @override
+    def prepare_data(self) -> None:
+        for dataloader in self.dataloaders.values():
+            _call(dataloader, "prepare_data")
+
+    def _dataloader(self, stage: str) -> Iterable[Any] | None:
+        stage = "train" if stage == "fit" else "val" if stage == "validation" else stage
+        return self.dataloaders.get(stage)
+
+    @override
+    def setup(self, stage: str) -> None:
+        device = self.trainer and self.trainer.strategy and self.trainer.strategy.root_device
+        _call(self._dataloader(stage), "set_device", device)
+        _call(self._dataloader(stage), "setup", stage)
+
+    # will be used once https://github.com/Lightning-AI/pytorch-lightning/pull/19601 is in effect
+    def on_exception(self, exception: BaseException) -> None:
+        for dataloader in self.dataloaders.values():
+            _call(dataloader, "on_exception", exception)
+
+    @override
+    def teardown(self, stage: str | None) -> None:
+        # this is it, terminating everything regardless of which stage we received this
+        for dataloader in self.dataloaders.values():
+            _call(dataloader, "teardown", stage)
+
+    @override
+    def train_dataloader(self) -> Iterable[Any] | None:
+        return self._new_dataloader("train")
+
+    @override
+    def val_dataloader(self) -> Iterable[Any] | None:
+        return self._new_dataloader("val")
+
+    @override
+    def predict_dataloader(self) -> Iterable[Any] | None:
+        return self._new_dataloader("predict")
+
+    @override
+    def test_dataloader(self) -> Iterable[Any] | None:
+        return self._new_dataloader("test")
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        """Called when saving a checkpoint, implement to generate and save datamodule state for ShardedDataLoader.
+
+        This method iterates over each stage's dataloader to retrieve its state using the `state_dict()` method,
+        and saves it along with the RNG states. If a dataloader does not implement the `PersistStates` protocol,
+        it sets its `state_dict` attribute to `vanilla_dataloader_state_dict` to allow saving its state.
+
+        Returns:
+            dict[str, Any]: A dictionary containing the dataloader states and RNG states.
+        """
+        dataloader_states = {}
+        for stage, dataloader in self.dataloaders.items():
+            dataloader_states[stage] = get_rng_states()
+            if isinstance(dataloader, PersistStates):
+                try:
+                    result_dict = dataloader.state_dict()
+                    dataloader_states[stage].update(result_dict)
+                except Exception as e:
+                    logger.warning(f"{dataloader} states can't be persisted yet: {e}.")
+        return dataloader_states
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Called when loading a checkpoint to reload the DataModule's state for a ShardedDataLoader.
+
+        This method iterates over each stage's dataloader, loads its state from the provided `state_dict`,
+        and sets the RNG states. If a dataloader does not implement the `RestoreStates` protocol,
+        it sets its `load_state_dict` attribute to `vanilla_dataloader_load_state_dict` to allow loading its state.
+
+        Args:
+            state_dict (Dict[str, Any]): A dictionary containing the dataloader states and RNG states.
+        """
+        for stage, dataloader in self.dataloaders.items():
+            set_rng_states(state_dict[stage])
+            if isinstance(dataloader, RestoreStates):
+                try:
+                    dataloader.load_state_dict(state_dict[stage])
+                except Exception as e:
+                    logger.warning(f"{dataloader} states can't be restored yet: {e}.")
+
+    @override
+    def on_save_checkpoint(self, checkpoint: dict[str, Any]) -> None:
+        for dataloader in self.dataloaders.values():
+            _call(dataloader, "on_save_checkpoint", checkpoint)
+
+    @override
+    def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
+        for dataloader in self.dataloaders.values():
+            _call(dataloader, "on_load_checkpoint", checkpoint)

fkat/data/datasets/__init__.py

@@ -0,0 +1,19 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .sized import SizedDataset  # isort: skip
+from .dict import DictDataset
+from .map import MapDataset, IterableMapDataset
+from .json import JsonDataset, IterableJsonDataset
+from .parquet import ParquetDataset, IterableParquetDataset
+
+
+__all__ = [
+    "SizedDataset",
+    "MapDataset",
+    "IterableMapDataset",
+    "DictDataset",
+    "JsonDataset",
+    "IterableJsonDataset",
+    "ParquetDataset",
+    "IterableParquetDataset",
+]

fkat/data/datasets/dict.py

@@ -0,0 +1,78 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+
+from typing_extensions import override
+
+from fkat.data.datasets import SizedDataset
+from fkat.utils.config import to_primitive_container
+
+
+class DictDataset(SizedDataset[tuple[str, Any], dict[str, Any]]):
+    """:class:`Dataset` that can get samples from one of the :class:`Dataset` using a mapping."""
+
+    def __init__(
+        self,
+        datasets: dict[str, SizedDataset[Any, dict[str, Any]]],
+        key: str = "dataset",
+    ) -> None:
+        """Create a :class:`Dataset` that can get samples from one of the :class:`Dataset` using a mapping.
+
+        Args:
+            datasets (Dict[str, SizedDataset[Any, Dict[str, Any]]]): A mapping from labels to :class:`Dataset`\\s.
+            key (str): The name of the field to reflect the :class:`Dataset` the samples were provided from.
+                Defaults to "dataset".
+
+        Returns:
+            None
+        """
+        self.datasets = to_primitive_container(datasets)
+        self.len = sum(len(dataset) for dataset in datasets.values())
+        self.key = key
+
+    @override
+    def __len__(self) -> int:
+        """Get :class:`Dataset` size.
+
+        Returns:
+            int: :class:`Dataset` size.
+        """
+        return self.len
+
+    def _wrap(self, name: str, item: dict[str, Any]) -> dict[str, Any]:
+        if not isinstance(item, dict) or self.key in item:
+            raise RuntimeError(f"Datasets must return a dict without {self.key} key")
+        item[self.key] = name
+        return item
+
+    def __getitems__(self, name_and_idxs: tuple[str, list[Any]]) -> list[dict[str, Any]]:
+        """Get a batch of samples from the target :class:`Dataset` at the specified indices.
+
+        Args:
+            name_and_idxs (Tuple[str, List[Any]]): Samples' :class:`Dataset` and indices.
+
+        Returns:
+            List[Dict[str, Any]]: A batch of samples.
+        """
+        name, idxs = name_and_idxs
+        if getitems := getattr(self.datasets[name], "__getitems__", None):
+            batch = getitems(idxs)
+        else:
+            batch = [self.datasets[name][idx] for idx in idxs]
+        for b in batch:
+            self._wrap(name, b)
+        return batch
+
+    @override
+    def __getitem__(self, idx: tuple[str, Any]) -> dict[str, Any]:
+        """Get a sample from the target :class:`Dataset` at the specified index.
+
+        Args:
+            idx (Tuple[str, Any]): Sample :class:`Dataset` and index.
+
+        Returns:
+            Dict[str, Any]: A sample.
+        """
+        name, idx_ = idx
+        sample = self.datasets[name][idx_]
+        return self._wrap(name, sample)

fkat/data/datasets/json.py

@@ -0,0 +1,176 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import awswrangler as s3wr
+from typing import Any
+from collections.abc import Iterator
+
+import pyarrow as pa
+import pyarrow.json as pj
+from pyarrow.fs import FileSystem, S3FileSystem  # type: ignore[possibly-unbound-import]
+import pandas as pd
+import numpy as np
+from fkat.data.datasets import SizedDataset
+from fkat.utils.pyarrow import iter_rows as pa_iter_rows
+from fkat.utils.pandas import iter_rows as pd_iter_rows
+from fkat.utils.boto3 import session
+from torch.utils.data import IterableDataset
+
+
+class IterableJsonDataset(IterableDataset[dict[str, Any]]):
+    """
+    An :class:`IterableDataset` backed by Json data.
+
+    Args:
+        uri (str | list[str]): URI of Parquet data.
+        read_options: pyarrow.json.ReadOptions, optional
+            Options for the JSON reader (see ReadOptions constructor for defaults).
+        parse_options: pyarrow.json.ParseOptions, optional
+            Options for the JSON parser
+            (see ParseOptions constructor for defaults).
+        memory_pool: MemoryPool, optional
+            Pool to allocate Table memory from.
+        chunk_size (int): An iterable of DataFrames is returned with maximum rows equal to the received INTEGER.
+        replace_nan (bool): Whether to replace np.nan as None.
+            Default to ``True``
+        s3wr_args (dict): config for s3wr.s3.read_json,
+            refer to https://aws-sdk-pandas.readthedocs.io/en/3.5.1/stubs/awswrangler.s3.read_parquet.html
+
+    """
+
+    def __init__(
+        self,
+        uri: str | list[str],
+        read_options: pa.json.ReadOptions | None = None,
+        parse_options: pa.json.ParseOptions | None = None,
+        memory_pool: pa.MemoryPool | None = None,
+        chunk_size: int = 10000,
+        replace_nan: bool = True,
+        **s3wr_args: Any,
+    ) -> None:
+        fs: FileSystem
+        path: str
+        if isinstance(uri, str):
+            fs, path = FileSystem.from_uri(uri)
+        else:
+            fs, path = FileSystem.from_uri(uri[0])
+        self.s3_file = isinstance(fs, S3FileSystem)
+        if self.s3_file:
+            self.uri = uri
+            self.chunk_size = chunk_size
+            self.replace_nan = replace_nan
+            self.s3wr_args = s3wr_args
+        else:
+            assert isinstance(uri, str), "IterableJsonDataset can only accept uri as str"
+            with fs.open_input_file(path) as f:
+                self.tbl = pj.read_json(
+                    f, read_options=read_options, parse_options=parse_options, memory_pool=memory_pool
+                )
+            self.chunk_size = chunk_size
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Creates dataset iterator.
+        Returns:
+            Iterator[dict[str, Any]]: dataset iterator
+        """
+        if self.s3_file:
+            return pd_iter_rows(
+                s3wr.s3.read_json(
+                    self.uri,
+                    lines=True,
+                    chunksize=self.chunk_size,
+                    boto3_session=session(clients=["s3"]),
+                    path_suffix="json",
+                    **self.s3wr_args,
+                ),
+                self.replace_nan,
+            )
+        else:
+            return pa_iter_rows(self.tbl, chunk_size=self.chunk_size)
+
+
+class JsonDataset(SizedDataset[int, dict[str, Any]]):
+    """
+    Create a :class:`Dataset` from JSON data at the specified URI.
+
+    Args:
+        uri (str | list[str]): URI of JSON data.
+        read_options (pa.json.ReadOptions | None): JSON read options.
+        parse_options (pa.json.ParseOptions | None): JSON parse options.
+        memory_pool (pa.MemoryPool | None): JSON processing memory pool configuration.
+        replace_nan (bool): Whether to replace np.nan as None.
+            Default to ``True``
+        s3wr_args (Any): config for s3wr.s3.read_json,
+            refer to https://aws-sdk-pandas.readthedocs.io/en/3.5.1/stubs/awswrangler.s3.read_json.html
+    """
+
+    def __init__(
+        self,
+        uri: str | list[str],
+        read_options: pa.json.ReadOptions | None = None,
+        parse_options: pa.json.ParseOptions | None = None,
+        memory_pool: pa.MemoryPool | None = None,
+        replace_nan: bool = True,
+        **s3wr_args: Any,
+    ) -> None:
+        fs: FileSystem
+        path: str
+        if isinstance(uri, str):
+            fs, path = FileSystem.from_uri(uri)
+        else:
+            fs, path = FileSystem.from_uri(uri[0])
+        if isinstance(fs, S3FileSystem):
+            self.df = s3wr.s3.read_json(
+                uri, lines=True, boto3_session=session(clients=["s3"]), path_suffix="json", **s3wr_args
+            )
+            if replace_nan:
+                self.df = self.df.replace({np.nan: None})
+        else:
+            path_list: list[str] = []
+            if isinstance(uri, str):
+                path_list.append(path)
+            else:
+                for each in uri:
+                    _, path = FileSystem.from_uri(each)
+                    path_list.append(path)
+            df = []
+            for each in path_list:
+                with fs.open_input_file(each) as f:
+                    tbl = pj.read_json(
+                        f, read_options=read_options, parse_options=parse_options, memory_pool=memory_pool
+                    )
+                    df.append(tbl.to_pandas())
+            self.df = pd.concat(df)
+
+    def __len__(self) -> int:
+        """Get :class:`Dataset` size.
+
+        Returns:
+            int: :class:`Dataset` size.
+        """
+        return len(self.df)
+
+    def __getitems__(self, idxs: list[int]) -> list[dict[str, Any]]:
+        """Get a batch of samples at the specified indices.
+
+        Args:
+            idxs (list[int]): Samples' indices.
+
+        Returns:
+            list[dict[str, Any]]: A batch of samples.
+        """
+        series = self.df.iloc[idxs]
+        samples = [series.iloc[i].to_dict() for i in range(len(idxs))]
+        return samples
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Get a sample at the specified index.
+
+        Args:
+            idx (int): Sample index.
+
+        Returns:
+            dict[str, Any]: A sample.
+        """
+        series = self.df.iloc[idx]
+        sample = series.to_dict()
+        return sample