Fix multi-host execution: --hosts parsing (#322) and MPI cluster collector staging (#303)

mlpstorage

@@ -1,2 +1,12 @@
-#! /bin/bash
-uv run python3 -m mlpstorage_py.main $*
+#!/usr/bin/env bash
+# Wrapper for running mlpstorage directly from a checkout via `uv run`.
+# Preferred install path is `uv pip install -e .` (uses [project.scripts]).
+#
+# IMPORTANT fixes vs prior version (see issue #322):
+#   - "$@" (not $*) preserves argument quoting; args with spaces stay intact.
+#   - `--` separator prevents `uv run` from intercepting flags meant for mlpstorage
+#     (e.g. --project, --directory, --no-sync, --python).
+#   - --project pins the uv project to this checkout regardless of $PWD.
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec uv run --project "$SCRIPT_DIR" -- python3 -m mlpstorage_py.main "$@"

mlpstorage_py/cli/common_args.py

@@ -51,11 +51,15 @@
     ),
     'client_hosts': (
         "Space-separated list of IP addresses or hostnames of the participating hosts. "
-        "\nExample: '--hosts 192.168.1.1 192.168.1.2 192.168.1.3' or '--hosts host1 host2 host3'. Slots can "
+        "\nExample: '--hosts 192.168.1.1 192.168.1.2 192.168.1.3' or '--hosts host1 host2 host3'. "
+        "Comma-separated values are also accepted: '--hosts host1,host2,host3'. "
+        "Slots can "
         "be specified by appending ':<num_slots>' to a hostname like so: '--hosts host1:2 host2:2'. This "
         "example will run 2 accelerators on each host. If slots are not specified the number of processes "
         "will be equally distributed across the hosts with any remainder being distributed evenly on the "
-        "remaining hosts in the order they are listed."
+        "remaining hosts in the order they are listed. "
+        "\nDo NOT use '--hosts=h1 h2' (with '=' and space); argparse will only bind 'h1' to --hosts "
+        "and treat 'h2' as a stray positional argument. Use '--hosts h1 h2' or '--hosts=h1,h2' instead."
     ),
     'category': "Benchmark category to be submitted.",
     'results_dir': "Directory where the benchmark results will be saved.",

mlpstorage_py/cli_parser.py

@@ -6,6 +6,7 @@
 """
 
 import argparse
+import re
 import sys
 
 from mlpstorage_py import VERSION
@@ -243,14 +244,29 @@ def update_args(args):
         flattened_mpi_params = [item for sublist in args.mpi_params for item in sublist]
         setattr(args,'mpi_params', flattened_mpi_params)
 
-    if hasattr(args, 'hosts'):
-        print(f'Hosts is: {args.hosts}')
-        # hosts can be comma separated string or a list of strings. If it's a string, it is still a list of length 1
-        if len(args.hosts) == 1 and isinstance(args.hosts[0], str):
-            setattr(args, 'hosts', args.hosts[0].split(','))
-        print(f'Hosts is: {args.hosts}')
-
-    if not hasattr(args, "num_client_hosts") and hasattr(args, "hosts"):
+    if hasattr(args, 'hosts') and args.hosts is not None:
+        # Accept any of the following equivalent forms and normalize to a clean list:
+        #   --hosts h1 h2 h3              -> ['h1', 'h2', 'h3']
+        #   --hosts h1,h2,h3              -> ['h1', 'h2', 'h3']
+        #   --hosts 'h1 h2 h3'            -> ['h1', 'h2', 'h3']   (quoted, e.g. from YAML)
+        #   --hosts='h1,h2,h3'            -> ['h1', 'h2', 'h3']   (DLIO subprocess form)
+        #   --hosts='h1 h2 h3'            -> ['h1', 'h2', 'h3']   (quoted after '=')
+        # This defends against the argparse + nargs='+' + '=' interaction documented in
+        # https://github.com/mlcommons/storage/issues/322.
+        raw = args.hosts if isinstance(args.hosts, list) else [args.hosts]
+        normalized = []
+        for item in raw:
+            if not isinstance(item, str):
+                continue
+            for tok in re.split(r'[,\s]+', item.strip()):
+                if tok:
+                    normalized.append(tok)
+        if not normalized:
+            print("ERROR: --hosts is empty after parsing", file=sys.stderr)
+            sys.exit(EXIT_CODE.INVALID_ARGUMENTS)
+        args.hosts = normalized
+
+    if hasattr(args, 'hosts') and getattr(args, 'num_client_hosts', None) is None:
         setattr(args, "num_client_hosts", len(args.hosts))