Allow launcher_extra to split quoted values

legate/driver/args.py

@@ -162,7 +162,9 @@
     action="append",
     default=[],
     required=False,
-    help="Specify extra flags for Nsight Systems",
+    help="Specify extra flags for Nsight Systems (can appear more than once). "
+    "Multiple arguments may be provided together in a quoted string "
+    "(arguments with spaces inside must be additionally quoted)",
 )
 
 logging = parser.add_argument_group("Logging")

legate/driver/config.py

@@ -17,6 +17,7 @@
 """
 from __future__ import annotations
 
+import shlex
 from argparse import Namespace
 from dataclasses import dataclass
 from functools import cached_property
@@ -44,6 +45,16 @@ class MultiNode(DataclassMixin):
     launcher: LauncherType
     launcher_extra: list[str]
 
+    def __post_init__(self, **kw: dict[str, Any]) -> None:
+        # fix up launcher_extra to automaticaly handle quoted strings with
+        # internal whitespace, have to use __setattr__ for frozen
+        # https://docs.python.org/3/library/dataclasses.html#frozen-instances
+        if self.launcher_extra:
+            ex: list[str] = sum(
+                (shlex.split(x) for x in self.launcher_extra), []
+            )
+            object.__setattr__(self, "launcher_extra", ex)
+
     @property
     def ranks(self) -> int:
         return self.nodes * self.ranks_per_node
@@ -84,6 +95,14 @@ class Profiling(DataclassMixin):
     nsys_targets: str  # TODO: multi-choice
     nsys_extra: list[str]
 
+    def __post_init__(self, **kw: dict[str, Any]) -> None:
+        # fix up nsys_extra to automaticaly handle quoted strings with
+        # internal whitespace, have to use __setattr__ for frozen
+        # https://docs.python.org/3/library/dataclasses.html#frozen-instances
+        if self.nsys_extra:
+            ex: list[str] = sum((shlex.split(x) for x in self.nsys_extra), [])
+            object.__setattr__(self, "nsys_extra", ex)
+
 
 @dataclass(frozen=True)
 class Logging(DataclassMixin):

legate/util/shared_args.py

@@ -96,7 +96,8 @@
         default=[],
         required=False,
         help="additional argument to pass to the launcher (can appear more "
-        "than once)",
+        "than once). Multiple arguments may be provided together in a quoted "
+        "string (arguments with spaces inside must be additionally quoted)",
     ),
 )
 

tests/unit/legate/driver/test_config.py

@@ -55,6 +55,56 @@ def test_fields(self) -> None:
     def test_mixin(self) -> None:
         assert issubclass(m.MultiNode, DataclassMixin)
 
+    @pytest.mark.parametrize(
+        "extra",
+        (["a"], ["a", "b c"], ["a", "b c", "d e"], ["a", "b c", "d e", "f"]),
+    )
+    def test_launcher_extra_fixup_basic(self, extra) -> None:
+        mn = m.MultiNode(
+            nodes=1,
+            ranks_per_node=1,
+            not_control_replicable=False,
+            launcher="launcher",
+            launcher_extra=extra,
+        )
+        assert mn.launcher_extra == sum((x.split() for x in extra), [])
+
+    def test_launcher_extra_fixup_complex(self) -> None:
+        mn = m.MultiNode(
+            nodes=1,
+            ranks_per_node=1,
+            not_control_replicable=False,
+            launcher="launcher",
+            launcher_extra=[
+                "-H g0002,g0002 -X SOMEENV --fork",
+                "-bind-to none",
+            ],
+        )
+        assert mn.launcher_extra == [
+            "-H",
+            "g0002,g0002",
+            "-X",
+            "SOMEENV",
+            "--fork",
+            "-bind-to",
+            "none",
+        ]
+
+    def test_launcher_extra_fixup_quoted(self) -> None:
+        mn = m.MultiNode(
+            nodes=1,
+            ranks_per_node=1,
+            not_control_replicable=False,
+            launcher="launcher",
+            launcher_extra=[
+                "-f 'some path with spaces/foo.txt'",
+            ],
+        )
+        assert mn.launcher_extra == [
+            "-f",
+            "some path with spaces/foo.txt",
+        ]
+
 
 class TestBinding:
     def test_fields(self) -> None:
@@ -111,6 +161,56 @@ def test_fields(self) -> None:
     def test_mixin(self) -> None:
         assert issubclass(m.Profiling, DataclassMixin)
 
+    @pytest.mark.parametrize(
+        "extra",
+        (["a"], ["a", "b c"], ["a", "b c", "d e"], ["a", "b c", "d e", "f"]),
+    )
+    def test_nsys_extra_fixup_basic(self, extra) -> None:
+        p = m.Profiling(
+            profile=True,
+            nvprof=True,
+            nsys=True,
+            nsys_targets="foo,bar",
+            nsys_extra=extra,
+        )
+        assert p.nsys_extra == sum((x.split() for x in extra), [])
+
+    def test_nsys_extra_fixup_complex(self) -> None:
+        p = m.Profiling(
+            profile=True,
+            nvprof=True,
+            nsys=True,
+            nsys_targets="foo,bar",
+            nsys_extra=[
+                "-H g0002,g0002 -X SOMEENV --fork",
+                "-bind-to none",
+            ],
+        )
+        assert p.nsys_extra == [
+            "-H",
+            "g0002,g0002",
+            "-X",
+            "SOMEENV",
+            "--fork",
+            "-bind-to",
+            "none",
+        ]
+
+    def test_nsys_extra_fixup_quoted(self) -> None:
+        p = m.Profiling(
+            profile=True,
+            nvprof=True,
+            nsys=True,
+            nsys_targets="foo,bar",
+            nsys_extra=[
+                "-f 'some path with spaces/foo.txt'",
+            ],
+        )
+        assert p.nsys_extra == [
+            "-f",
+            "some path with spaces/foo.txt",
+        ]
+
 
 class TestLogging:
     def test_fields(self) -> None: