Raising error for inconsistent add_ml_model and add_script parameters (…

…#324)

Added error when user requests CPU with devices >1 within add_ml_model and add_script methods. Included more strict typing and tests.  

[ committed by @juliaputko ]
[ reviewed by @ashao, @ankona ]

smartsim/_core/utils/redis.py

@@ -158,9 +158,8 @@ def db_is_active(hosts: t.List[str], ports: t.List[int], num_shards: int) -> boo
 
 def set_ml_model(db_model: DBModel, client: Client) -> None:
     logger.debug(f"Adding DBModel named {db_model.name}")
-    devices = db_model._enumerate_devices()  # pylint: disable=protected-access
 
-    for device in devices:
+    for device in db_model.devices:
         try:
             if db_model.is_file:
                 client.set_model_from_file(
@@ -194,9 +193,7 @@ def set_ml_model(db_model: DBModel, client: Client) -> None:
 def set_script(db_script: DBScript, client: Client) -> None:
     logger.debug(f"Adding DBScript named {db_script.name}")
 
-    devices = db_script._enumerate_devices()  # pylint: disable=protected-access
-
-    for device in devices:
+    for device in db_script.devices:
         try:
             if db_script.is_file:
                 client.set_script_from_file(

smartsim/entity/dbobject.py

@@ -28,6 +28,7 @@
 
 from pathlib import Path
 from .._core.utils import init_default
+from ..error import SSUnsupportedError
 
 
 __all__ = ["DBObject", "DBModel", "DBScript"]
@@ -43,7 +44,7 @@ def __init__(
         name: str,
         func: t.Optional[str],
         file_path: t.Optional[str],
-        device: str,
+        device: t.Literal["CPU", "GPU"],
         devices_per_node: int,
     ) -> None:
         self.name = name
@@ -55,6 +56,11 @@ def __init__(
             self.file = self._check_filepath(file_path)
         self.device = self._check_device(device)
         self.devices_per_node = devices_per_node
+        self._check_devices(device, devices_per_node)
+
+    @property
+    def devices(self) -> t.List[str]:
+        return self._enumerate_devices()
 
     @property
     def is_file(self) -> bool:
@@ -95,8 +101,8 @@ def _check_filepath(file: str) -> Path:
         return file_path
 
     @staticmethod
-    def _check_device(device: str) -> str:
-        device = device.upper()
+    def _check_device(device: t.Literal["CPU", "GPU"]) -> str:
+        device = t.cast(t.Literal["CPU", "GPU"], device.upper())
         if not device.startswith("CPU") and not device.startswith("GPU"):
             raise ValueError("Device argument must start with either CPU or GPU")
         return device
@@ -109,21 +115,31 @@ def _enumerate_devices(self) -> t.List[str]:
         :return: list of device names
         :rtype: list[str]
         """
-        devices = []
-        if ":" in self.device and self.devices_per_node > 1:
-            msg = (
-                "Cannot set devices_per_node>1 if a device numeral is specified, "
-                f"the device was set to {self.device} and "
-                f"devices_per_node=={self.devices_per_node}"
-            )
-            raise ValueError(msg)
-        if self.device in ["CPU", "GPU"] and self.devices_per_node > 1:
-            for device_num in range(self.devices_per_node):
-                devices.append(f"{self.device}:{str(device_num)}")
-        else:
-            devices = [self.device]
 
-        return devices
+        if self.device == "GPU" and self.devices_per_node > 1:
+            return [
+                f"{self.device}:{str(device_num)}"
+                for device_num in range(self.devices_per_node)
+            ]
+
+        return [self.device]
+
+    @staticmethod
+    def _check_devices(
+        device: t.Literal["CPU", "GPU"], devices_per_node: int
+    ) -> None:
+        if devices_per_node == 1:
+            return
+
+        if ":" in device:
+            msg = "Cannot set devices_per_node>1 if a device numeral is specified, "
+            msg += f"the device was set to {device} and \
+                devices_per_node=={devices_per_node}"
+            raise ValueError(msg)
+        if device == "CPU":
+            raise SSUnsupportedError(
+                "Cannot set devices_per_node>1 if CPU is specified under devices"
+            )
 
 
 class DBScript(DBObject):
@@ -132,7 +148,7 @@ def __init__(
         name: str,
         script: t.Optional[str] = None,
         script_path: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU", "GPU"] = "CPU",
         devices_per_node: int = 1,
     ):
         """TorchScript code represenation
@@ -185,7 +201,7 @@ def __init__(
         backend: str,
         model: t.Optional[str] = None,
         model_file: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU", "GPU"] = "CPU",
         devices_per_node: int = 1,
         batch_size: int = 0,
         min_batch_size: int = 0,

smartsim/entity/ensemble.py

@@ -335,7 +335,7 @@ def add_ml_model(
         backend: str,
         model: t.Optional[str] = None,
         model_path: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
         batch_size: int = 0,
         min_batch_size: int = 0,
@@ -395,7 +395,7 @@ def add_script(
         name: str,
         script: t.Optional[str] = None,
         script_path: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript to launch with every entity belonging to this ensemble
@@ -439,7 +439,7 @@ def add_function(
         self,
         name: str,
         function: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript function to launch with every entity belonging to this ensemble

smartsim/entity/model.py

@@ -405,7 +405,7 @@ def add_ml_model(
         backend: str,
         model: t.Optional[str] = None,
         model_path: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
         batch_size: int = 0,
         min_batch_size: int = 0,
@@ -467,7 +467,7 @@ def add_script(
         name: str,
         script: t.Optional[str] = None,
         script_path: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript to launch with this Model instance
@@ -511,7 +511,7 @@ def add_function(
         self,
         name: str,
         function: t.Optional[str] = None,
-        device: str = "CPU",
+        device: t.Literal["CPU","GPU"] = "CPU",
         devices_per_node: int = 1,
     ) -> None:
         """TorchScript function to launch with this Model instance

tests/backends/test_dbmodel.py

@@ -36,6 +36,8 @@
 from smartsim.error.errors import SSUnsupportedError
 from smartsim.log import get_logger
 
+from smartsim.entity.dbobject import DBModel
+
 logger = get_logger(__name__)
 
 should_run_tf = True
@@ -793,3 +795,24 @@ def test_colocated_db_model_errors(fileutils, wlmutils, mlutils):
 
     with pytest.raises(SSUnsupportedError):
         colo_ensemble.add_model(colo_model)
+
+def test_inconsistent_params_db_model():
+    """Test error when devices_per_node parameter>1 when devices is set to CPU in DBModel"""
+
+    # Create and save ML model to filesystem
+    model, inputs, outputs = create_tf_cnn()
+    with pytest.raises(SSUnsupportedError) as ex:
+        db_model = DBModel(
+            "cnn",
+            "TF",
+            model=model,
+            device="CPU",
+            devices_per_node=2,
+            tag="test",
+            inputs=inputs,
+            outputs=outputs,
+        )
+    assert (
+            ex.value.args[0] 
+            == "Cannot set devices_per_node>1 if CPU is specified under devices"
+        )

tests/backends/test_dbscript.py

@@ -33,6 +33,8 @@
 from smartsim.error.errors import SSUnsupportedError
 from smartsim.log import get_logger
 
+from smartsim.entity.dbobject import DBScript
+
 logger = get_logger(__name__)
 
 should_run = True
@@ -578,3 +580,19 @@ def test_db_script_errors(fileutils, wlmutils, mlutils):
     # an in-memory script
     with pytest.raises(SSUnsupportedError):
         colo_ensemble.add_model(colo_model)
+
+def test_inconsistent_params_db_script(fileutils):
+    """Test error when devices_per_node>1 and when devices is set to CPU in DBScript constructor"""
+
+    torch_script = fileutils.get_test_conf_path("torchscript.py")
+    with pytest.raises(SSUnsupportedError) as ex:
+        db_script = DBScript(
+            name="test_script_db",
+            script_path = torch_script,
+            device="CPU",
+            devices_per_node=2,
+        )
+    assert (
+            ex.value.args[0] 
+            == "Cannot set devices_per_node>1 if CPU is specified under devices"
+        )