Add optional filtering to LLM Engines

ci/iwyu/mappings.imp

@@ -49,6 +49,7 @@
 # pybind11
 { "include": [ "<pybind11/detail/common.h>", "private", "<pybind11/pytypes.h>", "public" ] },
 { "include": [ "<pybind11/cast.h>", "private", "<pybind11/pybind11.h>", "public" ] },
+{ "include": [ "<pybind11/attr.h>", "private", "<pybind11/pybind11.h>", "public" ] },
 
 # rxcpp
 # Hide includes that are exported by <rxcpp/rx.hpp>

morpheus/_lib/cudf_helpers.pyx

@@ -14,12 +14,15 @@
 # limitations under the License.
 
 import cudf
+from cudf.core.dtypes import StructDtype
 
 from libcpp.string cimport string
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
 from cudf._lib.column cimport Column
+from cudf._lib.cpp.io.types cimport column_name_info
+from cudf._lib.cpp.io.types cimport table_metadata
 from cudf._lib.cpp.io.types cimport table_with_metadata
 from cudf._lib.cpp.table.table_view cimport table_view
 from cudf._lib.cpp.types cimport size_type
@@ -56,6 +59,18 @@ cdef public api:
 
  object make_table_from_table_info_data(TableInfoData table_info, object owner):
 
+ cdef table_metadata tbl_meta
+
+ num_index_cols_meta = 0
+ cdef column_name_info child_info
+ for i, name in enumerate(owner._column_names, num_index_cols_meta):
+ child_info.name = name.encode()
+ tbl_meta.schema_info.push_back(child_info)
+ _set_col_children_metadata(
+ owner[name]._column,
+ tbl_meta.schema_info[i]
+ )
+
  index_names = None
 
  if (table_info.index_names.size() > 0):
@@ -89,7 +104,11 @@ cdef public api:
  import traceback
  print("error while converting libcudf table to cudf dataframe:", traceback.format_exc())
 
- return cudf.DataFrame._from_data(data, index)
+ df = cudf.DataFrame._from_data(data, index)
+
+ update_struct_field_names(df, tbl_meta.schema_info)
+
+ return df
 
 
  TableInfoData make_table_info_data_from_table(object table):
@@ -173,3 +192,59 @@ cdef public api:
  source_column_idx += 1
 
  return dict(zip(column_names, data_columns)), index
+
+cdef _set_col_children_metadata(Column col,
+ column_name_info& col_meta):
+ cdef column_name_info child_info
+ if isinstance(col.dtype, cudf.StructDtype):
+ for i, (child_col, name) in enumerate(
+ zip(col.children, list(col.dtype.fields))
+ ):
+ child_info.name = name.encode()
+ col_meta.children.push_back(child_info)
+ _set_col_children_metadata(
+ child_col, col_meta.children[i]
+ )
+ elif isinstance(col.dtype, cudf.ListDtype):
+ for i, child_col in enumerate(col.children):
+ col_meta.children.push_back(child_info)
+ _set_col_children_metadata(
+ child_col, col_meta.children[i]
+ )
+ else:
+ return
+
+cdef update_struct_field_names(
+ table,
+ vector[column_name_info]& schema_info
+):
+ for i, (name, col) in enumerate(table._data.items()):
+ table._data[name] = update_column_struct_field_names(
+ col, schema_info[i]
+ )
+
+
+cdef Column update_column_struct_field_names(
+ Column col,
+ column_name_info& info
+):
+ cdef vector[string] field_names
+
+ if col.dtype != "object" and col.children:
+ children = list(col.children)
+ for i, child in enumerate(children):
+ children[i] = update_column_struct_field_names(
+ child,
+ info.children[i]
+ ) 
+ col.set_base_children(tuple(children))
+
+ if isinstance(col.dtype, StructDtype):
+ field_names.reserve(len(col.base_children))
+ for i in range(info.children.size()):
+ field_names.push_back(info.children[i].name)
+ col = col._rename_fields(
+ field_names
+ )
+
+ return col

morpheus/_lib/cudf_helpers/__init__.pyi

