Repository: caikit/caikit
Pull request #131 (2023-05-23 16:00:45)

Jsonable dicts
## Description

This PR adds a new data model utility to support type hinting and proto serialization for arbitrary json-serializable dicts

## Changes

* Add the new `caikit.core.data_model.json_dict` module containing the `JsonDict` type definition and `dict_to_struct`/`struct_to_dict` converter utilities
* Add support in `ModuleBase` to detect `Struct`/`JsonDict` fields and convert accordingly in `to_proto`/`from_proto`

gabe-l-hart: Ok, this PR led to some _wild_ behavior. For some reason, the usage of `struct_pb2.Struct` led to consistent segfaults in CI that were occasionally reproducible in a local build environment. This led us down several rabbit holes, including refactoring how we use `tox` (labels like `py39` instead of `3.9`) and removing bad versions of `grpcio` that were recently released. Eventually, this seemed to all resolve by removing the unused `import` of `struct_pb2` from `base.py`. We can't confirm explicitly that this was the cause, but it seems that after removing that import, things work cleanly.

Patch:
diff --git a/caikit/core/data_model/base.py b/caikit/core/data_model/base.py
index 471dc519..bf9fda64 100644
--- a/caikit/core/data_model/base.py
+++ b/caikit/core/data_model/base.py
@@ -36,7 +36,7 @@ import alog
 
 # Local
 from ..toolkit.errors import error_handler
-from . import enums
+from . import enums, json_dict
 
 log = alog.use_channel("DATAM")
 error = error_handler.get(log)
@@ -618,9 +618,12 @@ class DataBase(metaclass=_DataBaseMetaClass):
 
             elif field in cls._fields_message:
                 if proto.HasField(field):
-                    contained_class = cls.get_class_for_proto(proto_attr)
-                    contained_obj = contained_class.from_proto(proto_attr)
-                    kwargs[field] = contained_obj
+                    if proto_attr.DESCRIPTOR.full_name == "google.protobuf.Struct":
+                        kwargs[field] = json_dict.struct_to_dict(proto_attr)
+                    else:
+                        contained_class = cls.get_class_for_proto(proto_attr)
+                        contained_obj = contained_class.from_proto(proto_attr)
+                        kwargs[field] = contained_obj
 
             elif field in cls._fields_message_repeated:
                 elements = []
@@ -752,7 +755,12 @@ class DataBase(metaclass=_DataBaseMetaClass):
 
             elif field in self._fields_message:
                 subproto = getattr(proto, field)
-                attr.fill_proto(subproto)
+                if subproto.DESCRIPTOR.full_name == "google.protobuf.Struct":
+                    subproto.CopyFrom(
+                        json_dict.dict_to_struct(attr, subproto.__class__)
+                    )
+                else:
+                    attr.fill_proto(subproto)
 
             elif field in self._fields_message_repeated:
                 subproto = getattr(proto, field)
diff --git a/caikit/core/data_model/dataobject.py b/caikit/core/data_model/dataobject.py
index caa27379..6dcc2cd4 100644
--- a/caikit/core/data_model/dataobject.py
+++ b/caikit/core/data_model/dataobject.py
@@ -25,6 +25,7 @@ import dataclasses
 # Third Party
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
+from google.protobuf import struct_pb2
 from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 import numpy as np
 
@@ -37,12 +38,25 @@ import py_to_proto
 from ..toolkit.errors import error_handler
 from . import enums
 from .base import DataBase, _DataBaseMetaClass
+from .json_dict import JsonDict
 
 ## Globals #####################################################################
 
 log = alog.use_channel("SCHEMA")
 error = error_handler.get(log)
 
