fixes per review comments, fix bug in from_run_args

Author: ankona

doc/dragon.rst

@@ -104,9 +104,9 @@ as a list of device indices.
 
 .. note::
 
-    SmartSim submits jobs in the order they are received. On a heterogeneous system, SmartSim
-    will attempt to allocate non-GPU nodes first. However, a process may be allocated to a GPU
-    node if only GPU nodes are available, regardless of the requested features.
+    SmartSim submits jobs in the order they are received. However, a process may
+    be allocated to a GPU node if only GPU nodes are available, regardless of
+    the requested features.
 
     To ensure a process is allocated to a specific node, configure a hostname constraint.
 

smartsim/_core/launcher/dragon/dragonBackend.py

@@ -429,7 +429,7 @@ def _stop_steps(self) -> None:
 
     @staticmethod
     def create_run_policy(
-        request: DragonRunRequest, node_name: str
+        request: DragonRequest, node_name: str
     ) -> "dragon_policy.Policy":
         if isinstance(request, DragonRunRequest):
             run_request: DragonRunRequest = request
@@ -439,18 +439,30 @@ def create_run_policy(
             cpu_affinity: t.List[int] = []
             gpu_affinity: t.List[int] = []
 
+            # Customize policy only if the client requested it, otherwise use default
             if run_request.policy is not None:
+                # Don't change device from default unless specified
+                if run_request.policy.device == "gpu":
+                    device = dragon_policy.Policy.Device.GPU
+                elif run_request.policy.device == "cpu":
+                    device = dragon_policy.Policy.Device.CPU
+
+                # affinities are not mutually exclusive. If specified, both are used
                 if run_request.policy.cpu_affinity:
                     affinity = dragon_policy.Policy.Affinity.SPECIFIC
                     cpu_affinity = run_request.policy.cpu_affinity
-                    device = dragon_policy.Policy.Device.CPU
 
                 if run_request.policy.gpu_affinity:
                     affinity = dragon_policy.Policy.Affinity.SPECIFIC
                     gpu_affinity = run_request.policy.gpu_affinity
-                    device = dragon_policy.Policy.Device.GPU
 
-            if affinity != dragon_policy.Policy.Affinity.DEFAULT:
+                # TODO: if device == "cpu" and gpu_affinity is set, will it fail?
+                # Should we override device in the gpu_affinity setter or raise?
+
+            if (
+                affinity != dragon_policy.Policy.Affinity.DEFAULT
+                or device != dragon_policy.Policy.Device.DEFAULT
+            ):
                 return dragon_policy.Policy(
                     placement=dragon_policy.Policy.Placement.HOST_NAME,
                     host_name=node_name,

smartsim/_core/schemas/dragonRequests.py

@@ -26,9 +26,10 @@
 
 import typing as t
 
-from pydantic import BaseModel, Field, NonNegativeInt, PositiveInt
+from pydantic import BaseModel, Field, NonNegativeInt, PositiveInt, ValidationError
 
 import smartsim._core.schemas.utils as _utils
+from smartsim.error.errors import SmartSimError
 
 # Black and Pylint disagree about where to put the `...`
 # pylint: disable=multiple-statements
@@ -42,31 +43,50 @@ class DragonRequest(BaseModel): ...
 class DragonRunPolicy(BaseModel):
     """Policy specifying hardware constraints when running a Dragon job"""
 
-    device: t.Literal["cpu", "gpu"] = Field(default="cpu")
+    device: t.Literal["cpu", "gpu", ""] = Field(default="")
+    """Specify a device category on which to run the job. Uses system default
+    if not specified"""
     cpu_affinity: t.List[NonNegativeInt] = Field(default_factory=list)
+    """List of CPU indices to which the job should be pinned"""
     gpu_affinity: t.List[NonNegativeInt] = Field(default_factory=list)
+    """List of GPU indices to which the job should be pinned"""
 
     @staticmethod
     def from_run_args(
         run_args: t.Dict[str, t.Union[int, str, float, None]]
     ) -> "DragonRunPolicy":
+        """Create a DragonRunPolicy from a dictionary of run arguments"""
         features: str = str(run_args.get("node-feature", ""))
 
-        device = "gpu" if "gpu" in features else "cpu"
+        device = ""
+        if "gpu" in features:
+            device = "gpu"
+        elif "cpu" in features:
+            device = "cpu"
 
-        gpu_args = str(run_args.get("gpu-affinity", ""))
-        cpu_args = str(run_args.get("cpu-affinity", ""))
-        gpu_affinity = [x for x in gpu_args.split(",") if x]
-        cpu_affinity = [x for x in cpu_args.split(",") if x]
+        gpu_args = ""
+        if gpu_arg_value := run_args.get("gpu-affinity", None):
+            gpu_args = str(gpu_arg_value)
 
-        if device == "cpu" and not (cpu_affinity or gpu_affinity):
+        cpu_args = ""
+        if cpu_arg_value := run_args.get("cpu-affinity", None):
+            cpu_args = str(cpu_arg_value)
+
+        # run args converted to a string must be split back into a list[int]
+        gpu_affinity = [int(x.strip()) for x in gpu_args.split(",") if x]
+        cpu_affinity = [int(x.strip()) for x in cpu_args.split(",") if x]
+
+        if device == "" and not (cpu_affinity or gpu_affinity):
             return DragonRunPolicy()
 
-        return DragonRunPolicy(
-            device=device,
-            cpu_affinity=cpu_affinity,
-            gpu_affinity=gpu_affinity,
-        )
+        try:
+            return DragonRunPolicy(
+                device=device,
+                cpu_affinity=cpu_affinity,
+                gpu_affinity=gpu_affinity,
+            )
+        except ValidationError as ex:
+            raise SmartSimError("Unable to build DragonRunPolicy") from ex
 
 
 class DragonRunRequestView(DragonRequest):

smartsim/settings/dragonRunSettings.py

@@ -83,9 +83,12 @@ def set_tasks_per_node(self, tasks_per_node: int) -> None:
 
     @override
     def set_node_feature(self, feature_list: t.Union[str, t.List[str]]) -> None:
-        """Add a node feature requirement
+        """Specify the node feature for this job
 
-        :param tasks_per_node: number of tasks per node
+        :param feature_list: a collection of strings representing the
+        required node features.
+
+        Currently supported node features are: "gpu"
         """
         if isinstance(feature_list, str):
             feature_list = feature_list.strip().split()

tests/test_dragon_client.py

@@ -180,7 +180,7 @@ def test_dragon_client_main(
         assert isinstance(request_arg, DragonRunRequest)
 
         policy = request_arg.policy
-        assert policy.device == "cpu" if not policy.gpu_affinity else "gpu"
+        assert policy.device == "" if not policy.gpu_affinity else "gpu"
 
         # make sure each policy has been read in correctly with valid affinity indices
         assert len(policy.cpu_affinity) == len(set(policy.cpu_affinity))

tests/test_dragon_launcher.py

@@ -701,7 +701,7 @@ def test_run_step_success(test_dir: str) -> None:
     assert req_name.startswith(step0.name)
 
     req_policy_device = dragon_run_request.policy.device
-    assert req_policy_device == "cpu"  # default device is cpu
+    assert req_policy_device == ""  # default device is cpu
 
     req_policy_cpu_affinity = dragon_run_request.policy.cpu_affinity
     assert not req_policy_cpu_affinity  # default should be empty list