feat: improve context and tensor implemention

.github/workflows/build_tensor_wheel.yml → .github/workflows/build_rust_paillier_wheel.yml

@@ -1,4 +1,4 @@
-name: Build FATE-CPU-Tensor
+name: Build Rust Paillier
 
 on:
   workflow_dispatch:
@@ -44,13 +44,13 @@ jobs:
       with:
         manylinux: auto
         command: build
-        args: --release -o dist -m rust/fate_tensor/Cargo.toml
+        args: --release -o dist -m rust/tensor/rust_paillier/Cargo.toml
     - name: macos-maturin
       if: matrix.os == 'macos'
       uses: messense/maturin-action@v1
       with:
         command: build
-        args: --release --no-sdist -o dist -m rust/fate_tensor/Cargo.toml
+        args: --release --no-sdist -o dist -m rust/tensor/rust_paillier/Cargo.toml
     - name: Upload wheels
       uses: actions/upload-artifact@v2
       with:

python/fate_arch/tensor/__init__.py

@@ -1,19 +1,20 @@
 from ._dataloader import LabeledDataloaderWrapper, UnlabeledDataloaderWrapper
-from ._federation import ARBITER, GUEST, HOST
-from ._context import Context, CipherKind
-from ._tensor import (
-    FPTensor,
-    PHETensor,
-)
+from ._parties import Parties, PreludeParty
+from ._tensor import CipherKind, Context, FPTensor, PHETensor
+
+ARBITER = PreludeParty.ARBITER
+GUEST = PreludeParty.GUEST
+HOST = PreludeParty.HOST
 
 __all__ = [
     "FPTensor",
     "PHETensor",
+    "Parties",
     "ARBITER",
     "GUEST",
     "HOST",
     "Context",
     "LabeledDataloaderWrapper",
     "UnlabeledDataloaderWrapper",
-    "CipherKind"
+    "CipherKind",
 ]

python/fate_arch/tensor/_federation.py

@@ -1,74 +1,7 @@
-from typing import List, Union
+class FederationDeserializer:
+    def do_deserialize(self, ctx, party):
+        ...
 
-from fate_arch.common import Party
-from fate_arch.session import get_parties
-
-
-def _get_role_parties(role: str):
-    return get_parties().roles_to_parties([role], strict=False)
-
-
-class _RoleIndexedParty:
-    def __init__(self, role: str, index: int) -> None:
-        assert index >= 0, "index should >= 0"
-        self._role = role
-        self._index = index
-
-    @property
-    def party(self) -> Party:
-        parties = _get_role_parties(self._role)
-        if 0 <= self._index < len(parties):
-            return parties[self._index]
-        raise KeyError(
-            f"index `{self._index}` out of bound `0 <= index < {len(parties)}`"
-        )
-
-
-class _Parties:
-    def __init__(
-        self, parties: List[Union[str, Party, _RoleIndexedParty, "_Parties"]]
-    ) -> None:
-        self._parties = parties
-
-    def _reverse(self):
-        self._parties.reverse()
-        return self
-
-    @property
-    def parties(self) -> List[Party]:
-        flatten = []
-        for p in self._parties:
-            if isinstance(p, str) and (p == "guest" or p == "host" or p == "arbiter"):
-                flatten.extend(_get_role_parties(p))
-            elif isinstance(p, Party):
-                flatten.append(p)
-            elif isinstance(p, _RoleIndexedParty):
-                flatten.append(p.party)
-            elif isinstance(p, _Parties):
-                flatten.extend(p.parties)
-        return flatten
-
-    def __add__(self, other) -> "_Parties":
-        if isinstance(other, Party):
-            return _Parties([self, other])
-        elif isinstance(other, list):
-            return _Parties([self, *other])
-        else:
-            raise ValueError(f"can't add `{other}`")
-
-    def __radd__(self, other) -> "_Parties":
-        return self.__add__(other)._reverse()
-
-
-class _Role(_Parties):
-    def __init__(self, role: str) -> None:
-        self._role = role
-        super().__init__([role])
-
-    def __getitem__(self, key) -> _RoleIndexedParty:
-        return _RoleIndexedParty(self._role, key)
-
-
-ARBITER = _Role("arbiter")
-GUEST = _Role("guest")
-HOST = _Role("host")
+    @classmethod
+    def make_frac_key(cls, base_key, frac_key):
+        return f"{base_key}__frac__{frac_key}"