Fix braceexpand memory explosion for complex urls (#336)

* fix braceexpand memory explosion for complex urls

* add tests

Author: gabrielilharco

src/training/data.py

@@ -9,7 +9,6 @@
 from dataclasses import dataclass
 from multiprocessing import Value
 
-import braceexpand
 import numpy as np
 import pandas as pd
 import torch
@@ -73,8 +72,8 @@ def set_epoch(self, epoch):
 
 
 def get_dataset_size(shards):
-    shards_list = list(braceexpand.braceexpand(shards))
-    dir_path = os.path.dirname(shards)
+    shards_list = wds.shardlists.expand_urls(shards)
+    dir_path = os.path.dirname(shards_list[0])
     sizes_filename = os.path.join(dir_path, 'sizes.json')
     len_filename = os.path.join(dir_path, '__len__')
     if os.path.exists(sizes_filename):

tests/test_num_shards.py

@@ -0,0 +1,20 @@
+import pytest
+
+from training.data import get_dataset_size
+
+@pytest.mark.parametrize(
+    "shards,expected_size",
+    [
+        ('/path/to/shard.tar', 1),
+        ('/path/to/shard_{000..000}.tar', 1),
+        ('/path/to/shard_{000..009}.tar', 10),
+        ('/path/to/shard_{000..009}_{000..009}.tar', 100),
+        ('/path/to/shard.tar::/path/to/other_shard_{000..009}.tar', 11),
+        ('/path/to/shard_{000..009}.tar::/path/to/other_shard_{000..009}.tar', 20),
+        (['/path/to/shard.tar'], 1),
+        (['/path/to/shard.tar', '/path/to/other_shard.tar'], 2),
+    ]
+)
+def test_num_shards(shards, expected_size):
+    _, size = get_dataset_size(shards)
+    assert size == expected_size, f'Expected {expected_size} for {shards} but found {size} instead.'