refactor(generators): unify generators to work with any storage backend (argonne-lcf#329)

Every new storage backend required copy-pasting each generator into an
_XXX sibling file: npz_generator_s3.py, npy_generator_s3.py and so on.
The only difference was whether to write the output locally on disk,
directly via numpy/PIL, or via the storage interface.
This makes the pattern unsustainable: two duplicated formats today, more
with each new backend — incurring a significant maintenance burden.

Since all generators already had a storage instance and used it to
generate file names, we can leverage it.

The only set of generators now can check if the stroage is locally available
via `islocalfs` and use some optimisation, if any. If the storage is not local,
the sample serializes to io.BytesIO, call buf.getvalue(), and
delegate to self.storage.put_data().
All storage backends receive plain bytes as designed by the storage interface,
removing type inspection, seek() and getvalue() calls scattered across backends.

- FileStorage.put_data was never called, had text-mode open and a double
  get_uri call (once from the generator, once inside put_data itself).
  Now it is the default write path for LOCAL_FS, used by almost every
  workload config. get_data aligned to binary mode ("rb") for consistency.
- AIStoreStorage.put_data: remove isinstance dispatch, accept bytes directly.
- S3TorchStorage.put_data: remove data.getvalue() — just write data.
- generator_factory: removed S3/AIStore branching for NPZ, NPY, JPEG.
- factory referenced jpeg_generator_s3.JPEGGeneratorS3 which never existed;
  JPEG + S3/AIStore would crash at import time.

After this patch, adding a new storage backend requires no changes in any
generator. Adding a new data format automatically works with all backends.

Signed-off-by: Denis Barakhtanov <dbarahtanov@enakta.com>
Co-authored-by: Denis Barakhtanov <denis.barahtanov@gmail.com>

Author: enakta

dlio_benchmark/data_generator/generator_factory.py

@@ -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,29 +33,14 @@ def get_generator(type):
             from dlio_benchmark.data_generator.csv_generator import CSVGenerator
             return CSVGenerator()
         elif type == FormatType.NPZ:
-            # Use S3 generators for both S3 and AIStore
-            if _args.storage_type in (StorageType.S3, StorageType.AISTORE):
-                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:
-            # Use S3 generators for both S3 and AIStore
-            if _args.storage_type in (StorageType.S3, StorageType.AISTORE):
-                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:
-            # Use S3 generators for both S3 and AIStore
-            if _args.storage_type in (StorageType.S3, StorageType.AISTORE):
-                from dlio_benchmark.data_generator.jpeg_generator_s3 import JPEGGeneratorS3
-                return JPEGGeneratorS3()
-            else:
-                from dlio_benchmark.data_generator.jpeg_generator import JPEGGenerator
-                return JPEGGenerator()
+            from dlio_benchmark.data_generator.jpeg_generator import JPEGGenerator
+            return JPEGGenerator()
         elif type == FormatType.PNG:
             from dlio_benchmark.data_generator.png_generator import PNGGenerator
             return PNGGenerator()

dlio_benchmark/data_generator/jpeg_generator.py

@@ -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()

dlio_benchmark/data_generator/npy_generator.py

@@ -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()

dlio_benchmark/data_generator/npy_generator_s3.py

@@ -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()

dlio_benchmark/data_generator/npz_generator.py

@@ -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()

dlio_benchmark/data_generator/npz_generator_s3.py

@@ -16,7 +16,6 @@
 """
 
 import os
-import io
 import logging
 
 try:
@@ -89,12 +88,12 @@ def _clean_key(self, id):
         Extract the object key from a full S3/AIS URI.
 
         Why this is needed:
-        - S3 generators (NPYGeneratorS3, NPYReaderS3) pass full URIs like:
+        - Generators call storage.get_uri(file_list[i]) which pass full URIs like:
           "s3://dlio-benchmark-native/train/img_08_of_16.npy"
           or "ais://dlio-benchmark-native/train/img_08_of_16.npy"
         - AIStore SDK expects just the object key:
           "train/img_08_of_16.npy"
-        - This method strips the "s3://" or "ais://" prefix and bucket name
+        - This method strips the scheme and bucket name from the URI
 
         Handles: 
           s3://bucket/path/file.ext -> path/file.ext
@@ -226,22 +225,12 @@ def put_data(self, id, data, offset=None, length=None):
 
             key = self._clean_key(id)
 
-            # Convert data to bytes
-            if isinstance(data, io.BytesIO):
-                data.seek(0)
-                body = data.read()
-            elif isinstance(data, bytes):
-                body = data
-            else:
-                body = bytes(data)
-
-            # Put object
             obj = self.bucket.object(key)
-            obj.get_writer().put_content(body)
+            obj.get_writer().put_content(data)
 
             # TODO: add offset and length support
 
-            logging.debug(f"Successfully uploaded: {key} ({len(body)} bytes)")
+            logging.debug(f"Successfully uploaded: {key} ({len(data)} bytes)")
             return True
         except Exception as e:
             logging.error(f"Error putting data to {id}: {e}")

dlio_benchmark/data_generator/png_generator.py

@@ -90,12 +90,16 @@ def delete_node(self, id):
     # TODO Handle partial read and writes
     @dlp.log
     def put_data(self, id, data, offset=None, length=None):
-        with open(self.get_uri(id), "w") as fd:
+        # id is the fully-resolved path (callers call get_uri() before put_data).
+        # Do NOT call self.get_uri(id) here — that would double-prefix the namespace.
+        with open(id, "wb") as fd:
             fd.write(data)
 
     @dlp.log
     def get_data(self, id, data, offset=None, length=None):
-        with open(self.get_uri(id), "r") as fd:
+        # id is the fully-resolved path (callers call get_uri() before put_data).
+        # Do NOT call self.get_uri(id) here — that would double-prefix the namespace.
+        with open(id, "rb") as fd:
             data = fd.read()
         return data
 
@@ -105,3 +109,6 @@ def isfile(self, id):
 
     def get_basename(self, id):
         return os.path.basename(id)
+
+    def islocalfs(self):
+        return True

dlio_benchmark/storage/aistore_storage.py

@@ -131,7 +131,7 @@ def put_data(self, id, data, offset=None, length=None):
 
         bucket_name = parsed.netloc
         writer = self.s3_client.put_object(bucket_name, id)
-        writer.write(data.getvalue())
+        writer.write(data)
         writer.close()
         return None