merantix-momentum / squirrel-core Goto Github PK

Hey,

Thank you for working on this library!
I think it has a huge potential, especially for dataset creators to provide their dataset in an optimized deep-learning format that is well suited for distribution.
The performance of the MessagepackSerializer is amazing and being able to distribute subsets of the dataset (shards) is something I never wanted but really want to utilize in the future!

I have played around with some "MessagepackSerializer" configurations and according to some internal benchmarks, it would be helpful to allow the user to configure the compression algorithm.

Currently, the compression mode is "locked" to gzip.
I assume the main reason is due to the wide usage of gzip and to keep the code 'simple' as it makes it easy in the deserializer to know that the gzip compression was used:

Here I would like to note that given the extension, fsspec (default) could also infer the compression by inspecting the filename suffix. But I can see how this might cause problems if somebody would like to switch out fsspec with something else (although I would have no idea with what and why :D )

Other spots within the codebase that are coupled to this compression assumption are the methods from the SquirrelStore:

Or to show the significant parts:

• get:
  

  squirrel-core/squirrel/store/squirrel_store.py

  Lines 40 to 41 in 0ff6be7

  	fp = f"{self.url}/{key}.gz"
  	yield from self.serializer.deserialize_shard_from_file(fp, fs=self.fs, **kwargs)

• set:
  

  squirrel-core/squirrel/store/squirrel_store.py

  Lines 59 to 60 in 0ff6be7

  	fp = f"{self.url}/{key}.gz"
  	self.serializer.serialize_shard_to_file(value, fp, fs=self.fs, **kwargs)

• keys:
  

  squirrel-core/squirrel/store/squirrel_store.py

  Lines 66 to 67 in 0ff6be7

  	if k.endswith(".gz"):
  	yield k.rsplit(".gz", 1)[0]

In my internal benchmarks, I was able to greatly speed up the data loading by simply using no compression at all (None).
Although I am fully aware that the correct compression mode heavily depends on the specific hardware/use case.
But even in a network limited domain, I can see reasons to then prefer xz instead due to its better compression ratio and relatively similar decompression speed to gzip.

IMHO, I think it should be ok to not store any suffix at all for the squirrel store.
If I/a user looks inside of the squirrel store URL I think it is not mandatory to show what compression algorithm was used.
The user could/should use the designated driver/metadata that comes bundled with the dataset and let the driver handle the correct decompression.

If you don't agree I still think the gz extension doesn't have to be 'hardcoded' into these functions.
This is actually something that confused me when I was looking at the internals of the code base.
So instead, we could use something like:

comp = kwargs.get("compression", "gzip")
comp_to_ext_dict[comp] # just to show the concept

With these modifications, it should be possible to utilize different compression modes and make them easily configurable.
I would be very happy to create a PR and contribute to this project!

I'm using squirrel main, directly from source. I'm trying to seed the shuffle buffer of a Composable using an instantiated random.Random instance and expect the same output over several repeated instances. However this is not correct due to a bug in the squirrel.iterstream.iterators.py::get_random_range functionality. There an instantiated rng instance will be ignored if the seed is None.

Below I attach a test file which exposed the expected behavior. As of now test_1 and test_4 fail, while test_2 and test_3 are passing.

from squirrel.iterstream import IterableSource
import random

def get_itersource():
    data = list(range(100))
    source = IterableSource(data)
    return source


def test_1():

    iter_0 = get_itersource().shuffle(size=40, rng=random.Random(42)).collect()
    iter_1 = get_itersource().shuffle(size=40, rng=random.Random(42)).collect()

    assert iter_0 == iter_1, "test_1 failed"

    print("test_1 successful")

def test_2():

    iter_0 = get_itersource().shuffle(size=40, seed=42).collect()
    iter_1 = get_itersource().shuffle(size=40, seed=42).collect()

    assert iter_0 == iter_1, "test_2 failed"
    print("test_2 successful")


def test_3():

    iter_0 = get_itersource().shuffle(size=40, rng=random.Random(21), seed=42).collect()
    iter_1 = get_itersource().shuffle(size=40, rng=random.Random(21), seed=42).collect()

    assert iter_0 == iter_1, "test_3 failed"
    print("test_3 successful")


def test_4():

    iter_0 = get_itersource().shuffle(size=40, rng=random.Random(21)).collect()
    iter_1 = get_itersource().shuffle(size=40, seed=21).collect()

    assert iter_0 == iter_1, "test_4 failed"
    print("test_4 successful")

Using squirrel main directly from github.

When using shuffle inside a nested function/method, the forwarding of the shuffle defaults into the function does not work. According to the doc-strings size (int, optional): Buffer size for shuffling. Defaults to 1000. Skip the shuffle step if `size < 2`. so I would assume that feeding in None would create a default or similar.

Test file below.

import pytest
from squirrel.iterstream import IterableSource

def get_itersource():
    data = list(range(100))
    source = IterableSource(data)
    return source

@pytest.mark.parametrize('size', [None, 1, 200])
def test_1(size):

    iter = get_itersource().shuffle(size=size).collect()
    assert sorted(iter) == list(range(100))

