chugr 0.2.0.dev0__py3-none-any.whl

chugr/wds/loader.py

@@ -0,0 +1,187 @@
+import dataclasses
+import math
+import warnings
+from typing import Callable, List, Optional, Union
+
+import torch.utils.data
+import webdataset as wds
+
+from chugr.common import DistributedCfg, LoaderBundle, SharedCount, ShardSpec, rank_seed, seed_worker
+from .helpers import expand_urls
+from .pipeline import build_data_pipeline
+
+
+_SAMPLE_SHUFFLE_SIZE = 2000
+_SAMPLE_SHUFFLE_INITIAL = 500
+
+
+def create_loader_wds(
+        shards: Union[str, List[str], ShardSpec],
+        task_pipeline: Optional[List[Callable]],
+        num_samples: Optional[int] = None,
+        is_training: bool = False,
+        batch_size: Optional[int] = 1,
+        resampled: bool = False,
+        multi_interval: bool = True,
+        num_batches_round: str = 'ceil',
+        num_workers: int = 4,
+        persistent_workers: bool = True,
+        start_interval: int = 0,
+        seed: int = 0,
+        handler: Callable = wds.reraise_exception,
+        collate_fn: Optional[Callable] = None,
+        sample_shuffle_size: int = _SAMPLE_SHUFFLE_SIZE,
+        sample_shuffle_initial: int = _SAMPLE_SHUFFLE_INITIAL,
+        distributed: DistributedCfg = DistributedCfg(),
+        naflex_scheduler=None,
+        naflex_stage_factory: Optional[Callable] = None,
+):
+    """ Create a webdataset loader
+
+    Args:
+        shards:
+        task_pipeline:
+        is_training:
+        resampled: If True, shards are resampled with replacement.
+        multi_interval: If True, run loader in multi-interval mode (multi-epoch), num_samples is interpreted
+            as num_samples per interval (epoch if # samples == dataset length). Dataset length is set to a
+            fixed value to approximate at-least once sample visiting per interval.
+            If False, num_samples is treated as total samples to visit over training without paying attention
+            to # samples in underlying dataset. Dataset length is not accessible.
+        start_interval:
+        seed:
+        num_workers:
+        persistent_workers:
+        batch_size:
+        num_batches_round:
+        collate_fn:
+        sample_shuffle_size:
+        sample_shuffle_initial:
+        distributed:
+        naflex_scheduler: NaFlexBatchScheduler for schedule driven (variable seq_len & batch size) batching
+        naflex_stage_factory: interval counter -> pipeline stages (bucketing, NaFlexBatcher) replacing the
+            task pipeline and fixed size batching
+
+    Returns:
+
+    """
+    resampled = resampled and is_training
+
+    if not isinstance(shards, ShardSpec):
+        if isinstance(shards, str):
+            shards = expand_urls(shards)
+        shards = ShardSpec(
+            urls=shards,
+        )
+
+    num_shards = len(shards.urls)
+    # eval, an explicit num_samples below the (known) dataset size caps the eval set, otherwise all samples are used
+    max_samples = None
+    if not is_training and num_samples and not (shards.sizes and num_samples >= sum(shards.sizes)):
+        max_samples = num_samples
+    if num_samples is None:
+        if shards.sizes:
+            num_samples = sum(shards.sizes)
+        if is_training and not num_samples:
+            raise RuntimeError(
+                'The number of dataset samples must be specified for the training dataset '
+                'if no dataset length info is present.')
+
+    num_batches_per_worker = 0
+    num_workers_nonzero = max(num_workers, 1)
+    if is_training and naflex_scheduler is not None:
+        # schedule driven batching, each worker yields its share of the rank identical schedule per interval,
+        # the sample stream must not run dry -> resampled shards, or (w/o replacement) repeating shard lists
+        num_batches = naflex_scheduler.num_batches_for_workers(num_workers_nonzero)
+        num_samples = naflex_scheduler.num_samples_for_workers(num_workers_nonzero)
+        batch_size = 0
+    elif is_training:
+        assert batch_size >= 1, 'batching must be enabled for train, set batch_size>=1'
+        # We want to see the same # of batches on each member of the distributed group (GPU),
+        # this is enforced by making each worker produce the same # of batches regardless of the
+        # underlying iterator, so we estimate and make the iterator wrap around if end is hit.
+        # This will repeat some samples and may miss some sample per interval as shards may be
+        # uneven or allocated unevenly across all workers. There are ways improve on this naive
+        # approach, to get closer to the ideal of each sample in an interval (epoch) seen once,
+        # but difficult to achieve perfectly, and most improvements require full co-ordination across
+        # all workers via out-of-band IPC/RPC.
+
+        # roll over and repeat a few samples to get same number of full batches on each node
+        round_fn = math.floor if num_batches_round == 'floor' else math.ceil
+        global_batch_size = batch_size * distributed.world_size
+        num_batches = round_fn(num_samples / global_batch_size)
+        num_batches_per_worker = round_fn(num_batches / num_workers_nonzero)  # per dataloader worker
+        num_batches = num_batches_per_worker * num_workers_nonzero
+        num_samples = num_batches * global_batch_size
+    else:
+        # Eval / inference will exhaust the iterator if the size is not specified.
+        # Samples (or shards, w/ enough of them and no sample cap) are split across ranks, the counts per rank
+        # differ, and ranks may see different numbers of (partial) batches, num_batches is an estimate. Eval
+        # loops must not run collectives per batch.
+        num_samples = num_samples or 0
+        if batch_size:
+            num_batches = math.ceil(math.ceil(num_samples / distributed.world_size) / batch_size)
+        else:
+            num_batches = 0
+
+    if is_training:
+        # create a shared epoch store to sync epoch to dataloader worker proc
+        shared_interval_count = SharedCount(count=start_interval)
+        if not resampled:
+            assert num_shards >= num_workers_nonzero * distributed.world_size, 'number of shards must be >= total workers'
+    else:
+        shared_interval_count = None
+
+    if is_training and naflex_scheduler is not None:
+        task_pipeline = naflex_stage_factory(shared_interval_count)
+
+    # DataLoader worker seeds (resampled shards, sample shuffle, worker RNGs) must differ across ranks, the shard
+    # list shuffle (w/o resampling) uses the shared `seed` so the node / worker split stays consistent
+    worker_seed = rank_seed(seed, distributed.global_rank)
+    datapipe = build_data_pipeline(
+        shards=shards,
+        task_pipeline=task_pipeline,
+        is_training=is_training,
+        batch_size=batch_size,
+        resampled=resampled,
+        multi_interval=multi_interval,
+        seed=seed,
+        shared_interval_count=shared_interval_count,
+        num_batches_per_worker=num_batches_per_worker,
+        sample_shuffle_initial=sample_shuffle_initial,
+        sample_shuffle_size=sample_shuffle_size,
+        handler=handler,
+        collate_fn=collate_fn,
+        max_samples=max_samples,
+        repeat_shards=is_training and naflex_scheduler is not None and not resampled,
+        worker_seed=worker_seed,
+        rank=distributed.global_rank,
+        world_size=distributed.world_size,
+        num_workers=num_workers,
+    )
+
+    dl_generator = torch.Generator()
+    dl_generator.manual_seed(worker_seed)
+
+    dataloader = wds.WebLoader(
+        datapipe,
+        batch_size=None,  # batching done in data-pipeline
+        shuffle=False,  # shuffling done in data-pipeline
+        num_workers=num_workers,
+        persistent_workers=persistent_workers and num_workers > 0,
+        generator=dl_generator,
+        worker_init_fn=seed_worker,
+    )
+
+    return LoaderBundle(
+        loader=dataloader,
+        num_batches=num_batches,
+        num_samples=num_samples,
+        shared_interval=shared_interval_count,
+        mean_batch_size=(
+            naflex_scheduler.mean_batch_size_for_workers(num_workers_nonzero)
+            if is_training and naflex_scheduler is not None else None
+        ),
+    )
+
+

