Add model metadata to request schema (#624)

Add `Model` schema with model metadata.

[ committed by @AlyssaCote ]
[ approved by @ankona ]

doc/changelog.md

@@ -13,10 +13,11 @@ Jump to:
 
 Description
 
+- Add Model schema with model metadata included
+- Removed device from schemas, MessageHandler and tests
 - 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

@@ -43,6 +43,7 @@
 from smartsim.log import get_logger
 
 if t.TYPE_CHECKING:
+    from smartsim._core.mli.mli_schemas.model.model_capnp import Model
     from smartsim._core.mli.mli_schemas.response.response_capnp import StatusEnum
 
 logger = get_logger(__name__)
@@ -65,12 +66,12 @@ def deserialize_message(
     request = MessageHandler.deserialize_request(data_blob)
     # return request
     model_key: t.Optional[str] = None
-    model_bytes: t.Optional[bytes] = None
+    model_bytes: t.Optional[Model] = None
 
-    if request.model.which() == "modelKey":
-        model_key = request.model.modelKey.key
-    elif request.model.which() == "modelData":
-        model_bytes = request.model.modelData
+    if request.model.which() == "key":
+        model_key = request.model.key.key
+    elif request.model.which() == "data":
+        model_bytes = request.model.data
 
     callback_key = request.replyChannel.reply
 
@@ -91,11 +92,11 @@ def deserialize_message(
     # # end client
     input_meta: t.List[t.Any] = []
 
-    if request.input.which() == "inputKeys":
-        input_keys = [input_key.key for input_key in request.input.inputKeys]
-    elif request.input.which() == "inputData":
-        input_bytes = [data.blob for data in request.input.inputData]
-        input_meta = [data.tensorDescriptor for data in request.input.inputData]
+    if request.input.which() == "keys":
+        input_keys = [input_key.key for input_key in request.input.keys]
+    elif request.input.which() == "data":
+        input_bytes = [data.blob for data in request.input.data]
+        input_meta = [data.tensorDescriptor for data in request.input.data]
 
     inference_request = InferenceRequest(
         model_key=model_key,

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

@@ -30,6 +30,7 @@
 import smartsim.error as sse
 from smartsim._core.mli.comm.channel.channel import CommChannelBase
 from smartsim._core.mli.infrastructure.storage.featurestore import FeatureStore
+from smartsim._core.mli.mli_schemas.model.model_capnp import Model
 from smartsim.log import get_logger
 
 logger = get_logger(__name__)
@@ -48,7 +49,7 @@ def __init__(
         input_keys: t.Optional[t.List[str]] = None,
         input_meta: t.Optional[t.List[t.Any]] = None,
         output_keys: t.Optional[t.List[str]] = None,
-        raw_model: t.Optional[bytes] = None,
+        raw_model: t.Optional[Model] = None,
         batch_size: int = 0,
     ):
         """Initialize the object"""
@@ -158,7 +159,7 @@ def fetch_model(
             # model_key = hash(request.raw_model)
             # feature_store[model_key] = request.raw_model
             # short-circuit and return the directly supplied model
-            return FetchModelResult(request.raw_model)
+            return FetchModelResult(request.raw_model.data)
 
         if not request.model_key:
             raise sse.SmartSimError(

smartsim/_core/mli/message_handler.py

@@ -28,6 +28,7 @@
 import numpy as np
 
 from .mli_schemas.data import data_references_capnp
+from .mli_schemas.model import model_capnp
 from .mli_schemas.request import request_capnp
 from .mli_schemas.request.request_attributes import request_attributes_capnp
 from .mli_schemas.response import response_capnp
@@ -112,6 +113,25 @@ def build_tensor_key(key: str) -> data_references_capnp.TensorKey:
             raise ValueError("Error building tensor key.") from e
         return tensor_key
 
+    @staticmethod
+    def build_model(data: bytes, name: str, version: str) -> model_capnp.Model:
+        """
+        Builds a new Model message with the provided data, name, and version.
+
+        :param data: Model data
+        :param name: Model name
+        :param version: Model version
+        :raises ValueError: if building fails
+        """
+        try:
+            model = model_capnp.Model.new_message()
+            model.data = data
+            model.name = name
+            model.version = version
+        except Exception as e:
+            raise ValueError("Error building model.") from e
+        return model
+
     @staticmethod
     def build_model_key(key: str) -> data_references_capnp.ModelKey:
         """
@@ -187,7 +207,7 @@ def build_tf_response_attributes() -> (
     @staticmethod
     def _assign_model(
         request: request_capnp.Request,
-        model: t.Union[data_references_capnp.ModelKey, t.ByteString],
+        model: t.Union[data_references_capnp.ModelKey, model_capnp.Model],
     ) -> None:
         """
         Assigns a model to the supplied request.
@@ -197,16 +217,20 @@ def _assign_model(
         :raises ValueError: if building fails
         """
         try:
-            if isinstance(model, bytes):
-                request.model.modelData = model
+            class_name = model.schema.node.displayName.split(":")[-1]  # type: ignore
+            if class_name == "Model":
+                request.model.data = model  # type: ignore
+            elif class_name == "ModelKey":
+                request.model.key = model  # type: ignore
             else:
-                request.model.modelKey = model  # type: ignore
+                raise ValueError("""Invalid custom attribute class name.
+                        Expected 'Model' or 'ModelKey'.""")
         except Exception as e:
             raise ValueError("Error building model portion of request.") from e
 
     @staticmethod
     def _assign_reply_channel(
-        request: request_capnp.Request, reply_channel: t.ByteString
+        request: request_capnp.Request, reply_channel: bytes
     ) -> None:
         """
         Assigns a reply channel to the supplied request.
@@ -239,9 +263,9 @@ def _assign_inputs(
                 display_name = inputs[0].schema.node.displayName  # type: ignore
                 input_class_name = display_name.split(":")[-1]
                 if input_class_name == "Tensor":
-                    request.input.inputData = inputs  # type: ignore
+                    request.input.data = inputs  # type: ignore
                 elif input_class_name == "TensorKey":
-                    request.input.inputKeys = inputs  # type: ignore
+                    request.input.keys = inputs  # type: ignore
                 else:
                     raise ValueError(
                         "Invalid input class name. Expected 'Tensor' or 'TensorKey'."
@@ -324,8 +348,8 @@ def _assign_custom_request_attributes(
 
     @staticmethod
     def build_request(
-        reply_channel: t.ByteString,
-        model: t.Union[data_references_capnp.ModelKey, t.ByteString],
+        reply_channel: bytes,
+        model: t.Union[data_references_capnp.ModelKey, model_capnp.Model],
         inputs: t.Union[
             t.List[data_references_capnp.TensorKey], t.List[tensor_capnp.Tensor]
         ],
@@ -357,7 +381,7 @@ def build_request(
         return request
 
     @staticmethod
-    def serialize_request(request: request_capnp.RequestBuilder) -> t.ByteString:
+    def serialize_request(request: request_capnp.RequestBuilder) -> bytes:
         """
         Serializes a built request message.
 
@@ -366,7 +390,7 @@ def serialize_request(request: request_capnp.RequestBuilder) -> t.ByteString:
         return request.to_bytes()
 
     @staticmethod
-    def deserialize_request(request_bytes: t.ByteString) -> request_capnp.Request:
+    def deserialize_request(request_bytes: bytes) -> request_capnp.Request:
         """
         Deserializes a serialized request message.
 
@@ -499,14 +523,14 @@ def build_response(
         return response
 
     @staticmethod
-    def serialize_response(response: response_capnp.ResponseBuilder) -> t.ByteString:
+    def serialize_response(response: response_capnp.ResponseBuilder) -> bytes:
         """
         Serializes a built response message.
         """
         return response.to_bytes()
 
     @staticmethod
-    def deserialize_response(response_bytes: t.ByteString) -> response_capnp.Response:
+    def deserialize_response(response_bytes: bytes) -> response_capnp.Response:
         """
         Deserializes a serialized response message.
         """

smartsim/_core/mli/mli_schemas/data/data_references.capnp

@@ -32,4 +32,4 @@ struct ModelKey {
 
 struct TensorKey {
   key @0 :Text;
-}
+}

smartsim/_core/mli/mli_schemas/data/data_references_capnp.py

@@ -1,3 +1,29 @@
+# BSD 2-Clause License
+
+# Copyright (c) 2021-2024, Hewlett Packard Enterprise
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
 """This is an automatically generated stub for `data_references.capnp`."""
 
 import os

smartsim/_core/mli/mli_schemas/data/data_references_capnp.pyi

@@ -1,3 +1,29 @@
+# BSD 2-Clause License
+
+# Copyright (c) 2021-2024, Hewlett Packard Enterprise
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
 """This is an automatically generated stub for `data_references.capnp`."""
 
 # mypy: ignore-errors

smartsim/_core/mli/mli_schemas/model/model.capnp

@@ -0,0 +1,33 @@
+# BSD 2-Clause License
+
+# Copyright (c) 2021-2024, Hewlett Packard Enterprise
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+@0xaefb9301e14ba4bd;
+
+struct Model {
+  data @0 :Data;
+  name @1 :Text;
+  version @2 :Text;
+}

smartsim/_core/mli/mli_schemas/model/model_capnp.py

@@ -0,0 +1,38 @@
+# BSD 2-Clause License
+
+# Copyright (c) 2021-2024, Hewlett Packard Enterprise
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+"""This is an automatically generated stub for `model.capnp`."""
+
+import os
+
+import capnp  # type: ignore
+
+capnp.remove_import_hook()
+here = os.path.dirname(os.path.abspath(__file__))
+module_file = os.path.abspath(os.path.join(here, "model.capnp"))
+Model = capnp.load(module_file).Model
+ModelBuilder = Model
+ModelReader = Model