The use case: we have a store with samples, and 1..n other stores that each contain only features. These stores must have the same keys and same number of samples per shard.

IterableZipSource makes it possible to zip items from several iterables and use that as a source. For instance:

it1 = MessagepackDriver(url1).get_iter()
it2 = MessagepackDriver(url2).get_iter()

it3 = IterableZipSource(iterables=[it1, it2]).collect()

When reporting a bug, please be sure to include the following:

• A descriptive title
• An isolated way to reproduce the behavior (example: GitHub repository with code isolated to the issue that anyone can clone to observe the problem)

Add the following test to test/test_iterstream/test_stream.py

@pytest.mark.parametrize("probs", [[0.4, 0.6], None])
def test_iterablesamplersource_is_picklable(probs: t.Optional[t.List[float]]) -> None:
    """Smoke test IterableSamplerSource"""
    res_1 = IterableSource([0, 1, 2, 3])
    res_2 = IterableSource([4, 5, 6])
    src = IterableSamplerSource([res_1, res_2], probs=probs)
    import pickle

    pkl = pickle.dumps(src)
    src_2 = pickle.loads(pkl)

    assert src.collect() == src_2.collect()

Run: poetry run pytest -s test/test_iterstream/test_stream.py::test_iterablesamplersource_is_picklable

• What version of squirrel-datasets and squirrel you're using, and the platform(s) you're running it on
  commit-sha: 86f61c0

• What packages or other dependencies you're using
  no additional packages

• The behavior you expect to see and the actual behavior

Expected: Test passes
Actual: Fails with exception: TypeError: cannot pickle 'generator' object

Hey, I've stumbled across a potentially easy-to-misunderstand part of the MapDriver.get_iter documentation:

The documentation of max_workers states that by default None will be used and also mentions that this will cause async_map to be called but I missed these parts of the documentation and was surprised to see that so many threads were allocated.

I am/was not too familiar with the ThreadPoolExecutor interface and find it somewhat surprising that None equals numer_of_processors x 5 according to the ThreadPoolExecutor definition.
Maybe it would be helpful to explicitly state that by default ThreadPoolExecutor will be used with so many threads?
The documentation string reads a bit unintuitive as the starts out that max_worker defines how many items are fetched simultaneously with max_worker and then continues to state that otherwise map is used.
From that perspective, max_workers=None doesn't sound like it should be using any threads at all.
Without knowing the default values of ThreadPoolExecutor I would make it more explicit that to disable threading one has to set max_workers=0/1 and that by default many threads are used.

I am happy to add a PR with my suggested doc-string update if you agree! :)

When reporting a bug, please be sure to include the following:

• A descriptive title
• An isolated way to reproduce the behavior (example: GitHub repository with code isolated to the issue that anyone can clone to observe the problem)

This behavior is hard to reproduce with the current squirrel-version and requires #163 to be merged. However, let me try to describe the issue:

• Create two (or) more composables that contain split_by_worker_pytorch and/or split_by_rank_pytorch. Let's call them src_0 and src_1
• Construct an IterableSamplerSource samp_src (when it is picklable) using src_0 and src_1.
• Construct a pytorch dataloader dl.

A pseudocode would look like:

src_0 = IterableSource(samples_0).split_by_worker_pytorch()
src_1 = IterableSource(samples_0).split_by_worker_pytorch()

samp_src = IterableSamplerSource([src_0, src_1])

dl = torch.data.utils.Dataloader(samp_src.batched(batch_size).to_torch_iterable())

for b in dl:
   ....

• What version of squirrel-datasets and squirrel you're using, and the platform(s) you're running it on
• What packages or other dependencies you're using
  No additional packages
• The behavior you expect to see and the actual behavior

Expected: The for loop finishes successfully
Actual: Entering the for loop raises a PyTorchSplittingException

Current hypothesis:

The TorchIterable recursively checks the Composable source. It the source is a List[Composable] or simply an iterator the logic breaks and it will not find a split by rank / worker in the ancestors of the iterstream. Hence the Exception raises.

I tested this hypothesis locally on our internal tools. Disabling the _in_multi_rank_env and _in_multi_worker_env checks result in correct behavior of the iterstream.

The hosted documentation throws Mermaid syntax errors. e.g. https://squirrel-core.readthedocs.io/en/latest/advanced/store.html###

https://squirrel-core.readthedocs.io/en/latest/advanced/store.html mermaid figure raise a syntax error.

When reporting a bug, please be sure to include the following:

• A descriptive title
• An isolated way to reproduce the behavior (example: GitHub repository with code isolated to the issue that anyone can clone to observe the problem)
• What version of squirrel-datasets and squirrel you're using, and the platform(s) you're running it on
• What packages or other dependencies you're using
• The behavior you expect to see and the actual behavior

The current functionality of IterableSamplerSource is limited to subsampling datasets. However, there are many cases where you want to sub- and supersample a dataset mixture. A solution for this is to include a new sampler that allows multiplexing of datasets of vastly different sizes.

When you open an issue for a feature request, please add as much detail as possible:

• A descriptive title
• A description of the problem you're trying to solve, including why you think this is a problem
• An overview of the suggested solution
• If the feature changes current behavior, reasons why your solution is better