chugr/wds/pipeline.py

@@ -0,0 +1,182 @@
+from functools import partial
+from itertools import islice
+from typing import Callable, List, Optional, Sequence, Union
+
+import torch.utils.data
+import webdataset
+import webdataset as wds
+
+from chugr.common import ShardSpec, collate
+from .filters import detshuffle_v2
+from .shardlists import ResampledShardsV2, ShuffledShardList
+from .tariterators import tarfile_to_samples_nothrow
+
+
+def split_samples_by_node_and_worker(src, rank: int = 0, world_size: int = 1):
+    """ Split a sample stream that is identical on all ranks / workers, rank r keeps every world_size-th sample
+    (independent of the worker count), its workers split those. """
+    info = torch.utils.data.get_worker_info()
+    num_workers = info.num_workers if info is not None else 1
+    worker_id = info.id if info is not None else 0
+    yield from islice(islice(src, rank, None, world_size), worker_id, None, num_workers)
+
+
+def build_data_pipeline(
+        shards: ShardSpec,
+        task_pipeline: List[Callable],
+        is_training: bool = False,
+        batch_size: int = 0,
+        resampled: bool = False,
+        multi_interval: bool = False,
+        seed: int = 0,
+        shared_interval_count: int = -1,
+        num_batches_per_worker: int = 0,
+        sample_shuffle_initial: int = 1,
+        sample_shuffle_size: int = 1,
+        collate_fn: Optional[Callable] = None,
+        batched_task_pipeline: Optional[List[Callable]] = None,
+        handler=wds.reraise_exception,
+        max_samples: Optional[int] = None,
+        repeat_shards: bool = False,
+        worker_seed: Optional[int] = None,
+        rank: int = 0,
+        world_size: int = 1,
+        num_workers: int = 0,
+):
+    """
+
+    Args:
+        task_pipeline: task-specific pipeline for processing samples
+        shards: ShardSpec w/ shard url list and optional sizes and weights
+        is_training: Train mode. Enables shuffling of samples and forces consistent batch #s across workers.
+        batch_size: Batch for sample collation. Collation / batching disabled if batch_size == 0.
+        resampled: Enable resampling with replacement of shards.
+        multi_interval:
+        seed:
+        shared_interval_count:
+        num_batches_per_worker:
+        sample_shuffle_initial:
+        sample_shuffle_size:
+        collate_fn:
+        batched_task_pipeline: An optional task-specific pipeline for processing batched samples
+        handler: Exception handler
+        max_samples: Eval, limit the eval set to the first `max_samples` samples of the shard list. Exact for any
+            number of ranks / workers, every worker reads those samples and keeps its share, so leave it unset
+            to evaluate all samples (split by shard where possible).
+        repeat_shards: Train w/o resampling, cycle through the shard list (reshuffled each pass)
+        worker_seed: Seed for per worker randomness (shard resampling, sample shuffle) when the seed can't be
+            recovered from DataLoader workers (num_workers=0), must differ across ranks. Defaults to `seed`.
+        rank: Rank of this process, shards (train, eval w/ enough shards) or samples (eval) are split across
+            ranks. Passed explicitly, not derived from RANK / WORLD_SIZE env vars or torch.distributed state
+            (neither is reliably visible in dataloader workers for all launchers / start methods).
+        world_size: Number of ranks
+        num_workers: Number of dataloader workers (eval split)
+
+    Returns:
+
+    """
+    worker_seed = seed if worker_seed is None else worker_seed
+    if not isinstance(task_pipeline, (list, tuple)):
+        task_pipeline = [task_pipeline]
+    assert len(task_pipeline)
+
+    if resampled:
+        datapipe = [ResampledShardsV2(
+            shards.urls,
+            weights=shards.weights,
+            deterministic=True,
+            interval=shared_interval_count,
+            seed=worker_seed,
+        )]
+    else:
+        assert shards.weights is None, \
+            "upsampling_factors is only supported when sampling with replacement (resampled=False)."
+        if is_training:
+            datapipe = [ShuffledShardList(
+                shards.urls,
+                seed=seed,
+                interval=shared_interval_count,
+                repeat=repeat_shards,
+            )]
+        else:
+            datapipe = [wds.SimpleShardList(
+                shards.urls,
+            )]
+
+    # at this point we have an iterator over all the shards
+    if is_training:
+        if not resampled:
+            if world_size > 1:
+                # as wds.split_by_node, w/ the explicit rank
+                datapipe.append(wds.slice(rank, None, world_size))
+            datapipe.append(wds.split_by_worker)
+        # at this point, we have an iterator over the shards assigned to each worker at each node
+        datapipe.extend([
+            tarfile_to_samples_nothrow(handler=handler),
+            detshuffle_v2(
+                bufsize=sample_shuffle_size,
+                initial=sample_shuffle_initial,
+                seed=worker_seed,
+                interval=shared_interval_count,
+                unique_worker=True,
+            )
+            # wds.shuffle(
+            #     bufsize=sample_shuffle_size,
+            #     initial=sample_shuffle_initial,
+            # ),
+        ])
+    else:
+        if max_samples:
+            # capped eval set, the first max_samples samples of the shard list for any world size / worker count.
+            # Every worker reads (but doesn't decode / process) that head of the stream and keeps its share.
+            datapipe.extend([
+                wds.tarfile_to_samples(handler=handler),
+                wds.slice(max_samples),
+                partial(split_samples_by_node_and_worker, rank=rank, world_size=world_size),
+            ])
+        elif len(shards.urls) >= world_size * max(num_workers, 1):
+            # enough shards for every worker of every rank, split shards so each is read once
+            if world_size > 1:
+                datapipe.append(wds.slice(rank, None, world_size))
+            datapipe.extend([
+                wds.split_by_worker,
+                wds.tarfile_to_samples(handler=handler),
+            ])
+        else:
+            # few (or one) eval shards, split samples across ranks, every rank reads the samples of its workers'
+            # shards but only decodes / processes its share
+            datapipe.extend([
+                wds.split_by_worker,
+                wds.tarfile_to_samples(handler=handler),
+            ])
+            if world_size > 1:
+                datapipe.append(wds.slice(rank, None, world_size))
+
+    # task specific decode and map pipline (per-sample)
+    datapipe.extend(task_pipeline)
+
+    # collation, tensor output disabled with batch_size == 0 or None
+    if batch_size:
+        # NOTE torch default_collate handles dicts, wds default_collate does not
+        collate_fn = collate_fn or collate
+        datapipe.extend([
+            wds.batched(
+                batch_size,
+                partial=not is_training,
+                collation_fn=collate_fn,
+            )
+        ])
+
+        # task-specific batched pipeline (per-batch)
+        if batched_task_pipeline:
+            datapipe.extend(batched_task_pipeline)
+
+    datapipe = wds.DataPipeline(*datapipe)
+
+    if is_training and num_batches_per_worker > 0:
+        if multi_interval:
+            datapipe = datapipe.with_epoch(num_batches_per_worker)  # each worker is iterating over this
+        else:
+            datapipe = datapipe.repeat(nbatches=num_batches_per_worker)
+
+    return datapipe

