mplang-nightly 0.1.dev268__py3-none-any.whl → 0.1.dev270__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mplang/__init__.py +391 -17
- mplang/{v2/backends → backends}/__init__.py +9 -7
- mplang/{v2/backends → backends}/bfv_impl.py +6 -6
- mplang/{v2/backends → backends}/crypto_impl.py +6 -6
- mplang/{v2/backends → backends}/field_impl.py +5 -5
- mplang/{v2/backends → backends}/func_impl.py +4 -4
- mplang/{v2/backends → backends}/phe_impl.py +3 -3
- mplang/{v2/backends → backends}/simp_design.md +1 -1
- mplang/{v2/backends → backends}/simp_driver/__init__.py +5 -5
- mplang/{v2/backends → backends}/simp_driver/http.py +8 -8
- mplang/{v2/backends → backends}/simp_driver/mem.py +9 -9
- mplang/{v2/backends → backends}/simp_driver/ops.py +4 -4
- mplang/{v2/backends → backends}/simp_driver/state.py +2 -2
- mplang/{v2/backends → backends}/simp_driver/values.py +2 -2
- mplang/{v2/backends → backends}/simp_worker/__init__.py +3 -3
- mplang/{v2/backends → backends}/simp_worker/http.py +10 -10
- mplang/{v2/backends → backends}/simp_worker/mem.py +1 -1
- mplang/{v2/backends → backends}/simp_worker/ops.py +5 -5
- mplang/{v2/backends → backends}/simp_worker/state.py +2 -4
- mplang/{v2/backends → backends}/spu_impl.py +8 -8
- mplang/{v2/backends → backends}/spu_state.py +4 -4
- mplang/{v2/backends → backends}/store_impl.py +3 -3
- mplang/{v2/backends → backends}/table_impl.py +8 -8
- mplang/{v2/backends → backends}/tee_impl.py +6 -6
- mplang/{v2/backends → backends}/tensor_impl.py +6 -6
- mplang/{v2/cli.py → cli.py} +9 -9
- mplang/{v2/cli_guide.md → cli_guide.md} +12 -12
- mplang/{v2/dialects → dialects}/__init__.py +5 -5
- mplang/{v2/dialects → dialects}/bfv.py +6 -6
- mplang/{v2/dialects → dialects}/crypto.py +5 -5
- mplang/{v2/dialects → dialects}/dtypes.py +2 -2
- mplang/{v2/dialects → dialects}/field.py +3 -3
- mplang/{v2/dialects → dialects}/func.py +2 -2
- mplang/{v2/dialects → dialects}/phe.py +6 -6
- mplang/{v2/dialects → dialects}/simp.py +6 -6
- mplang/{v2/dialects → dialects}/spu.py +7 -7
- mplang/{v2/dialects → dialects}/store.py +2 -2
- mplang/{v2/dialects → dialects}/table.py +3 -3
- mplang/{v2/dialects → dialects}/tee.py +6 -6
- mplang/{v2/dialects → dialects}/tensor.py +5 -5
- mplang/{v2/edsl → edsl}/__init__.py +3 -3
- mplang/{v2/edsl → edsl}/context.py +6 -6
- mplang/{v2/edsl → edsl}/graph.py +5 -5
- mplang/{v2/edsl → edsl}/jit.py +2 -2
- mplang/{v2/edsl → edsl}/object.py +1 -1
- mplang/{v2/edsl → edsl}/primitive.py +5 -5
- mplang/{v2/edsl → edsl}/printer.py +1 -1
- mplang/{v2/edsl → edsl}/serde.py +1 -1
- mplang/{v2/edsl → edsl}/tracer.py +7 -7
- mplang/{v2/edsl → edsl}/typing.py +1 -1
- mplang/{v2/kernels → kernels}/ldpc.cpp +13 -13
- mplang/{v2/kernels → kernels}/okvs.cpp +4 -4
- mplang/{v2/kernels → kernels}/okvs_opt.cpp +46 -31
- mplang/{v2/kernels → kernels}/py_kernels.py +1 -1
- mplang/{v2/libs → libs}/collective.py +5 -5
- mplang/{v2/libs → libs}/device/__init__.py +1 -1
- mplang/{v2/libs → libs}/device/api.py +12 -12
- mplang/{v2/libs → libs}/ml/__init__.py +1 -1
- mplang/{v2/libs → libs}/ml/sgb.py +4 -4
- mplang/{v2/libs → libs}/mpc/__init__.py +3 -3
- mplang/{v2/libs → libs}/mpc/_utils.py +2 -2
- mplang/{v2/libs → libs}/mpc/analytics/aggregation.py +1 -1
- mplang/{v2/libs → libs}/mpc/analytics/groupby.py +2 -2
- mplang/{v2/libs → libs}/mpc/analytics/permutation.py +3 -3
- mplang/{v2/libs → libs}/mpc/ot/base.py +3 -3
- mplang/{v2/libs → libs}/mpc/ot/extension.py +2 -2
- mplang/{v2/libs → libs}/mpc/ot/silent.py +4 -4
- mplang/{v2/libs → libs}/mpc/psi/cuckoo.py +3 -3
- mplang/{v2/libs → libs}/mpc/psi/okvs.py +1 -1
- mplang/{v2/libs → libs}/mpc/psi/okvs_gct.py +19 -13
- mplang/{v2/libs → libs}/mpc/psi/oprf.py +3 -3
- mplang/libs/mpc/psi/rr22.py +303 -0
- mplang/{v2/libs → libs}/mpc/psi/unbalanced.py +4 -4
- mplang/{v2/libs → libs}/mpc/vole/gilboa.py +3 -3
- mplang/{v2/libs → libs}/mpc/vole/ldpc.py +2 -2
- mplang/{v2/libs → libs}/mpc/vole/silver.py +6 -6
- mplang/{v2/runtime → runtime}/interpreter.py +11 -11
- mplang/{v2/runtime → runtime}/value.py +2 -2
- mplang/{v1/runtime → utils}/__init__.py +18 -15
- mplang/{v1/utils → utils}/func_utils.py +1 -1
- {mplang_nightly-0.1.dev268.dist-info → mplang_nightly-0.1.dev270.dist-info}/METADATA +2 -2
- mplang_nightly-0.1.dev270.dist-info/RECORD +102 -0
- mplang/v1/__init__.py +0 -157
- mplang/v1/_device.py +0 -602
- mplang/v1/analysis/__init__.py +0 -37
- mplang/v1/analysis/diagram.py +0 -567
- mplang/v1/core/__init__.py +0 -157
- mplang/v1/core/cluster.py +0 -343
- mplang/v1/core/comm.py +0 -281
- mplang/v1/core/context_mgr.py +0 -50
- mplang/v1/core/dtypes.py +0 -335
- mplang/v1/core/expr/__init__.py +0 -80
- mplang/v1/core/expr/ast.py +0 -542
- mplang/v1/core/expr/evaluator.py +0 -581
- mplang/v1/core/expr/printer.py +0 -285
- mplang/v1/core/expr/transformer.py +0 -141
- mplang/v1/core/expr/utils.py +0 -78
- mplang/v1/core/expr/visitor.py +0 -85
- mplang/v1/core/expr/walk.py +0 -387
- mplang/v1/core/interp.py +0 -160
- mplang/v1/core/mask.py +0 -325
- mplang/v1/core/mpir.py +0 -965
- mplang/v1/core/mpobject.py +0 -117
- mplang/v1/core/mptype.py +0 -407
- mplang/v1/core/pfunc.py +0 -130
- mplang/v1/core/primitive.py +0 -877
- mplang/v1/core/table.py +0 -218
- mplang/v1/core/tensor.py +0 -75
- mplang/v1/core/tracer.py +0 -383
- mplang/v1/host.py +0 -130
- mplang/v1/kernels/__init__.py +0 -41
- mplang/v1/kernels/base.py +0 -125
- mplang/v1/kernels/basic.py +0 -240
- mplang/v1/kernels/context.py +0 -369
- mplang/v1/kernels/crypto.py +0 -122
- mplang/v1/kernels/fhe.py +0 -858
- mplang/v1/kernels/mock_tee.py +0 -72
- mplang/v1/kernels/phe.py +0 -1864
- mplang/v1/kernels/spu.py +0 -341
- mplang/v1/kernels/sql_duckdb.py +0 -44
- mplang/v1/kernels/stablehlo.py +0 -90
- mplang/v1/kernels/value.py +0 -626
- mplang/v1/ops/__init__.py +0 -35
- mplang/v1/ops/base.py +0 -424
- mplang/v1/ops/basic.py +0 -294
- mplang/v1/ops/crypto.py +0 -262
- mplang/v1/ops/fhe.py +0 -272
- mplang/v1/ops/jax_cc.py +0 -147
- mplang/v1/ops/nnx_cc.py +0 -168
- mplang/v1/ops/phe.py +0 -216
- mplang/v1/ops/spu.py +0 -151
- mplang/v1/ops/sql_cc.py +0 -303
- mplang/v1/ops/tee.py +0 -36
- mplang/v1/protos/v1alpha1/mpir_pb2.py +0 -63
- mplang/v1/protos/v1alpha1/mpir_pb2.pyi +0 -557
- mplang/v1/protos/v1alpha1/value_pb2.py +0 -34
- mplang/v1/protos/v1alpha1/value_pb2.pyi +0 -169
- mplang/v1/runtime/channel.py +0 -230
- mplang/v1/runtime/cli.py +0 -451
- mplang/v1/runtime/client.py +0 -456
- mplang/v1/runtime/communicator.py +0 -131
- mplang/v1/runtime/data_providers.py +0 -303
- mplang/v1/runtime/driver.py +0 -324
- mplang/v1/runtime/exceptions.py +0 -27
- mplang/v1/runtime/http_api.md +0 -56
- mplang/v1/runtime/link_comm.py +0 -196
- mplang/v1/runtime/server.py +0 -501
- mplang/v1/runtime/session.py +0 -270
- mplang/v1/runtime/simulation.py +0 -324
- mplang/v1/simp/__init__.py +0 -13
- mplang/v1/simp/api.py +0 -353
- mplang/v1/simp/mpi.py +0 -131
- mplang/v1/simp/party.py +0 -225
- mplang/v1/simp/random.py +0 -120
- mplang/v1/simp/smpc.py +0 -238
- mplang/v1/utils/__init__.py +0 -13
- mplang/v1/utils/crypto.py +0 -32
- mplang/v1/utils/spu_utils.py +0 -130
- mplang/v1/utils/table_utils.py +0 -185
- mplang/v2/__init__.py +0 -424
- mplang/v2/libs/mpc/psi/rr22.py +0 -344
- mplang_nightly-0.1.dev268.dist-info/RECORD +0 -180
- /mplang/{v2/backends → backends}/channel.py +0 -0
- /mplang/{v2/edsl → edsl}/README.md +0 -0
- /mplang/{v2/edsl → edsl}/registry.py +0 -0
- /mplang/{v2/kernels → kernels}/Makefile +0 -0
- /mplang/{v2/kernels → kernels}/__init__.py +0 -0
- /mplang/{v2/kernels → kernels}/gf128.cpp +0 -0
- /mplang/{v2/libs → libs}/device/cluster.py +0 -0
- /mplang/{v2/libs → libs}/mpc/analytics/__init__.py +0 -0
- /mplang/{v2/libs → libs}/mpc/analytics/groupby.md +0 -0
- /mplang/{v2/libs → libs}/mpc/common/constants.py +0 -0
- /mplang/{v2/libs → libs}/mpc/ot/__init__.py +0 -0
- /mplang/{v2/libs → libs}/mpc/psi/__init__.py +0 -0
- /mplang/{v2/libs → libs}/mpc/vole/__init__.py +0 -0
- /mplang/{v2/runtime → runtime}/__init__.py +0 -0
- /mplang/{v2/runtime → runtime}/dialect_state.py +0 -0
- /mplang/{v2/runtime → runtime}/object_store.py +0 -0
- {mplang_nightly-0.1.dev268.dist-info → mplang_nightly-0.1.dev270.dist-info}/WHEEL +0 -0
- {mplang_nightly-0.1.dev268.dist-info → mplang_nightly-0.1.dev270.dist-info}/entry_points.txt +0 -0
- {mplang_nightly-0.1.dev268.dist-info → mplang_nightly-0.1.dev270.dist-info}/licenses/LICENSE +0 -0
|
@@ -1,169 +0,0 @@
|
|
|
1
|
-
"""
|
|
2
|
-
@generated by mypy-protobuf. Do not edit manually!
|
|
3
|
-
isort:skip_file
|
|
4
|
-
Copyright 2025 Ant Group Co., Ltd.
|
|
5
|
-
|
|
6
|
-
Licensed under the Apache License, Version 2.0 (the "License");
|
|
7
|
-
you may not use this file except in compliance with the License.
|
|
8
|
-
You may obtain a copy of the License at
|
|
9
|
-
|
|
10
|
-
http://www.apache.org/licenses/LICENSE-2.0
|
|
11
|
-
|
|
12
|
-
Unless required by applicable law or agreed to in writing, software
|
|
13
|
-
distributed under the License is distributed on an "AS IS" BASIS,
|
|
14
|
-
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
15
|
-
See the License for the specific language governing permissions and
|
|
16
|
-
limitations under the License.
|
|
17
|
-
"""
|
|
18
|
-
import builtins
|
|
19
|
-
import collections.abc
|
|
20
|
-
import google.protobuf.descriptor
|
|
21
|
-
import google.protobuf.internal.containers
|
|
22
|
-
import google.protobuf.internal.enum_type_wrapper
|
|
23
|
-
import google.protobuf.message
|
|
24
|
-
import sys
|
|
25
|
-
import typing
|
|
26
|
-
|
|
27
|
-
if sys.version_info >= (3, 10):
|
|
28
|
-
import typing as typing_extensions
|
|
29
|
-
else:
|
|
30
|
-
import typing_extensions
|
|
31
|
-
|
|
32
|
-
DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
|
|
33
|
-
|
|
34
|
-
@typing_extensions.final
|
|
35
|
-
class ValueAttrProto(google.protobuf.message.Message):
|
|
36
|
-
"""Lightweight attribute proto used solely for ValueProto runtime metadata."""
|
|
37
|
-
|
|
38
|
-
DESCRIPTOR: google.protobuf.descriptor.Descriptor
|
|
39
|
-
|
|
40
|
-
class _AttrType:
|
|
41
|
-
ValueType = typing.NewType("ValueType", builtins.int)
|
|
42
|
-
V: typing_extensions.TypeAlias = ValueType
|
|
43
|
-
|
|
44
|
-
class _AttrTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[ValueAttrProto._AttrType.ValueType], builtins.type):
|
|
45
|
-
DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
|
|
46
|
-
UNDEFINED: ValueAttrProto._AttrType.ValueType # 0
|
|
47
|
-
FLOAT: ValueAttrProto._AttrType.ValueType # 1
|
|
48
|
-
INT: ValueAttrProto._AttrType.ValueType # 2
|
|
49
|
-
STRING: ValueAttrProto._AttrType.ValueType # 3
|
|
50
|
-
BOOL: ValueAttrProto._AttrType.ValueType # 4
|
|
51
|
-
BYTES: ValueAttrProto._AttrType.ValueType # 5
|
|
52
|
-
FLOATS: ValueAttrProto._AttrType.ValueType # 6
|
|
53
|
-
INTS: ValueAttrProto._AttrType.ValueType # 7
|
|
54
|
-
STRINGS: ValueAttrProto._AttrType.ValueType # 8
|
|
55
|
-
EMPTY: ValueAttrProto._AttrType.ValueType # 9
|
|
56
|
-
"""Represents an explicitly empty attribute value (e.g., empty list)"""
|
|
57
|
-
|
|
58
|
-
class AttrType(_AttrType, metaclass=_AttrTypeEnumTypeWrapper): ...
|
|
59
|
-
UNDEFINED: ValueAttrProto.AttrType.ValueType # 0
|
|
60
|
-
FLOAT: ValueAttrProto.AttrType.ValueType # 1
|
|
61
|
-
INT: ValueAttrProto.AttrType.ValueType # 2
|
|
62
|
-
STRING: ValueAttrProto.AttrType.ValueType # 3
|
|
63
|
-
BOOL: ValueAttrProto.AttrType.ValueType # 4
|
|
64
|
-
BYTES: ValueAttrProto.AttrType.ValueType # 5
|
|
65
|
-
FLOATS: ValueAttrProto.AttrType.ValueType # 6
|
|
66
|
-
INTS: ValueAttrProto.AttrType.ValueType # 7
|
|
67
|
-
STRINGS: ValueAttrProto.AttrType.ValueType # 8
|
|
68
|
-
EMPTY: ValueAttrProto.AttrType.ValueType # 9
|
|
69
|
-
"""Represents an explicitly empty attribute value (e.g., empty list)"""
|
|
70
|
-
|
|
71
|
-
TYPE_FIELD_NUMBER: builtins.int
|
|
72
|
-
F_FIELD_NUMBER: builtins.int
|
|
73
|
-
I_FIELD_NUMBER: builtins.int
|
|
74
|
-
S_FIELD_NUMBER: builtins.int
|
|
75
|
-
B_FIELD_NUMBER: builtins.int
|
|
76
|
-
RAW_BYTES_FIELD_NUMBER: builtins.int
|
|
77
|
-
FLOATS_FIELD_NUMBER: builtins.int
|
|
78
|
-
INTS_FIELD_NUMBER: builtins.int
|
|
79
|
-
STRS_FIELD_NUMBER: builtins.int
|
|
80
|
-
type: global___ValueAttrProto.AttrType.ValueType
|
|
81
|
-
f: builtins.float
|
|
82
|
-
i: builtins.int
|
|
83
|
-
s: builtins.str
|
|
84
|
-
b: builtins.bool
|
|
85
|
-
raw_bytes: builtins.bytes
|
|
86
|
-
@property
|
|
87
|
-
def floats(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.float]: ...
|
|
88
|
-
@property
|
|
89
|
-
def ints(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.int]: ...
|
|
90
|
-
@property
|
|
91
|
-
def strs(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: ...
|
|
92
|
-
def __init__(
|
|
93
|
-
self,
|
|
94
|
-
*,
|
|
95
|
-
type: global___ValueAttrProto.AttrType.ValueType = ...,
|
|
96
|
-
f: builtins.float = ...,
|
|
97
|
-
i: builtins.int = ...,
|
|
98
|
-
s: builtins.str = ...,
|
|
99
|
-
b: builtins.bool = ...,
|
|
100
|
-
raw_bytes: builtins.bytes = ...,
|
|
101
|
-
floats: collections.abc.Iterable[builtins.float] | None = ...,
|
|
102
|
-
ints: collections.abc.Iterable[builtins.int] | None = ...,
|
|
103
|
-
strs: collections.abc.Iterable[builtins.str] | None = ...,
|
|
104
|
-
) -> None: ...
|
|
105
|
-
def ClearField(self, field_name: typing_extensions.Literal["b", b"b", "f", b"f", "floats", b"floats", "i", b"i", "ints", b"ints", "raw_bytes", b"raw_bytes", "s", b"s", "strs", b"strs", "type", b"type"]) -> None: ...
|
|
106
|
-
|
|
107
|
-
global___ValueAttrProto = ValueAttrProto
|
|
108
|
-
|
|
109
|
-
@typing_extensions.final
|
|
110
|
-
class ValueProto(google.protobuf.message.Message):
|
|
111
|
-
"""Generic envelope for kernel-level transferable values.
|
|
112
|
-
|
|
113
|
-
DESIGN PRINCIPLES
|
|
114
|
-
* Small, stable schema: only descriptors needed for dynamic dispatch.
|
|
115
|
-
* Payload is opaque to the envelope; per-kind versioning lives in
|
|
116
|
-
value_version.
|
|
117
|
-
* Backward-compatible evolution: only append new optional fields.
|
|
118
|
-
|
|
119
|
-
Versioning Guidelines:
|
|
120
|
-
- value_version: per-kind semantic payload version (maintained by KernelValue
|
|
121
|
-
subclass).
|
|
122
|
-
- Adding fields: assign new unique field numbers; never reuse old numbers.
|
|
123
|
-
- Removing fields: reserve the field number & (optionally) name.
|
|
124
|
-
"""
|
|
125
|
-
|
|
126
|
-
DESCRIPTOR: google.protobuf.descriptor.Descriptor
|
|
127
|
-
|
|
128
|
-
@typing_extensions.final
|
|
129
|
-
class RuntimeAttrsEntry(google.protobuf.message.Message):
|
|
130
|
-
DESCRIPTOR: google.protobuf.descriptor.Descriptor
|
|
131
|
-
|
|
132
|
-
KEY_FIELD_NUMBER: builtins.int
|
|
133
|
-
VALUE_FIELD_NUMBER: builtins.int
|
|
134
|
-
key: builtins.str
|
|
135
|
-
@property
|
|
136
|
-
def value(self) -> global___ValueAttrProto: ...
|
|
137
|
-
def __init__(
|
|
138
|
-
self,
|
|
139
|
-
*,
|
|
140
|
-
key: builtins.str = ...,
|
|
141
|
-
value: global___ValueAttrProto | None = ...,
|
|
142
|
-
) -> None: ...
|
|
143
|
-
def HasField(self, field_name: typing_extensions.Literal["value", b"value"]) -> builtins.bool: ...
|
|
144
|
-
def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ...
|
|
145
|
-
|
|
146
|
-
KIND_FIELD_NUMBER: builtins.int
|
|
147
|
-
VALUE_VERSION_FIELD_NUMBER: builtins.int
|
|
148
|
-
PAYLOAD_FIELD_NUMBER: builtins.int
|
|
149
|
-
RUNTIME_ATTRS_FIELD_NUMBER: builtins.int
|
|
150
|
-
kind: builtins.str
|
|
151
|
-
"""Globally unique identifier for Value subclass, e.g. "mplang.ndarray"."""
|
|
152
|
-
value_version: builtins.int
|
|
153
|
-
"""Per-kind payload schema version (>=1)."""
|
|
154
|
-
payload: builtins.bytes
|
|
155
|
-
"""Primary data payload bytes (layout defined by each Value subclass)."""
|
|
156
|
-
@property
|
|
157
|
-
def runtime_attrs(self) -> google.protobuf.internal.containers.MessageMap[builtins.str, global___ValueAttrProto]:
|
|
158
|
-
"""Additional runtime metadata required to recreate the Value instance."""
|
|
159
|
-
def __init__(
|
|
160
|
-
self,
|
|
161
|
-
*,
|
|
162
|
-
kind: builtins.str = ...,
|
|
163
|
-
value_version: builtins.int = ...,
|
|
164
|
-
payload: builtins.bytes = ...,
|
|
165
|
-
runtime_attrs: collections.abc.Mapping[builtins.str, global___ValueAttrProto] | None = ...,
|
|
166
|
-
) -> None: ...
|
|
167
|
-
def ClearField(self, field_name: typing_extensions.Literal["kind", b"kind", "payload", b"payload", "runtime_attrs", b"runtime_attrs", "value_version", b"value_version"]) -> None: ...
|
|
168
|
-
|
|
169
|
-
global___ValueProto = ValueProto
|
mplang/v1/runtime/channel.py
DELETED
|
@@ -1,230 +0,0 @@
|
|
|
1
|
-
# Copyright 2025 Ant Group Co., Ltd.
|
|
2
|
-
#
|
|
3
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
-
# you may not use this file except in compliance with the License.
|
|
5
|
-
# You may obtain a copy of the License at
|
|
6
|
-
#
|
|
7
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
-
#
|
|
9
|
-
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
-
# See the License for the specific language governing permissions and
|
|
13
|
-
# limitations under the License.
|
|
14
|
-
|
|
15
|
-
"""SPU IChannel implementation that bridges to MPLang CommunicatorBase.
|
|
16
|
-
|
|
17
|
-
This module provides BaseChannel, which allows SPU to reuse MPLang's
|
|
18
|
-
existing communication layer (ThreadCommunicator/HttpCommunicator) instead
|
|
19
|
-
of creating separate BRPC connections.
|
|
20
|
-
"""
|
|
21
|
-
|
|
22
|
-
from __future__ import annotations
|
|
23
|
-
|
|
24
|
-
import logging
|
|
25
|
-
from typing import TYPE_CHECKING
|
|
26
|
-
|
|
27
|
-
import spu.libspu as libspu
|
|
28
|
-
|
|
29
|
-
if TYPE_CHECKING:
|
|
30
|
-
from mplang.v1.core.comm import CommunicatorBase
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
class BaseChannel(libspu.link.IChannel):
|
|
34
|
-
"""Bridge MPLang CommunicatorBase to SPU IChannel interface.
|
|
35
|
-
|
|
36
|
-
This adapter allows SPU to use MPLang's existing communication layer
|
|
37
|
-
(ThreadCommunicator or HttpCommunicator) instead of creating separate
|
|
38
|
-
BRPC connections.
|
|
39
|
-
|
|
40
|
-
Each BaseChannel represents a channel to ONE peer rank.
|
|
41
|
-
|
|
42
|
-
Communication Protocol:
|
|
43
|
-
- SPU calls send(tag, bytes_data) -> MPLang comm.send(peer, key, bytes_data)
|
|
44
|
-
- SPU calls recv(tag) -> bytes_data <- MPLang comm.recv(peer, key)
|
|
45
|
-
|
|
46
|
-
Tag Namespace:
|
|
47
|
-
All tags are prefixed with "spu:" to avoid collision with other
|
|
48
|
-
MPLang traffic on the same communicator.
|
|
49
|
-
"""
|
|
50
|
-
|
|
51
|
-
def __init__(
|
|
52
|
-
self,
|
|
53
|
-
comm: CommunicatorBase,
|
|
54
|
-
local_rank: int,
|
|
55
|
-
peer_rank: int,
|
|
56
|
-
tag_prefix: str = "spu",
|
|
57
|
-
):
|
|
58
|
-
"""Initialize channel to a specific peer.
|
|
59
|
-
|
|
60
|
-
Args:
|
|
61
|
-
comm: MPLang communicator instance (Thread/Http)
|
|
62
|
-
local_rank: Global rank of this party (for logging/debugging)
|
|
63
|
-
peer_rank: Global rank of the peer party
|
|
64
|
-
tag_prefix: Prefix for all tags to avoid collision (default: "spu")
|
|
65
|
-
"""
|
|
66
|
-
super().__init__()
|
|
67
|
-
self._comm = comm
|
|
68
|
-
self._local_rank = local_rank
|
|
69
|
-
self._peer_rank = peer_rank
|
|
70
|
-
self._tag_prefix = tag_prefix
|
|
71
|
-
|
|
72
|
-
logging.debug(
|
|
73
|
-
f"BaseChannel initialized: local_rank={local_rank}, "
|
|
74
|
-
f"peer_rank={peer_rank}, tag_prefix={tag_prefix}"
|
|
75
|
-
)
|
|
76
|
-
|
|
77
|
-
def _make_key(self, tag: str) -> str:
|
|
78
|
-
"""Create unique key for MPLang comm.
|
|
79
|
-
|
|
80
|
-
Prefixes the tag to avoid collision with non-SPU traffic.
|
|
81
|
-
|
|
82
|
-
Args:
|
|
83
|
-
tag: SPU-provided tag (e.g., "send_0", "recv_0")
|
|
84
|
-
|
|
85
|
-
Returns:
|
|
86
|
-
Prefixed key (e.g., "spu:send_0")
|
|
87
|
-
"""
|
|
88
|
-
return f"{self._tag_prefix}:{tag}"
|
|
89
|
-
|
|
90
|
-
def Send(self, tag: str, data: bytes) -> None:
|
|
91
|
-
"""Send bytes to peer (synchronous in SPU semantics).
|
|
92
|
-
|
|
93
|
-
Args:
|
|
94
|
-
tag: Message tag for matching send/recv pairs
|
|
95
|
-
data: Raw bytes to send
|
|
96
|
-
"""
|
|
97
|
-
key = self._make_key(tag)
|
|
98
|
-
logging.debug(
|
|
99
|
-
f"BaseChannel.Send: {self._local_rank} -> {self._peer_rank}, "
|
|
100
|
-
f"tag={tag}, key={key}, size={len(data)}"
|
|
101
|
-
)
|
|
102
|
-
|
|
103
|
-
# Send raw bytes directly
|
|
104
|
-
# Note: CommunicatorBase.send expects Any type, bytes is acceptable
|
|
105
|
-
self._comm.send(self._peer_rank, key, data)
|
|
106
|
-
|
|
107
|
-
def Recv(self, tag: str) -> bytes:
|
|
108
|
-
"""Receive bytes from peer (blocking).
|
|
109
|
-
|
|
110
|
-
Args:
|
|
111
|
-
tag: Message tag for matching send/recv pairs
|
|
112
|
-
|
|
113
|
-
Returns:
|
|
114
|
-
Raw bytes received
|
|
115
|
-
"""
|
|
116
|
-
key = self._make_key(tag)
|
|
117
|
-
logging.debug(
|
|
118
|
-
f"BaseChannel.Recv: {self._local_rank} <- {self._peer_rank}, "
|
|
119
|
-
f"tag={tag}, key={key}"
|
|
120
|
-
)
|
|
121
|
-
|
|
122
|
-
# Receive data (should be bytes)
|
|
123
|
-
data = self._comm.recv(self._peer_rank, key)
|
|
124
|
-
|
|
125
|
-
# Validate data type
|
|
126
|
-
if not isinstance(data, bytes):
|
|
127
|
-
raise TypeError(
|
|
128
|
-
f"Expected bytes from communicator, got {type(data).__name__}. "
|
|
129
|
-
f"Communicator must support raw bytes transmission for SPU channels."
|
|
130
|
-
)
|
|
131
|
-
|
|
132
|
-
logging.debug(
|
|
133
|
-
f"BaseChannel.Recv complete: {self._local_rank} <- {self._peer_rank}, "
|
|
134
|
-
f"tag={tag}, size={len(data)}"
|
|
135
|
-
)
|
|
136
|
-
return data
|
|
137
|
-
|
|
138
|
-
def SendAsync(self, tag: str, data: bytes) -> None:
|
|
139
|
-
"""Async send (MPLang's send is already async at network layer).
|
|
140
|
-
|
|
141
|
-
For HttpCommunicator, the underlying httpx.put() is non-blocking
|
|
142
|
-
at the HTTP client level. For ThreadCommunicator, send is instant
|
|
143
|
-
(memory transfer).
|
|
144
|
-
|
|
145
|
-
Args:
|
|
146
|
-
tag: Message tag
|
|
147
|
-
data: Raw bytes to send
|
|
148
|
-
"""
|
|
149
|
-
# Reuse synchronous send - it's already async underneath
|
|
150
|
-
self.Send(tag, data)
|
|
151
|
-
|
|
152
|
-
def SendAsyncThrottled(self, tag: str, data: bytes) -> None:
|
|
153
|
-
"""Throttled async send.
|
|
154
|
-
|
|
155
|
-
Currently maps to regular SendAsync. Future optimization could
|
|
156
|
-
implement rate limiting if needed.
|
|
157
|
-
|
|
158
|
-
Args:
|
|
159
|
-
tag: Message tag
|
|
160
|
-
data: Raw bytes to send
|
|
161
|
-
"""
|
|
162
|
-
self.SendAsync(tag, data)
|
|
163
|
-
|
|
164
|
-
def TestSend(self, timeout: int) -> None:
|
|
165
|
-
"""Test if this channel can send a dummy msg to peer.
|
|
166
|
-
|
|
167
|
-
Uses fixed 0 seq_id as dummy msg's id to make this function reentrant.
|
|
168
|
-
ConnectToMesh will retry on this multiple times.
|
|
169
|
-
|
|
170
|
-
Args:
|
|
171
|
-
timeout: Timeout in milliseconds
|
|
172
|
-
"""
|
|
173
|
-
# Send a handshake message to test connectivity
|
|
174
|
-
# Use fixed tag "__test__" to make this reentrant (idempotent)
|
|
175
|
-
test_data = b"\x00" # Minimal 1-byte message with seq_id=0
|
|
176
|
-
self.Send("__test__", test_data)
|
|
177
|
-
|
|
178
|
-
def TestRecv(self) -> None:
|
|
179
|
-
"""Wait for dummy msg from peer.
|
|
180
|
-
|
|
181
|
-
Timeout is controlled by recv_timeout_ms in link descriptor.
|
|
182
|
-
"""
|
|
183
|
-
# Receive the handshake message from peer
|
|
184
|
-
# This blocks until message arrives (timeout from desc.recv_timeout_ms)
|
|
185
|
-
test_data = self.Recv("__test__")
|
|
186
|
-
# Validate it's the expected handshake message
|
|
187
|
-
if test_data != b"\x00":
|
|
188
|
-
logging.warning(
|
|
189
|
-
f"TestRecv: unexpected handshake data from {self._peer_rank}, "
|
|
190
|
-
f"expected b'\\x00', got {test_data!r}"
|
|
191
|
-
)
|
|
192
|
-
|
|
193
|
-
def WaitLinkTaskFinish(self) -> None:
|
|
194
|
-
"""Wait for all pending async tasks.
|
|
195
|
-
|
|
196
|
-
For MPLang communicators:
|
|
197
|
-
- ThreadCommunicator: No-op (instant memory transfer)
|
|
198
|
-
- HttpCommunicator: No explicit wait needed (httpx handles it)
|
|
199
|
-
|
|
200
|
-
This is a no-op in current implementation.
|
|
201
|
-
"""
|
|
202
|
-
|
|
203
|
-
def Abort(self) -> None:
|
|
204
|
-
"""Abort communication (cleanup resources).
|
|
205
|
-
|
|
206
|
-
This could be extended to notify the communicator to drop pending
|
|
207
|
-
messages for this channel, but currently is a no-op.
|
|
208
|
-
"""
|
|
209
|
-
logging.warning(
|
|
210
|
-
f"BaseChannel.Abort called: {self._local_rank} <-> {self._peer_rank}"
|
|
211
|
-
)
|
|
212
|
-
# Future: Could call comm.abort_session() if implemented
|
|
213
|
-
|
|
214
|
-
def SetThrottleWindowSize(self, size: int) -> None:
|
|
215
|
-
"""Set throttle window size.
|
|
216
|
-
|
|
217
|
-
Not applicable to MPLang communicators. No-op.
|
|
218
|
-
|
|
219
|
-
Args:
|
|
220
|
-
size: Window size (ignored)
|
|
221
|
-
"""
|
|
222
|
-
|
|
223
|
-
def SetChunkParallelSendSize(self, size: int) -> None:
|
|
224
|
-
"""Set chunk parallel send size.
|
|
225
|
-
|
|
226
|
-
Not applicable to MPLang communicators. No-op.
|
|
227
|
-
|
|
228
|
-
Args:
|
|
229
|
-
size: Chunk size (ignored)
|
|
230
|
-
"""
|