Add demo model

iflytek · Sep 23, 2022 · 39384ac · 39384ac
1 parent f5f6861
commit 39384ac
Show file tree

Hide file tree

Showing 9 changed files with 264 additions and 3 deletions.
diff --git a/iflearner/business/hetero/builder/builders.py b/iflearner/business/hetero/builder/builders.py
@@ -0,0 +1,24 @@
+#  Copyright 2022 iFLYTEK. All Rights Reserved.
+#  #
+#  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.
+#  ==============================================================================
+
+from typing import Dict
+from iflearner.business.hetero.model.base_model import BaseModel
+from iflearner.business.hetero.builder.demo_builder import DemoBuilder
+from iflearner.business.hetero.builder.lr_builder import LRBuilder
+
+Builders: Dict[str, BaseModel] = {
+    "demo": DemoBuilder(),
+    "logistic_regression": LRBuilder(),
+}
diff --git a/iflearner/business/hetero/builder/demo_builder.py b/iflearner/business/hetero/builder/demo_builder.py
@@ -0,0 +1,59 @@
+#  Copyright 2022 iFLYTEK. All Rights Reserved.
+#  #
+#  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.
+#  ==============================================================================
+
+from iflearner.business.hetero.model.role import Role
+from iflearner.business.hetero.model.base_model import BaseModel
+from iflearner.business.hetero.model.demo import demo_guest, demo_host, demo_arbiter
+
+from iflearner.business.hetero.builder.model_builder import ModelBuilder
+
+
+class DemoBuilder(ModelBuilder):
+
+    def create_role_model_instance(self, role: str) -> BaseModel:
+        """Create a model instance base on specific role.
+
+        Args:
+            role (str): The role name.
+
+        Returns:
+            BaseModel: Return the base class.
+        """
+        if role == Role.guest:
+            return demo_guest.DemoGuest()
+        elif role == Role.host:
+            return demo_host.DemoHost()
+        elif role == Role.arbiter:
+            return demo_arbiter.DemoArbiter()
+
+        raise Exception(f"{role} is not existed.")
+
+    def get_role_model_flow_file(self, role: str) -> str:
+        """Get model flow file by role name.
+
+        Args:
+            role (str): The role name.
+
+        Returns:
+            str: Return the filename.
+        """
+        if role == Role.guest:
+            return "demo_guest_flow.yaml"
+        elif role == Role.host:
+            return "demo_host_flow.yaml"
+        elif role == Role.arbiter:
+            return "demo_arbiter_flow.yaml"
+
+        raise Exception(f"{role} is not existed.")
diff --git a/iflearner/business/hetero/driver.py b/iflearner/business/hetero/driver.py
@@ -19,7 +19,7 @@
 from loguru import logger
 
 from iflearner.business.hetero.parser import Parser
-from iflearner.business.hetero.builder.lr_builder import LRBuilder
+from iflearner.business.hetero.builder.builders import Builders
 from iflearner.communication.hetero.hetero_network import HeteroNetwork
 
 parser = Parser()
@@ -40,8 +40,11 @@ def __init__(self) -> None:
         """Init the class.
         """
         logger.add(f"log/{parser.model_name}_{parser.role_name}.log", backtrace=True, diagnose=True)
-        self._model = LRBuilder().create_role_model_instance(parser.role_name)
-        parser.parse_model_flow_file(join("model", parser.model_name, LRBuilder().get_role_model_flow_file(parser.role_name)))
+        if parser.model_name not in Builders:
+            raise Exception(f"{parser.model_name} is not existed.")
+
+        self._model = Builders[parser.model_name].create_role_model_instance(parser.role_name)
+        parser.parse_model_flow_file(join("model", parser.model_name, Builders[parser.model_name].get_role_model_flow_file(parser.role_name)))
 
         logger.info(f"Model flow: {parser.model_flow}")
         logger.info(f"Network config: {parser.network_config}")