chugr/wds/shardlists.py

@@ -0,0 +1,134 @@
+import random
+import sys
+
+from torch.utils.data import IterableDataset
+
+from chugr.common import SharedCount, get_pytorch_worker_seed
+from .helpers import expand_urls
+
+
+class ShuffledShardList(IterableDataset):
+    """An iterable dataset yielding a list of urls that is deterministically shuffled based on epoch."""
+
+    def __init__(
+            self,
+            urls,
+            seed=0,
+            interval=-1,
+            num_sub_intervals=None,
+            repeat: bool = False,
+    ):
+        """Iterate through the list of shards.
+
+        Args:
+            repeat: cycle through the shards indefinitely (each pass reshuffled w/ a seed shared by all
+                ranks / workers so node / worker splitting stays consistent), for schedule driven loaders that
+                need an infinite sample stream w/o sampling shards with replacement.
+        """
+        super().__init__()
+        self.repeat = repeat
+        self.urls = expand_urls(urls)
+        assert len(self.urls) and isinstance(self.urls[0], str)
+        self.seed = seed
+        self.interval = interval
+        self.num_sub_intervals = num_sub_intervals  # FIXME experimental feature
+
+    def __len__(self):
+        return len(self.urls)
+
+    def __iter__(self):
+        """Return an iterator over the shards."""
+        urls = self.urls.copy()
+
+        # Set epoch
+        if isinstance(self.interval, SharedCount):
+            interval = self.interval.get_value()
+        else:
+            # NOTE: this is interval tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.interval += 1
+            interval = self.interval
+
+        pass_idx = 0
+        while True:
+            pass_urls = list(urls)
+            if self.seed is not None:
+                # Shuffle with the same seed across all nodes/workers in each interval or super interval
+                if self.num_sub_intervals is None:
+                    seed = self.seed + interval
+                else:
+                    # Keep shuffling consistent across the super epochs
+                    seed = self.seed + (interval // self.num_sub_intervals)
+                random.Random(seed * 1000003 + pass_idx if pass_idx else seed).shuffle(pass_urls)
+
+            # Restrict to shards in the sub epoch if needed
+            if self.num_sub_intervals is not None:
+                pass_urls = pass_urls[interval % self.num_sub_intervals::self.num_sub_intervals]
+
+            # Yield shards
+            for url in pass_urls:
+                yield dict(url=url)
+
+            if not self.repeat:
+                break
+            pass_idx += 1
+
+
+class ResampledShardsV2(IterableDataset):
+    """An iterable dataset yielding a list of urls."""
+
+    def __init__(
+            self,
+            urls,
+            weights=None,
+            nshards=sys.maxsize,
+            worker_seed_fn=None,
+            deterministic=False,
+            interval=-1,
+            seed=None,
+    ):
+        """Sample shards from the shard list with replacement.
+
+        :param urls: a list of URLs as a Python list or brace notation string
+        """
+        super().__init__()
+        if weights is not None:
+            self.urls, self.weights = expand_urls(urls, weights)
+            assert len(self.urls) == len(self.weights), \
+                f"Number of urls {len(self.urls)} and weights {len(self.weights)} should match."
+        else:
+            self.urls = expand_urls(urls)
+            self.weights = None
+        assert isinstance(self.urls[0], str)
+        self.nshards = nshards
+        self.rng = random.Random()
+        self.worker_seed_fn = worker_seed_fn
+        self.deterministic = deterministic
+        self.interval = interval
+        self.seed = seed  # only used when seed cannot be recovered from DL workers
+
+    def __iter__(self):
+        """Return an iterator over the shards."""
+        if isinstance(self.interval, SharedCount):
+            interval = self.interval.get_value()
+        else:
+            # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.interval += 1
+            interval = self.interval
+
+        if self.deterministic:
+            # reset seed w/ interval if deterministic
+            if self.worker_seed_fn is None:
+                # pytorch worker seed should be deterministic (per-worker)
+                # It is init by process seed, rank, & worker id
+                seed = get_pytorch_worker_seed(interval, initial_seed=self.seed)
+            else:
+                seed = self.worker_seed_fn() + interval
+            self.rng.seed(seed)
+
+        for _ in range(self.nshards):
+            if self.weights is None:
+                yield dict(url=self.rng.choice(self.urls))
+            else:
+                yield dict(url=self.rng.choices(self.urls, weights=self.weights, k=1)[0])

chugr/wds/tariterators.py

@@ -0,0 +1,70 @@
+import re
+
+from webdataset.filters import pipelinefilter
+from webdataset.tariterators import base_plus_ext, valid_sample, url_opener, tar_file_expander
+
+from .helpers import log_and_continue
+
+BASE_RE = re.compile(r"^((?:.*/|)[^.]+)[.]([^/]*)$")
+
+
+def base_plus_ext(path):
+    """Split off all file extensions.
+
+    Returns base, allext.
+
+    Args:
+        path: path with extensions
+
+    Returns:
+        path with all extensions removed
+    """
+    match = re.match(BASE_RE, path)
+    if not match:
+        return None, None
+    return match.group(1), match.group(2)
+
+
+def group_by_keys_nothrow(
+        data,
+        keys=base_plus_ext,
+        lcase=True,
+        suffixes=None,
+        handler=None,
+):
+    """Return function over iterator that groups key, value pairs into samples.
+
+    :param keys: function that splits the key into key and extension (base_plus_ext)
+    :param lcase: convert suffixes to lower case (Default value = True)
+    """
+    current_sample = None
+    for filesample in data:
+        assert isinstance(filesample, dict)
+        fname, value = filesample["fname"], filesample["data"]
+        prefix, suffix = keys(fname)
+        if prefix is None:
+            continue
+        if lcase:
+            suffix = suffix.lower()
+        # FIXME wds version throws if suffix in current_sample, but we have a potential for
+        #  this happening in the current LAION400m dataset if a tar ends with same prefix as the next
+        #  begins, rare, but can happen since prefix aren't unique across tar files in that dataset
+        if current_sample is None or prefix != current_sample["__key__"] or suffix in current_sample:
+            if valid_sample(current_sample):
+                yield current_sample
+            current_sample = dict(__key__=prefix, __url__=filesample["__url__"])
+        if suffixes is None or suffix in suffixes:
+            current_sample[suffix] = value
+    if valid_sample(current_sample):
+        yield current_sample
+
+
+def tarfile_samples_nothrow(src, handler=log_and_continue):
+    # NOTE this is a re-impl of the wds impl with group_by_keys that doesn't throw
+    streams = url_opener(src, handler=handler)
+    files = tar_file_expander(streams, handler=handler)
+    samples = group_by_keys_nothrow(files, handler=handler)
+    return samples
+
+
+tarfile_to_samples_nothrow = pipelinefilter(tarfile_samples_nothrow)