Add nvme unit/perf tests (#993)

microsoft · Apr 22, 2021 · cf5ea89 · cf5ea89
1 parent 669028f
commit cf5ea89
Show file tree

Hide file tree

Showing 4 changed files with 356 additions and 24 deletions.
diff --git a/csrc/aio/py_test/ds_aio_basic.py b/csrc/aio/py_test/ds_aio_basic.py
@@ -8,7 +8,7 @@
 import torch
 import os
 import time
-from deepspeed.ops.aio import aio_read, aio_write
+from deepspeed.ops.aio import AsyncIOBuilder
 from multiprocessing import Pool, Barrier
 from test_ds_aio_utils import report_results, task_log, task_barrier
 
@@ -56,13 +56,13 @@ def post_basic(pool_params):
 def main_basic_read(pool_params):
     args, tid, ctxt = pool_params
     start_time = time.time()
-    aio_read(ctxt['buffer'],
-             ctxt['file'],
-             args.block_size,
-             args.queue_depth,
-             args.single_submit,
-             args.overlap_events,
-             args.validate)
+    AsyncIOBuilder().load().aio_read(ctxt['buffer'],
+                                     ctxt['file'],
+                                     args.block_size,
+                                     args.queue_depth,
+                                     args.single_submit,
+                                     args.overlap_events,
+                                     args.validate)
     end_time = time.time()
     ctxt['elapsed_sec'] += end_time - start_time
 
@@ -72,13 +72,13 @@ def main_basic_read(pool_params):
 def main_basic_write(pool_params):
     args, tid, ctxt = pool_params
     start_time = time.time()
-    aio_write(ctxt['buffer'],
-              ctxt['file'],
-              args.block_size,
-              args.queue_depth,
-              args.single_submit,
-              args.overlap_events,
-              args.validate)
+    AsyncIOBuilder().load().aio_write(ctxt['buffer'],
+                                      ctxt['file'],
+                                      args.block_size,
+                                      args.queue_depth,
+                                      args.single_submit,
+                                      args.overlap_events,
+                                      args.validate)
     end_time = time.time()
     ctxt['elapsed_sec'] += end_time - start_time
 

diff --git a/csrc/aio/py_test/ds_aio_handle.py b/csrc/aio/py_test/ds_aio_handle.py
@@ -8,8 +8,8 @@
 import torch
 import os
 import time
-from deepspeed.ops.aio import aio_handle
 from multiprocessing import Pool, Barrier
+from deepspeed.ops.aio import AsyncIOBuilder
 from test_ds_aio_utils import report_results, task_log, task_barrier
 
 
@@ -29,11 +29,11 @@ def pre_handle(args, tid, read_op):
     )
 
     io_parallel = args.io_parallel if args.io_parallel else 1
-    handle = aio_handle(args.block_size,
-                        args.queue_depth,
-                        args.single_submit,
-                        args.overlap_events,
-                        io_parallel)
+    handle = AsyncIOBuilder().load().aio_handle(args.block_size,
+                                                args.queue_depth,
+                                                args.single_submit,
+                                                args.overlap_events,
+                                                io_parallel)
     task_log(tid, f'created deepspeed aio handle')
 
     ctxt = {}

diff --git a/csrc/aio/py_test/test_ds_aio.py b/csrc/aio/py_test/test_ds_aio.py
@@ -12,12 +12,9 @@
 import sys
 from multiprocessing import Pool
 import multiprocessing as mp
-from deepspeed.ops.aio import aio_read, aio_write, aio_handle
 from ds_aio_basic import aio_basic_multiprocessing
 from ds_aio_handle import aio_handle_multiprocessing
 
-GB_DIVISOR = 1024**3
-
 
 def parse_arguments():
     parser = argparse.ArgumentParser()