mlcommons
diff --git a/‎.github/workflows/ci.yml‎
Lines changed: 23 additions & 0 deletions b/‎.github/workflows/ci.yml‎
Lines changed: 23 additions & 0 deletions
diff --git a/‎README.md‎
Lines changed: 12 additions & 2 deletions b/‎README.md‎
Lines changed: 12 additions & 2 deletions
diff --git a/‎dlio_benchmark/common/enumerations.py‎
Lines changed: 1 addition & 0 deletions b/‎dlio_benchmark/common/enumerations.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎dlio_benchmark/data_generator/generator_factory.py‎
Lines changed: 5 additions & 16 deletions b/‎dlio_benchmark/data_generator/generator_factory.py‎
Lines changed: 5 additions & 16 deletions
diff --git a/‎dlio_benchmark/data_generator/jpeg_generator.py‎
Lines changed: 5 additions & 1 deletion b/‎dlio_benchmark/data_generator/jpeg_generator.py‎
Lines changed: 5 additions & 1 deletion
diff --git a/‎dlio_benchmark/data_generator/npy_generator.py‎
Lines changed: 5 additions & 1 deletion b/‎dlio_benchmark/data_generator/npy_generator.py‎
Lines changed: 5 additions & 1 deletion
diff --git a/‎dlio_benchmark/data_generator/npy_generator_s3.py‎
Lines changed: 0 additions & 57 deletions b/‎dlio_benchmark/data_generator/npy_generator_s3.py‎
Lines changed: 0 additions & 57 deletions
diff --git a/‎dlio_benchmark/data_generator/npz_generator.py‎
Lines changed: 6 additions & 2 deletions b/‎dlio_benchmark/data_generator/npz_generator.py‎
Lines changed: 6 additions & 2 deletions
diff --git a/‎dlio_benchmark/data_generator/npz_generator_s3.py‎
Lines changed: 0 additions & 59 deletions b/‎dlio_benchmark/data_generator/npz_generator_s3.py‎
Lines changed: 0 additions & 59 deletions
diff --git a/‎dlio_benchmark/data_generator/png_generator.py‎
Lines changed: 5 additions & 1 deletion b/‎dlio_benchmark/data_generator/png_generator.py‎
Lines changed: 5 additions & 1 deletion
@@ -358,3 +358,26 @@ jobs:
         run: |
           source ${VENV_PATH}/bin/activate
           mpirun -np 1 pytest -k test_s3_checkpoint_step -v
+      # AIStore-specific tests (mock-based, no real cluster needed)
+      - name: test_aistore_gen_data
+        run: |
+          source ${VENV_PATH}/bin/activate
+          mpirun -np 1 pytest -k test_aistore_gen_data[npy-pytorch] -v
+          mpirun -np 1 pytest -k test_aistore_gen_data[npz-pytorch] -v
+      - name: test_aistore_train
+        run: |
+          source ${VENV_PATH}/bin/activate
+          mpirun -np 1 pytest -k test_aistore_train[npy-pytorch-True] -v
+          mpirun -np 1 pytest -k test_aistore_train[npz-pytorch-True] -v
+          mpirun -np 1 pytest -k test_aistore_train[npy-pytorch-False] -v
+          mpirun -np 1 pytest -k test_aistore_train[npz-pytorch-False] -v
+      - name: test_aistore_eval
+        run: |
+          source ${VENV_PATH}/bin/activate
+          mpirun -np 1 pytest -k test_aistore_eval -v
+      - name: test_aistore_multi_threads
+        run: |
+          source ${VENV_PATH}/bin/activate
+          mpirun -np 1 pytest -k test_aistore_multi_threads[pytorch-0] -v
+          mpirun -np 1 pytest -k test_aistore_multi_threads[pytorch-1] -v
+          mpirun -np 1 pytest -k test_aistore_multi_threads[pytorch-2] -v
@@ -17,6 +17,14 @@ pip install .
 dlio_benchmark ++workload.workflow.generate_data=True
 ```
 
+### Bare metal installation with AIStore support
+
+```bash
+git clone https://github.com/argonne-lcf/dlio_benchmark
+cd dlio_benchmark/
+pip install .[aistore]
+```
+
 ### Bare metal installation with profiler
 
 ```bash