+# Type mapping for type hints in @dataobject classes
+DATAOBJECT_PY_TO_PROTO_TYPES = {
+    JsonDict: struct_pb2.Struct,
+    np.int32: _descriptor.FieldDescriptor.TYPE_INT32,
+    np.int64: _descriptor.FieldDescriptor.TYPE_INT64,
+    np.uint32: _descriptor.FieldDescriptor.TYPE_UINT32,
+    np.uint64: _descriptor.FieldDescriptor.TYPE_UINT64,
+    np.float32: _descriptor.FieldDescriptor.TYPE_FLOAT,
+    np.float64: _descriptor.FieldDescriptor.TYPE_DOUBLE,
+    **PY_TO_PROTO_TYPES,
+}
+
 # Common package prefix
 CAIKIT_DATA_MODEL = "caikit_data_model"
 
@@ -206,16 +220,6 @@ def render_dataobject_protos(interfaces_dir: str):
 
 ## Implementation Details ######################################################
 
-DATAOBJECT_PY_TO_PROTO_TYPES = {
-    np.int32: _descriptor.FieldDescriptor.TYPE_INT32,
-    np.int64: _descriptor.FieldDescriptor.TYPE_INT64,
-    np.uint32: _descriptor.FieldDescriptor.TYPE_UINT32,
-    np.uint64: _descriptor.FieldDescriptor.TYPE_UINT64,
-    np.float32: _descriptor.FieldDescriptor.TYPE_FLOAT,
-    np.float64: _descriptor.FieldDescriptor.TYPE_DOUBLE,
-    **PY_TO_PROTO_TYPES,
-}
-
 
 def _dataobject_to_proto(*args, **kwargs):
     kwargs.setdefault("type_mapping", DATAOBJECT_PY_TO_PROTO_TYPES)
