Emit intra_op_parallelism and JSON booleans from ds_nvme_tune

deepspeed/nvme/perf_generate_param.py

@@ -10,6 +10,8 @@
 import json
 from .parse_nvme_stats import READ_SPEED, WRITE_SPEED, get_sorted_results
 from .perf_sweep_utils import BENCH_LOG_DIR, READ_LOG_DIR, WRITE_LOG_DIR
+from deepspeed.runtime.swap_tensor.constants import AIO, AIO_BLOCK_SIZE, AIO_QUEUE_DEPTH, AIO_INTRA_OP_PARALLELISM, \
+    AIO_SINGLE_SUBMIT, AIO_OVERLAP_EVENTS
 
 
 def parse_arguments():
@@ -41,11 +43,11 @@ def validate_args(args):
 def convert_to_param(key):
     assert len(key) == 6
     return {
-        "single_submit": "true" if key[0] == "single" else "false",
-        "overlap_events": "true" if key[1] == "overlap" else "false",
-        "num_threads": int(key[5]),
-        "queue_depth": int(key[3]),
-        "block_size": int(key[4])
+        AIO_SINGLE_SUBMIT: key[0] == "single",
+        AIO_OVERLAP_EVENTS: key[1] == "overlap",
+        AIO_INTRA_OP_PARALLELISM: int(key[5]),
+        AIO_QUEUE_DEPTH: int(key[3]),
+        AIO_BLOCK_SIZE: int(key[4])
     }
 
 
@@ -72,7 +74,7 @@ def generate_aio_param(read_log_dir, write_log_dir):
             optimal_perf = value
             optimal_key = key
 
-    aio_param = {"aio": convert_to_param(optimal_key)}
+    aio_param = {AIO: convert_to_param(optimal_key)}
 
     read_perf_keys = {key[1:]: key for key in read_results.keys()}
     write_perf_keys = {key[1:]: key for key in write_results.keys()}

docs/_pages/config-json.md

@@ -710,7 +710,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s
   "aio": {
     "block_size": 1048576,
     "queue_depth": 8,
-    "thread_count": 1,
+    "intra_op_parallelism": 1,
     "single_submit": false,
     "overlap_events": true
   }
@@ -727,7 +727,7 @@ Configuring the asynchronous I/O module for offloading parameter and optimizer s
 | ---------------- | ------- |
 | I/O queue depth. | 8       |
 
-***thread_count***: [integer]
+***intra_op_parallelism***: [integer]
 
 | Description                                                               | Default |
 | ------------------------------------------------------------------------- | ------- |

docs/_tutorials/deepnvme.md

@@ -232,8 +232,8 @@ Running DeepNVMe performance tuning on ['/local_nvme/']
 Best performance (GB/sec): read =  3.69, write =  3.18
 {
    "aio": {
-      "single_submit": "false",
-      "overlap_events": "true",
+      "single_submit": false,
+      "overlap_events": true,
       "intra_op_parallelism": 8,
       "queue_depth": 32,
       "block_size": 1048576

tests/unit/v1/nvme/test_ds_nvme_tune.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import json
+import os
+
+import pytest
+
+from deepspeed.nvme.perf_generate_param import generate_main
+from deepspeed.nvme.perf_sweep_utils import READ_LOG_DIR, WRITE_LOG_DIR
+from deepspeed.runtime.swap_tensor.aio_config import get_aio_config
+
+# Log names as written by perf_run_sweep.get_log_file: <op>_<submit>_<ordering>_ftd<N>_d<depth>_bs<size>_p<parallel>
+FAST_SWEEP_POINT = "single_overlap_ftd1_d32_bs1M_p8"
+SLOW_SWEEP_POINT = "block_sequential_ftd1_d128_bs8M_p2"
+
+
+def _write_sweep_logs(log_dir, speeds):
+    for sub_dir, op, metric in ((READ_LOG_DIR, "read", "E2E Read Speed"), (WRITE_LOG_DIR, "write", "E2E Write Speed")):
+        os.makedirs(os.path.join(log_dir, sub_dir))
+        for sweep_point, speed in speeds.items():
+            with open(os.path.join(log_dir, sub_dir, f"{op}_{sweep_point}.txt"), "w") as f:
+                f.write(f"{metric} = {speed} GB/sec\n")
+
+
+# ds_nvme_tune prints its recommendation as a ds_config "aio" section, so it must survive the config parser
+# unchanged: every key spelled as get_aio_config reads it, and the flags as JSON booleans.
+@pytest.mark.parametrize("winner, expected", [
+    (FAST_SWEEP_POINT, {
+        "block_size": 1048576,
+        "queue_depth": 32,
+        "intra_op_parallelism": 8,
+        "single_submit": True,
+        "overlap_events": True
+    }),
+    (SLOW_SWEEP_POINT, {
+        "block_size": 8388608,
+        "queue_depth": 128,
+        "intra_op_parallelism": 2,
+        "single_submit": False,
+        "overlap_events": False
+    }),
+])
+def test_recommended_aio_section_round_trips_through_config(tmpdir, capsys, winner, expected):
+    speeds = {FAST_SWEEP_POINT: 1.0, SLOW_SWEEP_POINT: 1.0}
+    speeds[winner] = 5.0
+    _write_sweep_logs(str(tmpdir), speeds)
+
+    generate_main(str(tmpdir))
+    printed = capsys.readouterr().out
+    recommended = json.loads(printed[printed.index("{"):])
+
+    aio_config = get_aio_config(recommended)
+    assert {key: aio_config[key] for key in expected} == expected