maxframe 0.1.0b5__cp310-cp310-win32.whl → 1.0.0rc1__cp310-cp310-win32.whl

maxframe/protocol.py

@@ -32,6 +32,7 @@ from .serialization.serializables import (
     EnumField,
     FieldTypes,
     Float64Field,
+    Int32Field,
     ListField,
     ReferenceField,
     Serializable,
@@ -71,6 +72,9 @@ class DagStatus(enum.Enum):
     CANCELLING = 4
     CANCELLED = 5
 
+    def is_terminated(self):
+        return self in (DagStatus.CANCELLED, DagStatus.SUCCEEDED, DagStatus.FAILED)
+
 
 class DimensionIndex(Serializable):
     is_slice: bool = BoolField("is_slice", default=None)
@@ -190,9 +194,9 @@ class ErrorInfo(JsonSerializable):
         "error_tracebacks", FieldTypes.list
     )
     raw_error_source: ErrorSource = EnumField(
-        "raw_error_source", ErrorSource, FieldTypes.int8
+        "raw_error_source", ErrorSource, FieldTypes.int8, default=None
     )
-    raw_error_data: Optional[Exception] = AnyField("raw_error_data")
+    raw_error_data: Optional[Exception] = AnyField("raw_error_data", default=None)
 
     @classmethod
     def from_exception(cls, exc: Exception):
@@ -201,20 +205,29 @@ class ErrorInfo(JsonSerializable):
         return cls(messages, tracebacks, ErrorSource.PYTHON, exc)
 
     def reraise(self):
-        if self.raw_error_source == ErrorSource.PYTHON:
+        if (
+            self.raw_error_source == ErrorSource.PYTHON
+            and self.raw_error_data is not None
+        ):
             raise self.raw_error_data
         raise RemoteException(self.error_messages, self.error_tracebacks, [])
 
     @classmethod
     def from_json(cls, serialized: dict) -> "ErrorInfo":
         kw = serialized.copy()
-        kw["raw_error_source"] = ErrorSource(serialized["raw_error_source"])
+        if kw.get("raw_error_source") is not None:
+            kw["raw_error_source"] = ErrorSource(serialized["raw_error_source"])
+        else:
+            kw["raw_error_source"] = None
+
         if kw.get("raw_error_data"):
             bufs = [base64.b64decode(s) for s in kw["raw_error_data"]]
             try:
                 kw["raw_error_data"] = pickle.loads(bufs[0], buffers=bufs[1:])
             except:
-                kw["raw_error_data"] = None
+                # both error source and data shall be None to make sure
+                # RemoteException is raised.
+                kw["raw_error_source"] = kw["raw_error_data"] = None
         return cls(**kw)
 
     def to_json(self) -> dict:
@@ -227,7 +240,12 @@ class ErrorInfo(JsonSerializable):
         if isinstance(self.raw_error_data, (PickleContainer, RemoteException)):
             err_data_bufs = self.raw_error_data.get_buffers()
         elif isinstance(self.raw_error_data, BaseException):
-            err_data_bufs = pickle_buffers(self.raw_error_data)
+            try:
+                err_data_bufs = pickle_buffers(self.raw_error_data)
+            except:
+                err_data_bufs = None
+                ret["raw_error_source"] = None
+
         if err_data_bufs:
             ret["raw_error_data"] = [
                 base64.b64encode(s).decode() for s in err_data_bufs
@@ -249,9 +267,17 @@ class DagInfo(JsonSerializable):
     error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
     start_timestamp: Optional[float] = Float64Field("start_timestamp", default=None)
     end_timestamp: Optional[float] = Float64Field("end_timestamp", default=None)
+    subdag_infos: Dict[str, "SubDagInfo"] = DictField(
+        "subdag_infos",
+        key_type=FieldTypes.string,
+        value_type=FieldTypes.reference,
+        default_factory=dict,
+    )
 
     @classmethod
-    def from_json(cls, serialized: dict) -> "DagInfo":
+    def from_json(cls, serialized: dict) -> Optional["DagInfo"]:
+        if serialized is None:
+            return None
         kw = serialized.copy()
         kw["status"] = DagStatus(kw["status"])
         if kw.get("tileable_to_result_infos"):
@@ -261,6 +287,10 @@ class DagInfo(JsonSerializable):
             }
         if kw.get("error_info"):
             kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        if kw.get("subdag_infos"):
