Remove device attribute from schemas (#619)

This PR removes `device` from the schemas, MessageHandler, and tests.

Author: AlyssaCote

doc/changelog.md

@@ -16,6 +16,7 @@ Description
 - Add ML worker manager, sample worker, and feature store
 - Added schemas and MessageHandler class for de/serialization of
   inference requests and response messages
+- Removed device from schemas, MessageHandler and tests
 
 
 ### Development branch

smartsim/_core/mli/infrastructure/control/workermanager.py

@@ -64,7 +64,6 @@ def deserialize_message(
 
     request = MessageHandler.deserialize_request(data_blob)
     # return request
-    device = request.device
     model_key: t.Optional[str] = None
     model_bytes: t.Optional[bytes] = None
 
@@ -106,7 +105,6 @@ def deserialize_message(
         input_keys=input_keys,
         raw_model=model_bytes,
         batch_size=0,
-        device=device,
     )
     return inference_request
 

smartsim/_core/mli/infrastructure/worker/worker.py

@@ -50,7 +50,6 @@ def __init__(
         output_keys: t.Optional[t.List[str]] = None,
         raw_model: t.Optional[bytes] = None,
         batch_size: int = 0,
-        device: t.Optional[str] = None,
     ):
         """Initialize the object"""
         self.model_key = model_key
@@ -61,7 +60,6 @@ def __init__(
         self.input_meta = input_meta or []
         self.output_keys = output_keys or []
         self.batch_size = batch_size
-        self.device = device
 
 
 class InferenceReply:

smartsim/_core/mli/message_handler.py

@@ -220,22 +220,6 @@ def _assign_reply_channel(
         except Exception as e:
             raise ValueError("Error building reply channel portion of request.") from e
 
-    @staticmethod
-    def _assign_device(
-        request: request_capnp.Request, device: "request_capnp.Device"
-    ) -> None:
-        """
-        Assigns a device to the supplied request.
-
-        :param request: Request being built
-        :param device: Device to be assigned
-        :raises ValueError: if building fails
-        """
-        try:
-            request.device = device
-        except Exception as e:
-            raise ValueError("Error building device portion of request.") from e
-
     @staticmethod
     def _assign_inputs(
         request: request_capnp.Request,
@@ -342,7 +326,6 @@ def _assign_custom_request_attributes(
     def build_request(
         reply_channel: t.ByteString,
         model: t.Union[data_references_capnp.ModelKey, t.ByteString],
-        device: "request_capnp.Device",
         inputs: t.Union[
             t.List[data_references_capnp.TensorKey], t.List[tensor_capnp.Tensor]
         ],
@@ -359,7 +342,6 @@ def build_request(
 
         :param reply_channel: Reply channel to be assigned to request
         :param model: Model to be assigned to request
-        :param device: Device to be assigned to request
         :param inputs: Inputs to be assigned to request
         :param outputs: Outputs to be assigned to request
         :param output_descriptors: Output descriptors to be assigned to request
@@ -368,7 +350,6 @@ def build_request(
         request = request_capnp.Request.new_message()
         MessageHandler._assign_reply_channel(request, reply_channel)
         MessageHandler._assign_model(request, model)
-        MessageHandler._assign_device(request, device)
         MessageHandler._assign_inputs(request, inputs)
         MessageHandler._assign_outputs(request, outputs)
         MessageHandler._assign_output_descriptors(request, output_descriptors)

smartsim/_core/mli/mli_schemas/request/request.capnp

@@ -30,12 +30,6 @@ using Tensors = import "../tensor/tensor.capnp";
 using RequestAttributes = import "request_attributes/request_attributes.capnp";
 using DataRef = import "../data/data_references.capnp";
 
-enum Device {
-  cpu @0;
-  gpu @1;
-  auto @2;
-}
-
 struct ChannelDescriptor {
   reply @0 :Data;
 }
@@ -46,16 +40,15 @@ struct Request {
     modelKey @1 :DataRef.ModelKey;
     modelData @2 :Data;
   }
-  device @3 :Device;
   input :union {
-    inputKeys @4 :List(DataRef.TensorKey);
-    inputData @5 :List(Tensors.Tensor);
+    inputKeys @3 :List(DataRef.TensorKey);
+    inputData @4 :List(Tensors.Tensor);
   }
-  output @6 :List(DataRef.TensorKey);
-  outputDescriptors @7 :List(Tensors.OutputDescriptor);
+  output @5 :List(DataRef.TensorKey);
+  outputDescriptors @6 :List(Tensors.OutputDescriptor);
   customAttributes :union {
-    torch @8 :RequestAttributes.TorchRequestAttributes;
-    tf @9 :RequestAttributes.TensorFlowRequestAttributes;
-    none @10 :Void;
+    torch @7 :RequestAttributes.TorchRequestAttributes;
+    tf @8 :RequestAttributes.TensorFlowRequestAttributes;
+    none @9 :Void;
   }
 }

smartsim/_core/mli/mli_schemas/request/request_capnp.pyi

@@ -33,8 +33,6 @@ from .request_attributes.request_attributes_capnp import (
     TorchRequestAttributesReader,
 )
 
-Device = Literal["cpu", "gpu", "auto"]
-
 class ChannelDescriptor:
     reply: bytes
     @staticmethod
@@ -215,7 +213,6 @@ class Request:
         def write_packed(file: BufferedWriter) -> None: ...
     replyChannel: ChannelDescriptor | ChannelDescriptorBuilder | ChannelDescriptorReader
     model: Request.Model | Request.ModelBuilder | Request.ModelReader
-    device: Device
     input: Request.Input | Request.InputBuilder | Request.InputReader
     output: Sequence[TensorKey | TensorKeyBuilder | TensorKeyReader]
     outputDescriptors: Sequence[

tests/mli/test_integrated_torch_worker.py

@@ -66,7 +66,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     model_bytes = persist_torch_model.read_bytes()
 #     input_tensor = torch.randn(2)
 
-#     expected_device = "cpu"
 #     expected_callback_channel = b"faux_channel_descriptor_bytes"
 #     callback_channel = mli.DragonCommChannel.find(expected_callback_channel)
 
@@ -77,7 +76,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     request = MessageHandler.build_request(
 #         reply_channel=callback_channel.descriptor,
 #         model=model_bytes,
-#         device=expected_device,
 #         inputs=[message_tensor_input],
 #         outputs=[],
 #         custom_attributes=None,
@@ -86,7 +84,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     msg_bytes = MessageHandler.serialize_request(request)
 
 #     inference_request = worker.deserialize(msg_bytes)
-#     assert inference_request.device == expected_device
 #     assert inference_request.callback._descriptor == expected_callback_channel
 
 
@@ -104,7 +101,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     # input_tensor = torch.randn(2)
 #     # feature_store[input_key] = input_tensor
 
-#     expected_device = "cpu"
 #     expected_callback_channel = b"faux_channel_descriptor_bytes"
 #     callback_channel = mli.DragonCommChannel.find(expected_callback_channel)
 
@@ -117,7 +113,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     request = MessageHandler.build_request(
 #         reply_channel=callback_channel.descriptor,
 #         model=message_model_key,
-#         device=expected_device,
 #         inputs=[message_tensor_input_key],
 #         outputs=[message_tensor_output_key],
 #         custom_attributes=None,
@@ -126,7 +121,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     msg_bytes = MessageHandler.serialize_request(request)
 
 #     inference_request = worker.deserialize(msg_bytes)
-#     assert inference_request.device == expected_device
 #     assert inference_request.callback._descriptor == expected_callback_channel
 
 
@@ -147,7 +141,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     # input_tensor = torch.randn(2)
 #     # feature_store[input_key] = input_tensor
 
-#     expected_device = "cpu"
 #     expected_callback_channel = b"faux_channel_descriptor_bytes"
 #     callback_channel = mli.DragonCommChannel.find(expected_callback_channel)
 
@@ -160,7 +153,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     request = MessageHandler.build_request(
 #         reply_channel=callback_channel.descriptor,
 #         model=model_bytes,
-#         device=expected_device,
 #         inputs=[message_tensor_input_key],
 #         # outputs=[message_tensor_output_key],
 #         outputs=[],
@@ -170,7 +162,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     msg_bytes = MessageHandler.serialize_request(request)
 
 #     inference_request = worker.deserialize(msg_bytes)
-#     assert inference_request.device == expected_device
 #     assert inference_request.callback._descriptor == expected_callback_channel
 
 
@@ -191,7 +182,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     input_tensor = torch.randn(2)
 #     # feature_store[input_key] = input_tensor
 
-#     expected_device = "cpu"
 #     expected_callback_channel = b"faux_channel_descriptor_bytes"
 #     callback_channel = mli.DragonCommChannel.find(expected_callback_channel)
 
@@ -207,7 +197,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     request = MessageHandler.build_request(
 #         reply_channel=callback_channel.descriptor,
 #         model=model_bytes,
-#         device=expected_device,
 #         inputs=[message_tensor_input],
 #         # outputs=[message_tensor_output_key],
 #         outputs=[message_tensor_output_key],
@@ -217,7 +206,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     msg_bytes = MessageHandler.serialize_request(request)
 
 #     inference_request = worker.deserialize(msg_bytes)
-#     assert inference_request.device == expected_device
 #     assert inference_request.callback._descriptor == expected_callback_channel
 
 
@@ -238,7 +226,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     input_tensor = torch.randn(2)
 #     # feature_store[input_key] = input_tensor
 
-#     expected_device = "cpu"
 #     expected_callback_channel = b"faux_channel_descriptor_bytes"
 #     callback_channel = mli.DragonCommChannel.find(expected_callback_channel)
 
@@ -254,7 +241,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     request = MessageHandler.build_request(
 #         reply_channel=callback_channel.descriptor,
 #         model=message_model_key,
-#         device=expected_device,
 #         inputs=[message_tensor_input],
 #         # outputs=[message_tensor_output_key],
 #         outputs=[],
@@ -264,7 +250,6 @@ def persist_torch_model(test_dir: str) -> pathlib.Path:
 #     msg_bytes = MessageHandler.serialize_request(request)
 
 #     inference_request = worker.deserialize(msg_bytes)
-#     assert inference_request.device == expected_device
 #     assert inference_request.callback._descriptor == expected_callback_channel
 
 

tests/mli/test_worker_manager.py

@@ -122,7 +122,6 @@ def mock_messages(
         # working set size > 1 has side-effects
         # only incurs cost when working set size has been exceeded
 
-        expected_device: t.Literal["cpu", "gpu"] = "cpu"
         channel_key = comm_channel_root_dir / f"{iteration_number}/channel.txt"
         callback_channel = FileSystemCommChannel(pathlib.Path(channel_key))
 
@@ -144,7 +143,6 @@ def mock_messages(
         request = MessageHandler.build_request(
             reply_channel=callback_channel.descriptor,
             model=message_model_key,
-            device=expected_device,
             inputs=[message_tensor_input_key],
             outputs=[message_tensor_output_key],
             custom_attributes=None,