diff --git a/iflearner/business/hetero/model/demo/demo_arbiter.py b/iflearner/business/hetero/model/demo/demo_arbiter.py
@@ -0,0 +1,45 @@
+#  Copyright 2022 iFLYTEK. All Rights Reserved.
+#  #
+#  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.
+#  ==============================================================================
+
+from loguru import logger
+from typing import List, Tuple
+from iflearner.business.hetero.model.role import Role
+from iflearner.business.hetero.model.base_model import BaseModel
+
+
+class DemoArbiter(BaseModel):
+    """A demo showing the model development.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._register_own_step("step1", self.handle_own_step1)
+        self._register_own_step("step2", self.handle_own_step2)
+
+        self._register_another_step(
+            Role.host, "step2", self.handle_host_step2)
+
+    def handle_own_step1(self) -> Tuple[str, bytes]:
+        logger.info("Arbiter step1")
+        return Role.guest, "Arbiter step1 completed.".encode("utf-8")
+
+    def handle_own_step2(self):
+        logger.info("Arbiter step2")
+        return Role.guest, "Arbiter step2 completed.".encode("utf-8")
+
+    def handle_host_step2(self, data: List[Tuple[str, bytes]]):
+        for item in data:
+            logger.info(f"{item[0]}, {item[1].decode('utf-8')}")
diff --git a/iflearner/business/hetero/model/demo/demo_arbiter_flow.yaml b/iflearner/business/hetero/model/demo/demo_arbiter_flow.yaml
@@ -0,0 +1,8 @@
+role: arbiter
+steps:
+  - name: step1
+    upstreams: null
+  - name: step2
+    upstreams:
+      - role: host
+        step: step2
diff --git a/iflearner/business/hetero/model/demo/demo_guest.py b/iflearner/business/hetero/model/demo/demo_guest.py
@@ -0,0 +1,56 @@
+#  Copyright 2022 iFLYTEK. All Rights Reserved.
+#  #
+#  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.
+#  ==============================================================================
+
+from loguru import logger
+from typing import List, Tuple
+from iflearner.business.hetero.model.role import Role
+from iflearner.business.hetero.model.base_model import BaseModel
+
+
+class DemoGuest(BaseModel):
+    """A demo showing the model development.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._register_own_step("step1", self.handle_own_step1)
+        self._register_own_step("step2", self.handle_own_step2)
+
+        self._register_another_step(
+            Role.host, "step1", self.handle_host_step1)
+        self._register_another_step(
+            Role.arbiter, "step1", self.handle_arbiter_step1)
+        self._register_another_step(
+            Role.arbiter, "step2", self.handle_arbiter_step2)
+
+    def handle_own_step1(self) -> Tuple[str, bytes]:
+        logger.info("Guest step1")
+        return Role.host, "Guest step1 completed.".encode("utf-8")
+
+    def handle_own_step2(self):
+        logger.info("Guest step2")
+
+    def handle_host_step1(self, data: List[Tuple[str, bytes]]):
+        for item in data:
+            logger.info(f"{item[0]}, {item[1].decode('utf-8')}")
+
+    def handle_arbiter_step1(self, data: List[Tuple[str, bytes]]):
+        for item in data:
+            logger.info(f"{item[0]}, {item[1].decode('utf-8')}")
+
+    def handle_arbiter_step2(self, data: List[Tuple[str, bytes]]):
+        for item in data:
+            logger.info(f"{item[0]}, {item[1].decode('utf-8')}")
diff --git a/iflearner/business/hetero/model/demo/demo_guest_flow.yaml b/iflearner/business/hetero/model/demo/demo_guest_flow.yaml
@@ -0,0 +1,13 @@
+role: guest
+steps:
+  - name: step1
+    upstreams:
+      - role: host
+        step: step1
+      - role: arbiter
+        step: step1   
+  - name: step2
+    upstreams:
+      - role: arbiter
+        step: step2   
+
diff --git a/iflearner/business/hetero/model/demo/demo_host.py b/iflearner/business/hetero/model/demo/demo_host.py
@@ -0,0 +1,45 @@
+#  Copyright 2022 iFLYTEK. All Rights Reserved.
+#  #
+#  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.
+#  ==============================================================================
+
+from loguru import logger
+from typing import List, Tuple
+from iflearner.business.hetero.model.role import Role
+from iflearner.business.hetero.model.base_model import BaseModel
+
+
+class DemoHost(BaseModel):
+    """A demo showing the model development.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._register_own_step("step1", self.handle_own_step1)
+        self._register_own_step("step2", self.handle_own_step2)
+
+        self._register_another_step(
+            Role.guest, "step1", self.handle_guest_step1)
+
+    def handle_own_step1(self) -> Tuple[str, bytes]:
+        logger.info("Host step1")
+        return Role.guest, "Host step1 completed.".encode("utf-8")
+
+    def handle_own_step2(self):
+        logger.info("Host step2")
+        return Role.arbiter, "Host step2 completed.".encode("utf-8")
+
+    def handle_guest_step1(self, data: List[Tuple[str, bytes]]):
+        for item in data:
+            logger.info(f"{item[0]}, {item[1].decode('utf-8')}")
diff --git a/iflearner/business/hetero/model/demo/demo_host_flow.yaml b/iflearner/business/hetero/model/demo/demo_host_flow.yaml
@@ -0,0 +1,8 @@
+role: host
+steps:
+  - name: step1
+    upstreams: null
+  - name: step2
+    upstreams:
+      - role: guest
+        step: step1