post-merge tweaks

Author: ankona

ex/high_throughput_inference/mock_app.py

@@ -115,7 +115,7 @@ def run_model(self, model: bytes | str, batch: torch.Tensor):
         self.measure_time("build_tensor_descriptor")
         built_model = None
         if isinstance(model, str):
-            model_arg = MessageHandler.build_model_key(model)
+            model_arg = MessageHandler.build_model_key(model)  # todo: this needs FSD
         else:
             model_arg = MessageHandler.build_model(model, "resnet-50", "1.0")
         request = MessageHandler.build_request(

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

@@ -34,7 +34,6 @@
 from ...comm.channel.channel import CommChannelBase
 from ...comm.channel.dragonchannel import DragonCommChannel
 from ...infrastructure.environmentloader import EnvironmentConfigLoader
-from ...infrastructure.storage.featurestore import FeatureStore
 from ...infrastructure.worker.worker import (
     InferenceReply,
     InferenceRequest,
@@ -54,32 +53,18 @@
 logger = get_logger(__name__)
 
 
-def build_failure_reply(status: "StatusEnum", message: str) -> Response:
+def build_failure_reply(status: "Status", message: str) -> ResponseBuilder:
     """Build a response indicating a failure occurred
     :param status: The status of the response
     :param message: The error message to include in the response"""
     return MessageHandler.build_response(
-        status=status,  # todo: need to indicate correct status
-        message=message,  # todo: decide what these will be
+        status=status,
+        message=message,
         result=None,
         custom_attributes=None,
     )
 
 
-def build_reply(worker: MachineLearningWorkerBase, reply: InferenceReply) -> Response:
-    """Builds a response for a successful inference request
-    :param worker: A worker to process the reply with
-    :param reply: The internal representation of the reply"""
-    results = worker.prepare_outputs(reply)
-
-    return MessageHandler.build_response(
-        status=reply.status_enum,
-        message=reply.message,
-        result=results,
-        custom_attributes=None,
-    )
-
-
 def exception_handler(
     exc: Exception, reply_channel: t.Optional[CommChannelBase], failure_message: str
 ) -> None:
@@ -143,13 +128,15 @@ def _check_feature_stores(self, request: InferenceRequest) -> bool:
         """Ensures that all feature stores required by the request are available
         :param request: The request to validate"""
         # collect all feature stores required by the request
-        fs_model = {request.model_key.descriptor}
+        fs_model: t.Set[str] = set()
+        if request.model_key:
+            fs_model = {request.model_key.descriptor}
         fs_inputs = {key.descriptor for key in request.input_keys}
         fs_outputs = {key.descriptor for key in request.output_keys}
 
         # identify which feature stores are requested and unknown
-        fs_desired = fs_model + fs_inputs + fs_outputs
-        fs_actual = {key for key in self._feature_stores}
+        fs_desired = fs_model.union(fs_inputs).union(fs_outputs)
+        fs_actual = {item.descriptor for item in self._feature_stores.values()}
         fs_missing = fs_desired - fs_actual
 
         # exit if all desired feature stores are not available
@@ -259,7 +246,7 @@ def _on_iteration(self) -> None:
                 interm = time.perf_counter()  # timing
                 try:
                     fetch_model_result = self._worker.fetch_model(
-                        request, self._feature_store
+                        request, self._feature_stores
                     )
                 except Exception as e:
                     exception_handler(
@@ -287,7 +274,7 @@ def _on_iteration(self) -> None:
             interm = time.perf_counter()  # timing
             try:
                 fetch_model_result = self._worker.fetch_model(
-                    request, self._feature_store
+                    request, self._feature_stores
                 )
             except Exception as e:
                 exception_handler(
@@ -310,7 +297,9 @@ def _on_iteration(self) -> None:
         timings.append(time.perf_counter() - interm)  # timing
         interm = time.perf_counter()  # timing
         try:
-            fetch_input_result = self._worker.fetch_inputs(request, self._feature_store)
+            fetch_input_result = self._worker.fetch_inputs(
+                request, self._feature_stores
+            )
         except Exception as e:
             exception_handler(e, request.callback, "Failed while fetching the inputs.")
             return
@@ -370,10 +359,16 @@ def _on_iteration(self) -> None:
         if reply.outputs is None or not reply.outputs:
             response = build_failure_reply("fail", "Outputs not found.")
         else:
-            if reply.outputs is None or not reply.outputs:
-                response = build_failure_reply("fail", "no-results")
-
-            response = build_reply(self._worker, reply)
+            reply.status_enum = "complete"
+            reply.message = "Success"
+
+            results = self._worker.prepare_outputs(reply)
+            response = MessageHandler.build_response(
+                status=reply.status_enum,
+                message=reply.message,
+                result=results,
+                custom_attributes=None,
+            )
 
         timings.append(time.perf_counter() - interm)  # timing
         interm = time.perf_counter()  # timing

smartsim/_core/mli/infrastructure/environmentloader.py

@@ -46,13 +46,10 @@ class EnvironmentConfigLoader:
     """
 
     def __init__(self) -> None:
-        self._feature_store_descriptor: t.Optional[str] = os.getenv(
-            "SSFeatureStore", None
-        )
         self._queue_descriptor: t.Optional[str] = os.getenv("SSQueue", None)
-        self.feature_store: t.Optional[FeatureStore] = None
-        self.feature_stores: t.Optional[t.Dict[FeatureStore]] = None
+        self.feature_stores: t.Optional[t.Dict[str, FeatureStore]] = None
         self.queue: t.Optional[DragonFLIChannel] = None
+        self._prefix = "SSFeatureStore"
 
     def _load_feature_store(self, env_var: str) -> FeatureStore:
         """Load a feature store from a descriptor
@@ -62,10 +59,12 @@ def _load_feature_store(self, env_var: str) -> FeatureStore:
 
         value = os.getenv(env_var)
         if not value:
-            raise SmartSimError(f"Empty feature store descriptor in environment: {env_var}")
+            raise SmartSimError(
+                f"Empty feature store descriptor in environment: {env_var}"
+            )
 
         try:
-            return pickle.loads(base64.b64decode(value))
+            return t.cast(FeatureStore, pickle.loads(base64.b64decode(value)))
         except:
             raise SmartSimError(
                 f"Invalid feature store descriptor in environment: {env_var}"
@@ -74,9 +73,8 @@ def _load_feature_store(self, env_var: str) -> FeatureStore:
     def get_feature_stores(self) -> t.Dict[str, FeatureStore]:
         """Loads multiple Feature Stores by scanning environment for variables
         prefixed with `SSFeatureStore`"""
-        prefix = "SSFeatureStore"
-        if self.feature_stores is None:
-            env_vars = [var for var in os.environ if var.startswith(prefix)]
+        if not self.feature_stores:
+            env_vars = [var for var in os.environ if var.startswith(self._prefix)]
             stores = [self._load_feature_store(var) for var in env_vars]
             self.feature_stores = {fs.descriptor: fs for fs in stores}
         return self.feature_stores

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

@@ -27,8 +27,6 @@
 import typing as t
 from abc import ABC, abstractmethod
 
-import numpy as np
-
 from .....error import SmartSimError
 from .....log import get_logger
 from ...comm.channel.channel import CommChannelBase
@@ -38,6 +36,7 @@
 
 if t.TYPE_CHECKING:
     from smartsim._core.mli.mli_schemas.response.response_capnp import Status
+    from smartsim._core.mli.mli_schemas.tensor.tensor_capnp import TensorDescriptor
 
 logger = get_logger(__name__)
 
@@ -81,13 +80,13 @@ class InferenceReply:
     def __init__(
         self,
         outputs: t.Optional[t.Collection[t.Any]] = None,
-        output_keys: t.Optional[t.Collection[str]] = None,
+        output_keys: t.Optional[t.Collection[FeatureStoreKey]] = None,
         status_enum: "Status" = "running",
         message: str = "In progress",
     ) -> None:
         """Initialize the object"""
         self.outputs: t.Collection[t.Any] = outputs or []
-        self.output_keys: t.Collection[t.Optional[str]] = output_keys or []
+        self.output_keys: t.Collection[t.Optional[FeatureStoreKey]] = output_keys or []
         self.status_enum = status_enum
         self.message = message
 
@@ -175,27 +174,25 @@ def deserialize_message(
         elif request.model.which() == "data":
             model_bytes = request.model.data
 
-        callback_key = request.replyChannel.reply
+        callback_key = request.replyChannel.descriptor
         comm_channel = channel_type(callback_key)
-
         input_keys: t.Optional[t.List[FeatureStoreKey]] = None
         input_bytes: t.Optional[t.List[bytes]] = None
-        input_meta: t.List[t.Any] = []
+        output_keys: t.Optional[t.List[FeatureStoreKey]] = None
+        input_meta: t.Optional[t.List[TensorDescriptor]] = None
 
         if request.input.which() == "keys":
             input_keys = [
-                FeatureStoreKey(input_key.key, input_key.featureStoreDescriptor)
-                for input_key in request.input.keys
+                FeatureStoreKey(value.key, value.featureStoreDescriptor)
+                for value 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]
+        elif request.input.which() == "descriptors":
+            input_meta = request.input.descriptors  # type: ignore
 
-        output_keys: t.List[FeatureStoreKey] = []
         if request.output:
             output_keys = [
-                FeatureStoreKey(output_key.key, output_key.featureStoreDescriptor)
-                for output_key in request.output
+                FeatureStoreKey(value.key, value.featureStoreDescriptor)
+                for value in request.output
             ]
 
         inference_request = InferenceRequest(
@@ -214,27 +211,19 @@ def deserialize_message(
     def prepare_outputs(reply: InferenceReply) -> t.List[t.Any]:
         prepared_outputs: t.List[t.Any] = []
         if reply.output_keys:
-            for fs_key in reply.output_keys:
-                if not fs_key:
+            for value in reply.output_keys:
+                if not value:
                     continue
-
-                msg_key = MessageHandler.build_tensor_key(fs_key.key, fs_key.descriptor)
+                msg_key = MessageHandler.build_tensor_key(value.key, value.descriptor)
                 prepared_outputs.append(msg_key)
         elif reply.outputs:
-            arrays: t.List[np.ndarray[t.Any, np.dtype[t.Any]]] = [
-                output.numpy() for output in reply.outputs
-            ]
-            for tensor in arrays:
-                # todo: need to have the output attributes specified in the req?
-                # maybe, add `MessageHandler.dtype_of(tensor)`?
-                # can `build_tensor` do dtype and shape?
-                msg_tensor = MessageHandler.build_tensor(
-                    tensor,
+            for _ in reply.outputs:
+                msg_tensor_desc = MessageHandler.build_tensor_descriptor(
                     "c",
                     "float32",
                     [1],
                 )
-                prepared_outputs.append(msg_tensor)
+                prepared_outputs.append(msg_tensor_desc)
         return prepared_outputs
 
     @staticmethod

smartsim/_core/mli/message_handler.py

@@ -439,6 +439,7 @@ def _assign_result(
         result: t.Union[
             t.List[tensor_capnp.TensorDescriptor],
             t.List[data_references_capnp.TensorKey],
+            None,
         ],
     ) -> None:
         """
@@ -504,7 +505,7 @@ def build_response(
         result: t.Union[
             t.List[tensor_capnp.TensorDescriptor],
             t.List[data_references_capnp.TensorKey],
-            None
+            None,
         ],
         custom_attributes: t.Union[
             response_attributes_capnp.TorchResponseAttributes,

tests/dragon/test_reply_building.py

@@ -30,10 +30,7 @@
 
 dragon = pytest.importorskip("dragon")
 
-from smartsim._core.mli.infrastructure.control.workermanager import (
-    build_failure_reply,
-    build_reply,
-)
+from smartsim._core.mli.infrastructure.control.workermanager import build_failure_reply
 from smartsim._core.mli.infrastructure.worker.worker import InferenceReply
 
 if t.TYPE_CHECKING:
@@ -64,28 +61,3 @@ def test_build_failure_reply_fails():
 
     assert "Error assigning status to response" in ex.value.args[0]
 
-
-@pytest.mark.parametrize(
-    "status, message",
-    [
-        pytest.param("complete", "Success", id="complete"),
-    ],
-)
-def test_build_reply(status: "Status", message: str):
-    "Ensures replies can be built successfully"
-    reply = InferenceReply()
-    reply.status_enum = status
-    reply.message = message
-    response = build_reply(reply)
-    assert response.status == status
-    assert response.message == message
-
-
-def test_build_reply_fails():
-    "Ensures ValueError is raised if a Status Enum is not used"
-    with pytest.raises(ValueError) as ex:
-        reply = InferenceReply()
-        reply.status_enum = "not a status enum"
-        response = build_reply(reply)
-
-    assert "Error assigning status to response" in ex.value.args[0]

tests/mli/test_worker_manager.py

@@ -32,6 +32,8 @@
 
 import pytest
 
+from tests.mli.featurestore import FileSystemFeatureStore
+
 torch = pytest.importorskip("torch")
 dragon = pytest.importorskip("dragon")
 
@@ -183,14 +185,11 @@ def test_worker_manager(prepare_environment: pathlib.Path) -> None:
     )
 
     # create a mock client application to populate the request queue
-    feature_stores = config_loader.get_feature_stores()
-    fs_list = list(feature_stores.values())
-
     msg_pump = mp.Process(
         target=mock_messages,
         args=(
             config_loader.get_queue(),
-            fs_list[0],
+            FileSystemFeatureStore(fs_path),
             fs_path,
             comm_path,
         ),