@@ -150,7 +158,9 @@ The YAML file is loaded through hydra (https://hydra.cc/). The default setting a
 
 * We assume the data/label pairs are stored in the same file. Storing data and labels in separate files will be supported in future.
 
-* File format support: we only support tfrecord, hdf5, npz, csv, jpg, jpeg formats. Other data formats can be extended. 
+* File format support: we only support tfrecord, hdf5, npz, csv, jpg, jpeg formats. Other data formats can be extended.
+
+* Storage backend support: we support local filesystem, AWS S3, and AIStore as storage backends. Other storage backends can be extended.
 
 * Data Loader support: we support reading datasets using TensorFlow tf.data data loader, PyTorch DataLoader, and a set of custom data readers implemented in ./reader. For TensorFlow tf.data data loader, PyTorch DataLoader  
   - We have complete support for tfrecord format in TensorFlow data loader. 
@@ -163,7 +173,7 @@ General new features needed including:
 * support for new workloads: if you think that your workload(s) would be interested to the public, and would like to provide the yaml file to be included in the repo, please submit an issue.  
 * support for new data loaders, such as DALI loader, MxNet loader, etc
 * support for new frameworks, such as MxNet
-* support for noval file systems or storage, such as AWS S3. 
+* support for novel file systems or storage, such as AWS S3, AIStore, etc.
 * support for loading new data formats. 
 
 If you would like to contribute, please submit an issue to https://github.com/argonne-lcf/dlio_benchmark/issues, and contact ALCF DLIO team, Huihuo Zheng at huihuo.zheng@anl.gov
 
@@ -58,6 +58,7 @@ class StorageType(Enum):
     LOCAL_FS = 'local_fs'
     PARALLEL_FS = 'parallel_fs'
     S3 = 's3'
+    AISTORE = 'aistore'
 
     def __str__(self):
         return self.value
 
@@ -14,9 +14,7 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
-from dlio_benchmark.utils.config import ConfigArguments
-
-from dlio_benchmark.common.enumerations import FormatType, StorageType
+from dlio_benchmark.common.enumerations import FormatType
 from dlio_benchmark.common.error_code import ErrorCodes
 
 class GeneratorFactory(object):
@@ -25,7 +23,6 @@ def __init__(self):
 
     @staticmethod
     def get_generator(type):
-        _args = ConfigArguments.get_instance()
         if type == FormatType.TFRECORD:
             from dlio_benchmark.data_generator.tf_generator import TFRecordGenerator
             return TFRecordGenerator()
@@ -36,19 +33,11 @@ def get_generator(type):
             from dlio_benchmark.data_generator.csv_generator import CSVGenerator
             return CSVGenerator()
         elif type == FormatType.NPZ:
-            if _args.storage_type == StorageType.S3:
-                from dlio_benchmark.data_generator.npz_generator_s3 import NPZGeneratorS3
-                return NPZGeneratorS3()
-            else:
-                from dlio_benchmark.data_generator.npz_generator import NPZGenerator
-                return NPZGenerator()
+            from dlio_benchmark.data_generator.npz_generator import NPZGenerator
+            return NPZGenerator()
         elif type == FormatType.NPY:
-            if _args.storage_type == StorageType.S3:
-                from dlio_benchmark.data_generator.npy_generator_s3 import NPYGeneratorS3
-                return NPYGeneratorS3()
-            else:
-                from dlio_benchmark.data_generator.npy_generator import NPYGenerator
-                return NPYGenerator()            
+            from dlio_benchmark.data_generator.npy_generator import NPYGenerator
+            return NPYGenerator()
         elif type == FormatType.JPEG:
             from dlio_benchmark.data_generator.jpeg_generator import JPEGGenerator
             return JPEGGenerator()
 
@@ -14,6 +14,7 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
+import io
 import numpy as np
 import PIL.Image as im
 
@@ -53,5 +54,8 @@ def generate(self):
                 self.logger.info(f"Generated file {i}/{self.total_files_to_generate}")
             out_path_spec = self.storage.get_uri(self._file_list[i])
             progress(i+1, self.total_files_to_generate, "Generating JPEG Data")
-            img.save(out_path_spec, format='JPEG', bits=8)
+            output = out_path_spec if self.storage.islocalfs() else io.BytesIO()
+            img.save(output, format='JPEG', bits=8)
+            if not self.storage.islocalfs():
+                self.storage.put_data(out_path_spec, output.getvalue())
         np.random.seed()
@@ -14,6 +14,7 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
+import io
 import numpy as np
 
 from dlio_benchmark.data_generator.data_generator import DataGenerator
@@ -49,5 +50,8 @@ def generate(self):
 
             out_path_spec = self.storage.get_uri(self._file_list[i])
             progress(i+1, self.total_files_to_generate, "Generating NPY Data")
-            np.save(out_path_spec, records)
+            output = out_path_spec if self.storage.islocalfs() else io.BytesIO()
+            np.save(output, records)
+            if not self.storage.islocalfs():
+                self.storage.put_data(out_path_spec, output.getvalue())
         np.random.seed()
@@ -14,6 +14,7 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
+import io
 import numpy as np
 
 from dlio_benchmark.common.enumerations import Compression
@@ -48,8 +49,11 @@ def generate(self):
                 records = gen_random_tensor(shape=(dim_, dim[2*i+1], self.num_samples), dtype=self._args.record_element_dtype, rng=rng)
             out_path_spec = self.storage.get_uri(self._file_list[i])
             progress(i+1, self.total_files_to_generate, "Generating NPZ Data")
+            output = out_path_spec if self.storage.islocalfs() else io.BytesIO()
             if self.compression != Compression.ZIP:
-                np.savez(out_path_spec, x=records, y=record_labels)
+                np.savez(output, x=records, y=record_labels)
             else:
-                np.savez_compressed(out_path_spec, x=records, y=record_labels)
+                np.savez_compressed(output, x=records, y=record_labels)
+            if not self.storage.islocalfs():
+                self.storage.put_data(out_path_spec, output.getvalue())
         np.random.seed()
@@ -14,6 +14,7 @@
    See the License for the specific language governing permissions and
    limitations under the License.
 """
+import io
 import numpy as np
 import PIL.Image as im
 
@@ -49,5 +50,8 @@ def generate(self):
                 self.logger.info(f"Generated file {i}/{self.total_files_to_generate}")
             out_path_spec = self.storage.get_uri(self._file_list[i])
             progress(i+1, self.total_files_to_generate, "Generating PNG Data")
-            img.save(out_path_spec, format='PNG', bits=8)
+            output = out_path_spec if self.storage.islocalfs() else io.BytesIO()
+            img.save(output, format='PNG', bits=8)
+            if not self.storage.islocalfs():
+                self.storage.put_data(out_path_spec, output.getvalue())
         np.random.seed()