Merge branch 'branch-24.10' into fix-shared-process-pool-fixture-and-…

…atexit

python/morpheus/morpheus/_lib/include/morpheus/pybind11/json.hpp

@@ -37,7 +37,7 @@ struct type_caster<nlohmann::json>
     /**
      * This macro establishes a local variable 'value' of type nlohmann::json
      */
-    PYBIND11_TYPE_CASTER(nlohmann::json, _("object"));
+    PYBIND11_TYPE_CASTER(nlohmann::json, _("dict[str, typing.Any] | list | str | bool | int | float | None"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json
@@ -83,7 +83,7 @@ struct type_caster<nlohmann::json_dict>
     /**
      * This macro establishes a local variable 'value' of type nlohmann::json_dict
      */
-    PYBIND11_TYPE_CASTER(nlohmann::json_dict, _("dict[str, typing.Any]"));
+    PYBIND11_TYPE_CASTER(nlohmann::json_dict, _("dict[str, typing.Any] | None"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json_dict
@@ -92,18 +92,25 @@ struct type_caster<nlohmann::json_dict>
      */
     bool load(handle src, bool convert)
     {
-        if (!src || src.is_none())
+        if (!src)
         {
             return false;
         }
 
-        if (!PyDict_Check(src.ptr()))
+        if (src.is_none())
         {
-            return false;
+            value = nlohmann::json_dict(nullptr);
         }
+        else
+        {
+            if (!PyDict_Check(src.ptr()))
+            {
+                return false;
+            }
 
-        value = static_cast<const nlohmann::json_dict>(
-            mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+            value = static_cast<const nlohmann::json_dict>(
+                mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+        }
 
         return true;
     }
@@ -128,7 +135,7 @@ struct type_caster<nlohmann::json_list>
     /**
      * This macro establishes a local variable 'value' of type nlohmann::json_list
      */
-    PYBIND11_TYPE_CASTER(nlohmann::json_list, _("list[typing.Any]"));
+    PYBIND11_TYPE_CASTER(nlohmann::json_list, _("list"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an nlohmann::json_list
@@ -137,18 +144,25 @@ struct type_caster<nlohmann::json_list>
      */
     bool load(handle src, bool convert)
     {
-        if (!src || src.is_none())
+        if (!src)
         {
             return false;
         }
 
-        if (!PyList_Check(src.ptr()))
+        if (src.is_none())
         {
-            return false;
+            value = nlohmann::json_list(nullptr);
         }
+        else
+        {
+            if (!PyList_Check(src.ptr()))
+            {
+                return false;
+            }
 
-        value = static_cast<const nlohmann::json_list>(
-            mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+            value = static_cast<const nlohmann::json_list>(
+                mrc::pymrc::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+        }
 
         return true;
     }
@@ -173,7 +187,7 @@ struct type_caster<morpheus::utilities::json_t>
     /**
      * This macro establishes a local variable 'value' of type morpheus::utilities::json_t
      */
-    PYBIND11_TYPE_CASTER(morpheus::utilities::json_t, _("object"));
+    PYBIND11_TYPE_CASTER(morpheus::utilities::json_t, _("object | None"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an morpheus::utilities::json_t
@@ -219,7 +233,7 @@ struct type_caster<morpheus::utilities::json_dict_t>
     /**
      * This macro establishes a local variable 'value' of type morpheus::utilities::json_t_dict
      */
-    PYBIND11_TYPE_CASTER(morpheus::utilities::json_dict_t, _("dict[str, typing.Any]"));
+    PYBIND11_TYPE_CASTER(morpheus::utilities::json_dict_t, _("dict[str, typing.Any] | None"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an morpheus::utilities::json_t_dict
@@ -228,18 +242,25 @@ struct type_caster<morpheus::utilities::json_dict_t>
      */
     bool load(handle src, bool convert)
     {
-        if (!src || src.is_none())
+        if (!src)
         {
             return false;
         }
 
-        if (!PyDict_Check(src.ptr()))
+        if (src.is_none())
         {
-            return false;
+            value = morpheus::utilities::json_dict_t(nullptr);
         }
+        else
+        {
+            if (!PyDict_Check(src.ptr()))
+            {
+                return false;
+            }
 
-        value = static_cast<const morpheus::utilities::json_dict_t>(
-            morpheus::utilities::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+            value = static_cast<const morpheus::utilities::json_dict_t>(
+                morpheus::utilities::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+        }
 
         return true;
     }
@@ -264,7 +285,7 @@ struct type_caster<morpheus::utilities::json_list_t>
     /**
      * This macro establishes a local variable 'value' of type morpheus::utilities::json_t_list
      */
-    PYBIND11_TYPE_CASTER(morpheus::utilities::json_list_t, _("list[typing.Any]"));
+    PYBIND11_TYPE_CASTER(morpheus::utilities::json_list_t, _("list | None"));
 
     /**
      * Conversion part 1 (Python->C++): convert a PyObject into an morpheus::utilities::json_t_list
@@ -273,18 +294,25 @@ struct type_caster<morpheus::utilities::json_list_t>
      */
     bool load(handle src, bool convert)
     {
-        if (!src || src.is_none())
+        if (!src)
         {
             return false;
         }
 
-        if (!PyList_Check(src.ptr()))
+        if (src.is_none())
         {
-            return false;
+            value = morpheus::utilities::json_list_t(nullptr);
         }
+        else
+        {
+            if (!PyList_Check(src.ptr()))
+            {
+                return false;
+            }
 
-        value = static_cast<const morpheus::utilities::json_list_t>(
-            morpheus::utilities::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+            value = static_cast<const morpheus::utilities::json_list_t>(
+                morpheus::utilities::cast_from_pyobject(pybind11::reinterpret_borrow<pybind11::object>(src)));
+        }
 
         return true;
     }

python/morpheus/morpheus/_lib/messages/__init__.pyi

@@ -37,18 +37,18 @@ class ControlMessage():
     def __init__(self, arg0: ControlMessage) -> None: ...
     @typing.overload
     def __init__(self, arg0: dict) -> None: ...
-    def add_task(self, task_type: str, task: object) -> None: ...
+    def add_task(self, task_type: str, task: object | None) -> None: ...
     @typing.overload
-    def config(self) -> object: ...
+    def config(self) -> object | None: ...
     @typing.overload
-    def config(self, config: object) -> None: ...
+    def config(self, config: object | None) -> None: ...
     def copy(self) -> ControlMessage: ...
     def filter_timestamp(self, regex_filter: str) -> dict: 
         """
         Retrieve timestamps matching a regex filter within a given group.
         """
     def get_metadata(self, key: object = None, default_value: object = None) -> object: ...
-    def get_tasks(self) -> object: ...
+    def get_tasks(self) -> object | None: ...
     def get_timestamp(self, key: str, fail_if_nonexist: bool = False) -> object: 
         """
         Retrieve the timestamp for a given group and key. Returns None if the timestamp does not exist and fail_if_nonexist is False.
@@ -62,8 +62,8 @@ class ControlMessage():
     def payload(self, arg0: MessageMeta) -> None: ...
     @typing.overload
     def payload(self, meta: object) -> None: ...
-    def remove_task(self, task_type: str) -> object: ...
-    def set_metadata(self, key: str, value: object) -> None: ...
+    def remove_task(self, task_type: str) -> object | None: ...
+    def set_metadata(self, key: str, value: object | None) -> None: ...
     def set_timestamp(self, key: str, timestamp: object) -> None: 
         """
         Set a timestamp for a given key and group.

python/morpheus_llm/morpheus_llm/_lib/include/morpheus_llm/llm/llm_context.hpp

@@ -22,9 +22,9 @@
 #include "morpheus_llm/llm/llm_task.hpp"
 
 #include "morpheus/export.h"
+#include "morpheus/utilities/json_types.hpp"
 
 #include <mrc/types.hpp>
-#include <pymrc/utilities/json_values.hpp>
 
 #include <memory>
 #include <string>
@@ -111,9 +111,9 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
     /**
      * @brief Get all outputs for this context.
      *
-     * @return const mrc::pymrc::JSONValues&
+     * @return const utilities::json_t&
      */
-    const mrc::pymrc::JSONValues& all_outputs() const;
+    const utilities::json_t& all_outputs() const;
 
     /**
      * @brief Get full name of context containing parents up to root.
@@ -141,39 +141,39 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
     /**
      * @brief Get the input value from parent context corresponding to first internal input of this context.
      *
-     * @return mrc::pymrc::JSONValues
+     * @return utilities::json_t
      */
-    mrc::pymrc::JSONValues get_input() const;
+    const utilities::json_t& get_input() const;
 
     /**
      * @brief Get the parent output value corresponding to given internal input name.
      *
      * @param node_name internal input name
-     * @return mrc::pymrc::JSONValues
+     * @return utilities::json_t
      */
-    mrc::pymrc::JSONValues get_input(const std::string& node_name) const;
+    const utilities::json_t& get_input(const std::string& node_name) const;
 
     /**
      * @brief Get parent output values corresponding to all internal input names.
      *
-     * @return mrc::pymrc::JSONValues
+     * @return utilities::json_t
      */
-    mrc::pymrc::JSONValues get_inputs() const;
+    utilities::json_t get_inputs() const;
 
     /**
      * @brief Set output mappings for this context.
      *
      * @param outputs output mappings
      */
-    void set_output(mrc::pymrc::JSONValues&& outputs);
+    void set_output(utilities::json_t outputs);
 
     /**
      * @brief Set an output value for this context.
      *
      * @param output_name output name
      * @param output output value
      */
-    void set_output(const std::string& output_name, mrc::pymrc::JSONValues&& output);
+    void set_output(const std::string& output_name, utilities::json_t output);
 
     /**
      * @brief Set the output names to propagate from this context when using pop.
@@ -187,9 +187,9 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
     /**
      * @brief Get all outputs for this context.
      *
-     * @return const mrc::pymrc::JSONValues&
+     * @return const utilities::json_t&
      */
-    const mrc::pymrc::JSONValues& view_outputs() const;
+    const utilities::json_t& view_outputs() const;
 
   private:
     input_mappings_t::const_iterator find_input(const std::string& node_name, bool throw_if_not_found = true) const;
@@ -201,7 +201,7 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
 
     std::shared_ptr<LLMContextState> m_state;
 
-    mrc::pymrc::JSONValues m_outputs;
+    utilities::json_t m_outputs;
 
     mrc::Promise<void> m_outputs_promise;
     mrc::SharedFuture<void> m_outputs_future;

python/morpheus_llm/morpheus_llm/_lib/include/morpheus_llm/llm/llm_lambda_node.hpp

@@ -79,7 +79,7 @@ class LLMLambdaNode : public LLMNodeBase
         {
             const auto& arg = context->get_input();
 
-            auto output = co_await this->m_function(arg.view_json().get<std::tuple_element_t<0, args_tuple_t>>());
+            auto output = co_await this->m_function(arg.get<std::tuple_element_t<0, args_tuple_t>>());
 
             nlohmann::json outputs_json = std::move(output);
 
@@ -92,7 +92,7 @@ class LLMLambdaNode : public LLMNodeBase
         {
             auto args = context->get_inputs();
 
-            auto outputs = co_await this->m_function(args.view_json());
+            auto outputs = co_await this->m_function(args);
 
             nlohmann::json outputs_json = std::move(outputs);
 

python/morpheus_llm/morpheus_llm/_lib/llm/__init__.pyi

@@ -64,16 +64,16 @@ class LLMContext():
     @typing.overload
     def __init__(self, task: LLMTask, message: morpheus._lib.messages.ControlMessage) -> None: ...
     @typing.overload
-    def get_input(self) -> object: ...
+    def get_input(self) -> object | None: ...
     @typing.overload
-    def get_input(self, node_name: str) -> object: ...
-    def get_inputs(self) -> object: ...
+    def get_input(self, node_name: str) -> object | None: ...
+    def get_inputs(self) -> object | None: ...
     def message(self) -> morpheus._lib.messages.ControlMessage: ...
     def push(self, name: str, inputs: typing.List[InputMap]) -> LLMContext: ...
     @typing.overload
-    def set_output(self, output_name: str, output: object) -> None: ...
+    def set_output(self, output_name: str, output: object | None) -> None: ...
     @typing.overload
-    def set_output(self, outputs: object) -> None: ...
+    def set_output(self, outputs: object | None) -> None: ...
     def task(self) -> LLMTask: ...
     @property
     def full_name(self) -> str:
@@ -96,9 +96,9 @@ class LLMContext():
         :type: LLMContext
         """
     @property
-    def view_outputs(self) -> object:
+    def view_outputs(self) -> object | None:
         """
-        :type: object
+        :type: object | None
         """
     pass
 class LLMNodeBase():

python/morpheus_llm/morpheus_llm/_lib/llm/module.cpp

@@ -32,18 +32,20 @@
 #include "py_llm_lambda_node.hpp"
 
 #include "morpheus/messages/control.hpp"  // IWYU pragma: keep
+#include "morpheus/pybind11/json.hpp"     // IWYU pragma: keep
 #include "morpheus/utilities/cudf_util.hpp"
+#include "morpheus/utilities/json_types.hpp"
 #include "morpheus/version.hpp"
 
 #include <mrc/segment/object.hpp>  // for Object, ObjectProperties
 #include <mrc/utils/string_utils.hpp>
 #include <nlohmann/detail/exceptions.hpp>  // for nlohmann::detail::out_of_range
+#include <nlohmann/json.hpp>               // for basic_json
 #include <pybind11/functional.h>           // IWYU pragma: keep
 #include <pybind11/pybind11.h>  // for arg, init, class_, module_, str_attr_accessor, PYBIND11_MODULE, pybind11
 #include <pybind11/stl.h>       // IWYU pragma: keep
 #include <pymrc/coro.hpp>       // IWYU pragma: keep
-#include <pymrc/utilities/json_values.hpp>  // for JSONValues
-#include <pymrc/utils.hpp>                  // for pymrc::import
+#include <pymrc/utils.hpp>      // for pymrc::import
 
 #include <memory>
 #include <sstream>
@@ -198,9 +200,9 @@ PYBIND11_MODULE(llm, _module)
              py::overload_cast<const std::string&>(&LLMContext::get_input, py::const_),
              py::arg("node_name"))
         .def("get_inputs", &LLMContext::get_inputs)
-        .def("set_output", py::overload_cast<mrc::pymrc::JSONValues&&>(&LLMContext::set_output), py::arg("outputs"))
+        .def("set_output", py::overload_cast<utilities::json_t>(&LLMContext::set_output), py::arg("outputs"))
         .def("set_output",
-             py::overload_cast<const std::string&, mrc::pymrc::JSONValues&&>(&LLMContext::set_output),
+             py::overload_cast<const std::string&, utilities::json_t>(&LLMContext::set_output),
              py::arg("output_name"),
              py::arg("output"))
         .def("push", &LLMContext::push, py::arg("name"), py::arg("inputs"));