+            kw["subdag_infos"] = {
+                k: SubDagInfo.from_json(v) for k, v in kw["subdag_infos"].items()
+            }
         return DagInfo(**kw)
 
     def to_json(self) -> dict:
@@ -279,6 +309,8 @@ class DagInfo(JsonSerializable):
             }
         if self.error_info:
             ret["error_info"] = self.error_info.to_json()
+        if self.subdag_infos:
+            ret["subdag_infos"] = {k: v.to_json() for k, v in self.subdag_infos.items()}
         return ret
 
 
@@ -302,7 +334,9 @@ class SessionInfo(JsonSerializable):
     error_info: Optional[ErrorInfo] = ReferenceField("error_info", default=None)
 
     @classmethod
-    def from_json(cls, serialized: dict) -> "SessionInfo":
+    def from_json(cls, serialized: dict) -> Optional["SessionInfo"]:
+        if serialized is None:
+            return None
         kw = serialized.copy()
         if kw.get("dag_infos"):
             kw["dag_infos"] = {
@@ -320,7 +354,10 @@ class SessionInfo(JsonSerializable):
             "idle_timestamp": self.idle_timestamp,
         }
         if self.dag_infos:
-            ret["dag_infos"] = {k: v.to_json() for k, v in self.dag_infos.items()}
+            ret["dag_infos"] = {
+                k: v.to_json() if v is not None else None
+                for k, v in self.dag_infos.items()
+            }
         if self.error_info:
             ret["error_info"] = self.error_info.to_json()
         return ret
@@ -342,7 +379,25 @@ class ExecuteDagRequest(Serializable):
     )
 
 
-class SubDagInfo(Serializable):
+class SubDagSubmitInstanceInfo(JsonSerializable):
+    submit_reason: str = StringField("submit_reason")
+    instance_id: str = StringField("instance_id")
+    subquery_id: Optional[int] = Int32Field("subquery_id", default=None)
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "SubDagSubmitInstanceInfo":
+        return SubDagSubmitInstanceInfo(**serialized)
+
+    def to_json(self) -> dict:
+        ret = {
+            "submit_reason": self.submit_reason,
+            "instance_id": self.instance_id,
+            "subquery_id": self.subquery_id,
+        }
+        return ret
+
+
+class SubDagInfo(JsonSerializable):
     subdag_id: str = StringField("subdag_id")
     status: DagStatus = EnumField("status", DagStatus, FieldTypes.int8, default=None)
     progress: float = Float64Field("progress", default=None)
@@ -355,9 +410,52 @@ class SubDagInfo(Serializable):
         FieldTypes.reference,
         default_factory=dict,
     )
+    start_timestamp: Optional[float] = Float64Field("start_timestamp", default=None)
+    end_timestamp: Optional[float] = Float64Field("end_timestamp", default=None)
+    submit_instances: List[SubDagSubmitInstanceInfo] = ListField(
+        "submit_instances",
+        FieldTypes.reference,
+        default_factory=list,
+    )
+
+    @classmethod
+    def from_json(cls, serialized: dict) -> "SubDagInfo":
+        kw = serialized.copy()
+        kw["status"] = DagStatus(kw["status"])
+        if kw.get("tileable_to_result_infos"):
+            kw["tileable_to_result_infos"] = {
+                k: ResultInfo.from_json(s)
+                for k, s in kw["tileable_to_result_infos"].items()
+            }
+        if kw.get("error_info"):
+            kw["error_info"] = ErrorInfo.from_json(kw["error_info"])
+        if kw.get("submit_instances"):
+            kw["submit_instances"] = [
+                SubDagSubmitInstanceInfo.from_json(s) for s in kw["submit_instances"]
+            ]
+        return SubDagInfo(**kw)
+
+    def to_json(self) -> dict:
+        ret = {
+            "subdag_id": self.subdag_id,
+            "status": self.status.value,
+            "progress": self.progress,
+            "start_timestamp": self.start_timestamp,
+            "end_timestamp": self.end_timestamp,
+        }
+        if self.error_info:
+            ret["error_info"] = self.error_info.to_json()
+        if self.tileable_to_result_infos:
+            ret["tileable_to_result_infos"] = {
+                k: v.to_json() for k, v in self.tileable_to_result_infos.items()
+            }
+        if self.submit_instances:
+            ret["submit_instances"] = [i.to_json() for i in self.submit_instances]
+        return ret
 
 
 class ExecuteSubDagRequest(Serializable):