diff --git a/caikit/core/data_model/json_dict.py b/caikit/core/data_model/json_dict.py
new file mode 100644
index 00000000..7d117a81
--- /dev/null
+++ b/caikit/core/data_model/json_dict.py
@@ -0,0 +1,138 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""This module holds common utilities for managing arbitrary JSON serializable
+dicts as protobuf Struct objects
+"""
+# Standard
+from typing import Dict, List, Optional, Union
+
+# Third Party
+from google.protobuf import struct_pb2
+from google.protobuf.message_factory import GetMessageClass
+
+# Type hints for JSON serializable dicts
+JsonDictValue = Union[
+    int,
+    float,
+    str,
+    bool,
+    type(None),
+    List["JsonDictValue"],
+    "JsonDict",
+]
+JsonDict = Dict[str, JsonDictValue]
+
+
+def dict_to_struct(
+    dictionary: JsonDict,
+    struct_class: Optional[type] = None,
+    value_class: Optional[type] = None,
+    list_value_class: Optional[type] = None,
+) -> struct_pb2.Struct:
+    """Convert a python dict to a protobuf Struct"""
+    if struct_class is None:
+        struct_class = struct_pb2.Struct
+        value_class = struct_pb2.Value
+        list_value_class = struct_pb2.ListValue
+    else:
+        if value_class is None:
+            value_class = GetMessageClass(
+                struct_class.DESCRIPTOR.file.pool.FindMessageTypeByName(
+                    "google.protobuf.Value"
+                )
+            )
+        if list_value_class is None:
+            list_value_class = GetMessageClass(
+                struct_class.DESCRIPTOR.file.pool.FindMessageTypeByName(
+                    "google.protobuf.ListValue"
+                )
+            )
+
+    return struct_class(
+        fields={
+            key: _value_to_struct_value(
+                value,
+                struct_class=struct_class,
+                value_class=value_class,
+                list_value_class=list_value_class,
+            )
+            for key, value in dictionary.items()
+        }
+    )
+
+
+def struct_to_dict(struct: struct_pb2.Struct) -> JsonDict:
+    """Convert a struct into the equivalent json dict"""
+    return {key: _struct_value_to_py(val) for key, val in struct.fields.items()}
+
+
+## Implementation Details ######################################################
+
+
+def _value_to_struct_value(value, struct_class, value_class, list_value_class):
+    """Recursive helper to convert python values to struct values"""
+    if value is None:
+        struct_value = value_class(null_value=struct_pb2.NullValue.NULL_VALUE)
+    elif isinstance(value, dict):
+        struct_value = value_class(
+            struct_value=dict_to_struct(
+                value,
+                struct_class=struct_class,
+                value_class=value_class,
+                list_value_class=list_value_class,
+            )
+        )
+    elif isinstance(value, list):
+        struct_value = value_class(
+            list_value=list_value_class(
+                values=(
+                    _value_to_struct_value(
+                        item,
+                        struct_class=struct_class,
+                        value_class=value_class,
+                        list_value_class=list_value_class,
+                    )
+                    for item in value
+                )
+            )
+        )
+    elif isinstance(value, bool):
+        struct_value = value_class(bool_value=value)
+    elif isinstance(value, int):
+        struct_value = value_class(number_value=value)
+    elif isinstance(value, float):
+        struct_value = value_class(number_value=value)
+    elif isinstance(value, str):
+        struct_value = value_class(string_value=value)
+    else:
+        raise ValueError(f"Unsupported value type: {type(value)}")
+    return struct_value
+
+
+def _struct_value_to_py(struct_value: struct_pb2.Value) -> JsonDictValue:
+    """Recursive helper to convert struct values to python values"""
+    which = struct_value.WhichOneof("kind")
+    if which in [None, "null_value"]:
+        return None
+    if which == "number_value":
+        val = struct_value.number_value
+        if int(val) == val:
+            return int(val)
+        return val
+    if which in ["string_value", "bool_value"]:
+        return getattr(struct_value, which)
+    if which == "struct_value":
+        return struct_to_dict(struct_value.struct_value)
+    if which == "list_value":
+        return [_struct_value_to_py(item) for item in struct_value.list_value.values]


Test patch:
diff --git a/tests/core/data_model/test_dataobject.py b/tests/core/data_model/test_dataobject.py
index 9c10198a..9b7df26d 100644
--- a/tests/core/data_model/test_dataobject.py
+++ b/tests/core/data_model/test_dataobject.py
@@ -25,7 +25,8 @@ import tempfile
 
 # Third Party
 from google.protobuf import descriptor as _descriptor
-from google.protobuf import descriptor_pool, message
+from google.protobuf import descriptor_pb2, descriptor_pool, message, struct_pb2
+from google.protobuf.message_factory import MessageFactory
 import numpy as np
 import pytest
 
@@ -43,6 +44,7 @@ from caikit.core.data_model.dataobject import (
     _AUTO_GEN_PROTO_CLASSES,
     render_dataobject_protos,
 )
+from caikit.core.data_model.json_dict import JsonDict
 from caikit.core.toolkit.isa import isprotobufenum
 
 ## Helpers #####################################################################
@@ -54,6 +56,17 @@ def temp_dpool():
     dpool = descriptor_pool.DescriptorPool()
     global_dpool = descriptor_pool._DEFAULT
     descriptor_pool._DEFAULT = dpool
+    fd = descriptor_pb2.FileDescriptorProto()
+    struct_pb2.DESCRIPTOR.CopyToProto(fd)
+    dpool.Add(fd)
+    # HACK! Doing this _appears_ to solve the mysterious segfault cause by using
+    #   Struct inside a temporary descriptor pool. The inspiration for this was
+    #   https://github.com/protocolbuffers/protobuf/issues/12047
+    mf = MessageFactory(dpool)
+    msgs = mf.GetMessages([fd.name])
+    _ = msgs["google.protobuf.Struct"]
+    _ = msgs["google.protobuf.Value"]
+    _ = msgs["google.protobuf.ListValue"]
     yield dpool
     # pylint: disable=duplicate-code
     descriptor_pool._DEFAULT = global_dpool
@@ -727,3 +740,38 @@ def test_np_dtypes():
         descriptor.fields_by_name["float64"].type
         == _descriptor.FieldDescriptor.TYPE_DOUBLE
     )
+
+
+@pytest.mark.parametrize("run_num", range(100))
+def test_dataobject_jsondict(temp_dpool, run_num):
+    """Make sure that a JsonDict type is handled correctly in a dataobject
+
+    NOTE: This test is repeated 100x due to a strange segfault in `upb` that it
+        can trigger. The workaround above in `temp_dpool` should solve it, but
+        we retain the repetition to catch anything that's missed.
+    """
+
+    @dataobject
+    class Foo(DataObjectBase):
+        js_dict: JsonDict
+
+    # Make sure the field has the right type
+    Struct = temp_dpool.FindMessageTypeByName("google.protobuf.Struct")
+    assert Foo._proto_class.DESCRIPTOR.fields_by_name["js_dict"].message_type == Struct
+
+    # Make sure dict is preserved on init
+    js_dict = {"foo": {"bar": [1, 2, 3]}}
+    foo = Foo(js_dict)
+    assert foo.js_dict == js_dict
+
+    # Make sure conversion to struct happens on to_proto
+    foo_proto = foo.to_proto()
+    assert set(foo_proto.js_dict.fields.keys()) == set(js_dict.keys())
+    assert foo_proto.js_dict.fields["foo"].struct_value
+    assert set(foo_proto.js_dict.fields["foo"].struct_value.fields.keys()) == set(
+        js_dict["foo"].keys()
+    )
+
+    # Make sure conversion back to dict happens on from_proto
+    foo2 = Foo.from_proto(foo_proto)
+    assert foo2.js_dict == foo.js_dict
diff --git a/tests/core/data_model/test_json_dict.py b/tests/core/data_model/test_json_dict.py
new file mode 100644
index 00000000..5a82b649
--- /dev/null
+++ b/tests/core/data_model/test_json_dict.py
@@ -0,0 +1,63 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for conversion between python dict and protobuf Struct
+"""
+
+# Third Party
+from google.protobuf import struct_pb2
+import pytest
+
+# Local
+from caikit.core.data_model.json_dict import dict_to_struct, struct_to_dict
+
+
+def test_dict_to_struct_to_dict():
+    """Make sure dict_to_struct can handle all variants"""
+    raw_dict = {
+        "int_val": 1,
+        "float_val": 0.42,
+        "str_val": "asdf",
+        "bool_val": False,
+        "null_val": None,
+        "list_val": [2, 3.14, "qwer", True, None, [1, 2, 3], {"nested": "val"}],
+        "dict_val": {"yep": "works"},
+    }
+
+    # Make sure the dict round trips correctly
+    struct = dict_to_struct(raw_dict)
+    round_trip = struct_to_dict(struct)
+    assert round_trip == raw_dict
+
+    # Make sure the struct representation looks right
+    assert set(struct.fields) == set(raw_dict)
+    assert all(
+        getattr(struct.fields[key], struct.fields[key].WhichOneof("kind")) == val
+        for key, val in raw_dict.items()
+        if not isinstance(val, (list, dict, type(None)))
+    )
+    assert struct.fields["null_val"].WhichOneof("kind") == "null_value"
+    assert struct.fields["null_val"].null_value == struct_pb2.NullValue.NULL_VALUE
+    assert isinstance(struct.fields["list_val"].list_value, struct_pb2.ListValue)
+    assert len(struct.fields["list_val"].list_value.values) == len(raw_dict["list_val"])
+    assert isinstance(struct.fields["dict_val"].struct_value, struct_pb2.Struct)
+    assert len(struct.fields["dict_val"].struct_value.fields) == len(
+        raw_dict["dict_val"]
+    )
+
+
+def test_dict_to_struct_invalid_value():
+    """Make sure that a ValueError is raised if a bad type is encountered"""
+    with pytest.raises(ValueError):
+        dict_to_struct({"foo": 1, "bar": {"baz": b"asdf"}})