@@ -1,9 +1,11 @@
 from __future__ import annotations
 import morpheus._lib.cudf_helpers
 import typing
+from cudf.core.dtypes import StructDtype
 import cudf
 
 __all__ = [
+ "StructDtype",
  "cudf"
 ]
 

morpheus/_lib/include/morpheus/llm/llm_context.hpp

@@ -29,12 +29,20 @@
 #include <string>
 #include <vector>
 
+// IWYU mistakenly believes that we could use the forward declares of LLMTask and ControlMessage in fwd.hpp, however an
+// an incomplete type decl cannot be used in a shared_ptr, and in the case of LLMTask in a struct that is used in a
+// shared_ptr.
+// IWYU pragma: no_include "morpheus/llm/fwd.hpp"
+
 namespace morpheus::llm {
 
 struct LLMContextState
 {
  LLMTask task;
  std::shared_ptr<ControlMessage> message;
+
+ // Optional row mask to be applied to the Dataframe by the extractor and task handler to filter rows
+ std::vector<bool> row_mask;
 };
 
 /**
@@ -64,6 +72,7 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
  * @param parent parent context
  * @param name new context name
  * @param inputs input mappings for new context
+ * @param row_mask row mask for new context
  */
  LLMContext(std::shared_ptr<LLMContext> parent, std::string name, input_mappings_t inputs);
 
@@ -190,6 +199,30 @@ class MORPHEUS_EXPORT LLMContext : public std::enable_shared_from_this<LLMContex
  */
  const mrc::pymrc::JSONValues& view_outputs() const;
 
+ /**
+ * @brief Set the row mask indicating which rows of the dataframe are being used to populate the inputs.
+ * This should only be called by the first node in an LLM Engine, typically the Extractor node.
+ *
+ * @param row_mask vector of bools
+ */
+ void set_row_mask(std::vector<bool>&& row_mask);
+
+ /**
+ * @brief Check if the row mask has been set.
+ *
+ * @return true if row mask has been set
+ * @return false if row mask has not been set
+ */
+ bool has_row_mask() const;
+
+ /**
+ * @brief Get the row mask indicating which rows of the dataframe the outputs should be written to.
+ * This should only be called by the task handler.
+ *
+ * @return vector of bools
+ */
+ const std::vector<bool>& get_row_mask() const;
+
  private:
  input_mappings_t::const_iterator find_input(const std::string& node_name, bool throw_if_not_found = true) const;
 

morpheus/_lib/include/morpheus/messages/meta.hpp

@@ -122,6 +122,8 @@ class MessageMeta
  */
  virtual std::optional<std::string> ensure_sliceable_index();
 
+ pybind11::object get_py_object() const;
+
  /**
  * @brief Create MessageMeta cpp object from a python object
  *

morpheus/_lib/llm/__init__.pyi

@@ -68,12 +68,15 @@ class LLMContext():
  @typing.overload
  def get_input(self, node_name: str) -> object: ...
  def get_inputs(self) -> object: ...
+ def get_row_mask(self) -> typing.List[bool]: ...
+ def has_row_mask(self) -> bool: ...
  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: ...
  @typing.overload
  def set_output(self, outputs: object) -> None: ...
+ def set_row_mask(self, row_mask: typing.List[bool]) -> None: ...
  def task(self) -> LLMTask: ...
  @property
  def full_name(self) -> str:

morpheus/_lib/llm/module.cpp

@@ -203,7 +203,10 @@ PYBIND11_MODULE(llm, _module)
  py::overload_cast<const std::string&, mrc::pymrc::JSONValues&&>(&LLMContext::set_output),
  py::arg("output_name"),
  py::arg("output"))
- .def("push", &LLMContext::push, py::arg("name"), py::arg("inputs"));
+ .def("push", &LLMContext::push, py::arg("name"), py::arg("inputs"))
+ .def("set_row_mask", &LLMContext::set_row_mask, py::arg("row_mask"))
+ .def("has_row_mask", &LLMContext::has_row_mask)
+ .def("get_row_mask", &LLMContext::get_row_mask);
 
  py::class_<LLMNodeBase, PyLLMNodeBase<>, std::shared_ptr<LLMNodeBase>>(_module, "LLMNodeBase")
  .def(py::init_alias<>())

morpheus/_lib/src/llm/llm_context.cpp

@@ -224,4 +224,33 @@ const mrc::pymrc::JSONValues& LLMContext::view_outputs() const
  return m_outputs;
 }
 
+void LLMContext::set_row_mask(std::vector<bool>&& row_mask)
+{
+ if (m_parent)
+ {
+ return m_parent->set_row_mask(std::move(row_mask));
+ }
+
+ m_state->row_mask = std::move(row_mask);
+}
+
+bool LLMContext::has_row_mask() const
+{
+ if (m_parent)
+ {
+ return m_parent->has_row_mask();
+ }
+
+ return !m_state->row_mask.empty();
+}
+
+const std::vector<bool>& LLMContext::get_row_mask() const
+{
+ if (m_parent)
+ {
+ return m_parent->get_row_mask();
+ }
+ return m_state->row_mask;
+}
+
 } // namespace morpheus::llm

morpheus/_lib/src/messages/meta.cpp

@@ -68,6 +68,11 @@ TableInfo MessageMeta::get_info() const
  return this->m_data->get_info();
 }
 
+py::object MessageMeta::get_py_object() const
+{
+ return this->m_data->get_py_object();
+}
+
 TableInfo MessageMeta::get_info(const std::string& col_name) const
 {
  auto full_info = this->m_data->get_info();

morpheus/_lib/src/messages/multi.cpp

@@ -264,50 +264,33 @@ std::vector<std::string> MultiMessageInterfaceProxy::get_meta_column_names(const
 
 pybind11::object MultiMessageInterfaceProxy::get_meta(MultiMessage& self)
 {
- // Need to release the GIL before calling `get_meta()`
- pybind11::gil_scoped_release no_gil;
-
- // Get the column and convert to cudf
- auto info = self.get_meta();
-
- // Convert to a python datatable. Automatically gets the GIL
- return CudfHelper::table_from_table_info(info);
+ return MultiMessageInterfaceProxy::get_meta(self, std::vector<std::string>{});
 }
 
 pybind11::object MultiMessageInterfaceProxy::get_meta(MultiMessage& self, std::string col_name)
 {
- TableInfo info;
-
- {
- // Need to release the GIL before calling `get_meta()`
- pybind11::gil_scoped_release no_gil;
-
- // Get the column and convert to cudf
- info = self.get_meta();
- }
-
- auto py_table = CudfHelper::table_from_table_info(info);
-
- // Now convert it to a series by selecting only the column
- return py_table[col_name.c_str()];
+ return MultiMessageInterfaceProxy::get_meta(self)[col_name.c_str()];
 }
 
 pybind11::object MultiMessageInterfaceProxy::get_meta(MultiMessage& self, std::vector<std::string> columns)
 {
- // Need to release the GIL before calling `get_meta()`
- pybind11::gil_scoped_release no_gil;
+ pybind11::object df = self.meta->get_py_object();
+
+ auto row_indexer = pybind11::slice(
+ pybind11::int_(self.mess_offset), pybind11::int_(self.mess_offset + self.mess_count), pybind11::none());
 
- // Get the column and convert to cudf
- auto info = self.get_meta(columns);
+ if (columns.empty())
+ {
+ return df.attr("iloc")[row_indexer];
+ }
 
- // Convert to a python datatable. Automatically gets the GIL
- return CudfHelper::table_from_table_info(info);
+ return df.attr("iloc")[row_indexer][py::cast(columns)];
 }
 
 pybind11::object MultiMessageInterfaceProxy::get_meta(MultiMessage& self, pybind11::none none_obj)
 {
  // Just offload to the overload without columns. This overload is needed to match the python interface
- return MultiMessageInterfaceProxy::get_meta(self);
+ return MultiMessageInterfaceProxy::get_meta(self, std::vector<std::string>{});
 }
 
 pybind11::object MultiMessageInterfaceProxy::get_meta_list(MultiMessage& self, pybind11::object col_name)