+    subdag_id: str = StringField("subdag_id")
     dag: TileableGraph = ReferenceField(
         "dag",
         on_serialize=SerializableGraph.from_graph,

maxframe/serialization/core.pxd

@@ -18,6 +18,9 @@ from libc.stdint cimport int32_t, uint64_t
 cdef class Serializer:
     cdef int _serializer_id
 
+    cpdef bint is_public_data_exist(self, dict context, object key)
+    cpdef put_public_data(self, dict context, object key, object value)
+    cpdef get_public_data(self, dict context, object key)
     cpdef serial(self, object obj, dict context)
     cpdef deserial(self, list serialized, dict context, list subs)
     cpdef on_deserial_error(

maxframe/serialization/core.pyi

@@ -29,6 +29,9 @@ class PickleContainer:
 
 class Serializer:
     serializer_id: int
+    def is_public_data_exist(self, context: Dict, key: Any) -> bool: ...
+    def put_public_data(self, context: Dict, key: Any, value: Any) -> None: ...
+    def get_public_data(self, context: Dict, key: Any) -> Any: ...
     def serial(self, obj: Any, context: Dict): ...
     def deserial(self, serialized: List, context: Dict, subs: List[Any]): ...
     def on_deserial_error(

maxframe/serialization/core.pyx

@@ -130,11 +130,30 @@ cdef Serializer get_deserializer(int32_t deserializer_id):
 
 cdef class Serializer:
     serializer_id = None
+    _public_data_context_key = 0x7fffffff - 1
 
     def __cinit__(self):
         # make the value can be referenced with C code
         self._serializer_id = self.serializer_id
 
+    cpdef bint is_public_data_exist(self, dict context, object key):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            return False
+        return key in public_dict
+
+    cpdef put_public_data(self, dict context, object key, object value):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            public_dict = context[self._public_data_context_key] = {}
+        public_dict[key] = value
+
+    cpdef get_public_data(self, dict context, object key):
+        cdef dict public_dict = context.get(self._public_data_context_key, None)
+        if public_dict is None:
+            return None
+        return public_dict.get(key)
+
     cpdef serial(self, object obj, dict context):
         """
         Returns intermediate serialization result of certain object.
@@ -993,17 +1012,20 @@ def serialize(obj, dict context = None):
     cdef list subs
     cdef bint final
     cdef _IdContextHolder id_context_holder = _IdContextHolder()
+    cdef tuple result
 
     context = context if context is not None else dict()
     serialized, subs, final = _serial_single(obj, context, id_context_holder)
     if final or not subs:
         # marked as a leaf node, return directly
-        return [{}, serialized], subs
-
-    serial_stack.append(_SerialStackItem(serialized, subs))
-    return _serialize_with_stack(
-        serial_stack, None, context, id_context_holder, result_bufs_list
-    )
+        result = [{}, serialized], subs
+    else:
+        serial_stack.append(_SerialStackItem(serialized, subs))
+        result = _serialize_with_stack(
+            serial_stack, None, context, id_context_holder, result_bufs_list
+        )
+    result[0][0]["_PUB"] = context.get(Serializer._public_data_context_key)
+    return result
 
 
 async def serialize_with_spawn(
@@ -1036,31 +1058,38 @@ async def serialize_with_spawn(
     cdef list subs
     cdef bint final
     cdef _IdContextHolder id_context_holder = _IdContextHolder()
+    cdef tuple result
 
     context = context if context is not None else dict()
     serialized, subs, final = _serial_single(obj, context, id_context_holder)
     if final or not subs:
         # marked as a leaf node, return directly
-        return [{}, serialized], subs
-
-    serial_stack.append(_SerialStackItem(serialized, subs))
+        result = [{}, serialized], subs
+    else:
+        serial_stack.append(_SerialStackItem(serialized, subs))
 
-    try:
-        result = _serialize_with_stack(
-            serial_stack, None, context, id_context_holder, result_bufs_list, spawn_threshold
-        )
-    except _SerializeObjectOverflow as ex:
-        result = await asyncio.get_running_loop().run_in_executor(
-            executor,
-            _serialize_with_stack,
-            serial_stack,
-            ex.cur_serialized,
-            context,
-            id_context_holder,
-            result_bufs_list,
-            0,
-            ex.num_total_serialized,
-        )
+        try:
+            result = _serialize_with_stack(
+                serial_stack,
+                None,
+                context,
+                id_context_holder,
+                result_bufs_list,
+                spawn_threshold,
+            )
+        except _SerializeObjectOverflow as ex:
+            result = await asyncio.get_running_loop().run_in_executor(
+                executor,
+                _serialize_with_stack,
+                serial_stack,
+                ex.cur_serialized,
+                context,
+                id_context_holder,
+                result_bufs_list,
+                0,
+                ex.num_total_serialized,
+            )
+    result[0][0]["_PUB"] = context.get(Serializer._public_data_context_key)
     return result
 
 

maxframe/serialization/exception.py

@@ -35,7 +35,7 @@ class RemoteException(MaxFrameError):
     def from_exception(cls, exc: Exception):
         try:
             buffers = pickle_buffers(exc)
-        except (TypeError, pickle.PicklingError):
+        except:
             logger.exception("Cannot pickle exception %s", exc)
             buffers = []
 

maxframe/serialization/pandas.py

@@ -176,11 +176,16 @@ class PdTimestampSerializer(Serializer):
 
 class PdTimedeltaSerializer(Serializer):
     def serial(self, obj: pd.Timedelta, context: Dict):
-        return [int(obj.seconds), obj.microseconds, obj.nanoseconds], [], True
+        return [int(obj.seconds), obj.microseconds, obj.nanoseconds, obj.days], [], True
 
     def deserial(self, serialized: List, context: Dict, subs: List):
+        days = 0 if len(serialized) < 4 else serialized[3]
+        seconds, microseconds, nanoseconds = serialized[:3]
         return pd.Timedelta(
-            seconds=serialized[0], microseconds=serialized[1], nanoseconds=serialized[2]
+            days=days,
+            seconds=seconds,
+            microseconds=microseconds,
+            nanoseconds=nanoseconds,
         )
 
 

maxframe/serialization/serializables/core.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import operator
 import weakref
-from typing import Dict, List, Tuple, Type
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Tuple, Type
 
 import msgpack
 
+from ...lib.mmh3 import hash
 from ..core import Placeholder, Serializer, buffered, load_type
 from .field import Field
 from .field_type import DictType, ListType, PrimitiveFieldType, TupleType
@@ -50,11 +51,16 @@ def _is_field_primitive_compound(field: Field):
 class SerializableMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type], properties: Dict):
         # All the fields including misc fields.
+        name_hash = hash(f"{properties.get('__module__')}.{name}")
         all_fields = dict()
+        # mapping field names to base classes
+        field_to_cls_hash = dict()
 
         for base in bases:
-            if hasattr(base, "_FIELDS"):
-                all_fields.update(base._FIELDS)
+            if not hasattr(base, "_FIELDS"):
+                continue
+            all_fields.update(base._FIELDS)
+            field_to_cls_hash.update(base._FIELD_TO_NAME_HASH)
 
         properties_without_fields = {}
         properties_field_slot_names = []
@@ -64,6 +70,8 @@ class SerializableMeta(type):
                 continue
 
             field = all_fields.get(k)
+            # record the field for the class being created
+            field_to_cls_hash[k] = name_hash
             if field is None:
                 properties_field_slot_names.append(k)
             else:
@@ -75,23 +83,40 @@ class SerializableMeta(type):
 
         # Make field order deterministic to serialize it as list instead of dict.
         field_order = list(all_fields)
-        all_fields = dict(sorted(all_fields.items(), key=operator.itemgetter(0)))
         primitive_fields = []
+        primitive_field_names = set()
         non_primitive_fields = []
-        for v in all_fields.values():
+        for field_name, v in all_fields.items():
             if _is_field_primitive_compound(v):
                 primitive_fields.append(v)
+                primitive_field_names.add(field_name)
             else:
                 non_primitive_fields.append(v)
 
+        # count number of fields for every base class
+        cls_to_primitive_field_count = defaultdict(lambda: 0)
+        cls_to_non_primitive_field_count = defaultdict(lambda: 0)
+        for field_name in field_order:
+            cls_hash = field_to_cls_hash[field_name]
+            if field_name in primitive_field_names:
+                cls_to_primitive_field_count[cls_hash] += 1
+            else:
+                cls_to_non_primitive_field_count[cls_hash] += 1
+
         slots = set(properties.pop("__slots__", set()))
         slots.update(properties_field_slot_names)
 
         properties = properties_without_fields
+        properties["_NAME_HASH"] = name_hash
         properties["_FIELDS"] = all_fields
         properties["_FIELD_ORDER"] = field_order
+        properties["_FIELD_TO_NAME_HASH"] = field_to_cls_hash
         properties["_PRIMITIVE_FIELDS"] = primitive_fields
+        properties["_CLS_TO_PRIMITIVE_FIELD_COUNT"] = dict(cls_to_primitive_field_count)
         properties["_NON_PRIMITIVE_FIELDS"] = non_primitive_fields
+        properties["_CLS_TO_NON_PRIMITIVE_FIELD_COUNT"] = dict(
+            cls_to_non_primitive_field_count
+        )
         properties["__slots__"] = tuple(slots)
 
         clz = type.__new__(mcs, name, bases, properties)
@@ -114,10 +139,14 @@ class Serializable(metaclass=SerializableMeta):
     _cache_primitive_serial = False
     _ignore_non_existing_keys = False
 
+    _NAME_HASH: int
     _FIELDS: Dict[str, Field]
     _FIELD_ORDER: List[str]
+    _FIELD_TO_NAME_HASH: Dict[str, int]
     _PRIMITIVE_FIELDS: List[str]
+    _CLS_TO_PRIMITIVE_FIELD_COUNT: Dict[int, int]
     _NON_PRIMITIVE_FIELDS: List[str]
+    _CLS_TO_NON_PRIMITIVE_FIELD_COUNT: Dict[int, int]
 
     def __init__(self, *args, **kwargs):
         fields = self._FIELDS
@@ -180,6 +209,10 @@ class SerializableSerializer(Serializer):
     Leverage DictSerializer to perform serde.
     """
 
+    @classmethod
+    def _get_obj_field_count_key(cls, obj: Serializable):
+        return f"FC_{obj._NAME_HASH}"
+
     @classmethod
     def _get_field_values(cls, obj: Serializable, fields):
         values = []
@@ -210,6 +243,18 @@ class SerializableSerializer(Serializer):
 
         compound_vals = self._get_field_values(obj, obj._NON_PRIMITIVE_FIELDS)
         cls_module = f"{type(obj).__module__}#{type(obj).__qualname__}"
+
+        field_count_key = self._get_obj_field_count_key(obj)
+        if not self.is_public_data_exist(context, field_count_key):
+            # store field distribution for current Serializable
+            counts = [
+                list(obj._CLS_TO_PRIMITIVE_FIELD_COUNT.items()),
+                list(obj._CLS_TO_NON_PRIMITIVE_FIELD_COUNT.items()),
+            ]
+            field_count_data = msgpack.dumps(counts)
+            self.put_public_data(
+                context, self._get_obj_field_count_key(obj), field_count_data
+            )
         return [cls_module, primitive_vals], [compound_vals], False
 
     @staticmethod
@@ -229,6 +274,62 @@ class SerializableSerializer(Serializer):
             else:
                 field.set(obj, value)
 
+    @classmethod
+    def _set_field_values(
+        cls,
+        obj: Serializable,
+        values: List[Any],
+        client_cls_to_field_count: Optional[Dict[str, int]],
+        is_primitive: bool = True,
+    ):
+        obj_class = type(obj)
+        if is_primitive:
+            server_cls_to_field_count = obj_class._CLS_TO_PRIMITIVE_FIELD_COUNT
+            server_fields = obj_class._PRIMITIVE_FIELDS
+        else:
+            server_cls_to_field_count = obj_class._CLS_TO_NON_PRIMITIVE_FIELD_COUNT
+            server_fields = obj_class._NON_PRIMITIVE_FIELDS
+
+        if client_cls_to_field_count:
+            field_num, server_field_num = 0, 0
+            for cls_hash, count in client_cls_to_field_count.items():
+                # cut values and fields given field distribution
+                # at client and server end
+                cls_fields = server_fields[server_field_num : field_num + count]
+                cls_values = values[field_num : field_num + count]
+                for field, value in zip(cls_fields, cls_values):
+                    if not is_primitive or value != {}:
+                        cls._set_field_value(obj, field, value)
+                field_num += count
+                server_field_num += server_cls_to_field_count[cls_hash]
+        else:
+            # todo remove this branch when all versions below v0.1.0b5 is eliminated
+            from .field import AnyField
+
+            # legacy serialization style, with all fields sorted by name
+            if is_primitive:
+                field_attr = "_legacy_deprecated_primitives"
+            else:
+                field_attr = "_legacy_deprecated_non_primitives"
+            deprecated_fields = []
+            deprecated_names = set()
+            if hasattr(obj_class, field_attr):
+                deprecated_names = set(getattr(obj_class, field_attr))
+                for field_name in deprecated_names:
+                    field = AnyField(tag=field_name)
+                    field.name = field_name
+                    deprecated_fields.append(field)
+            server_fields = sorted(
+                server_fields + deprecated_fields, key=lambda f: f.name
+            )
+            for field, value in zip(server_fields, values):
+                if not is_primitive or value != {}:
+                    try:
+                        cls._set_field_value(obj, field, value)
+                    except AttributeError:  # pragma: no cover
+                        if field.name not in deprecated_names:
+                            raise
+
     def deserial(self, serialized: List, context: Dict, subs: List) -> Serializable:
         obj_class_name, primitives = serialized
         obj_class = load_type(obj_class_name, Serializable)
@@ -238,14 +339,20 @@ class SerializableSerializer(Serializer):
 
         obj = obj_class.__new__(obj_class)
 
-        if primitives:
-            for field, value in zip(obj_class._PRIMITIVE_FIELDS, primitives):
-                if value != {}:
-                    self._set_field_value(obj, field, value)
+        field_count_data = self.get_public_data(
+            context, self._get_obj_field_count_key(obj)
+        )
+        if field_count_data is not None:
+            cls_to_prim_key, cls_to_non_prim_key = msgpack.loads(field_count_data)
+            cls_to_prim_key = dict(cls_to_prim_key)
+            cls_to_non_prim_key = dict(cls_to_non_prim_key)
+        else:
+            cls_to_prim_key, cls_to_non_prim_key = None, None
 
+        if primitives:
+            self._set_field_values(obj, primitives, cls_to_prim_key, True)
         if obj_class._NON_PRIMITIVE_FIELDS:
-            for field, value in zip(obj_class._NON_PRIMITIVE_FIELDS, subs[0]):
-                self._set_field_value(obj, field, value)
+            self._set_field_values(obj, subs[0], cls_to_non_prim_key, False)
         obj.__on_deserialize__()
         return obj