remove build_reply wrapper

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

@@ -65,22 +65,6 @@ def build_failure_reply(status: "Status", message: str) -> ResponseBuilder:
     )
 
 
-def build_reply(
-    worker: MachineLearningWorkerBase, reply: InferenceReply
-) -> ResponseBuilder:
-    """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:
@@ -377,7 +361,14 @@ def _on_iteration(self) -> None:
         else:
             reply.status_enum = "complete"
             reply.message = "Success"
-            response = build_reply(self._worker, reply)
+
+            results = self._worker.prepare_outputs(reply)
+            return 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

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]