Merge pull request #232 from acompany-develop/feature/takahashi/fix_benchmark

Reduced memory in benchmark

Author: mdonaka

.github/workflows/test-server-all.yml

@@ -198,7 +198,8 @@ jobs:
             string,
             sum,
             token,
-            variance
+            variance,
+            get_computation_result
           ]
     steps:
       - name: Checkout

packages/client/libclient-py/quickmpc/qmpc.py

@@ -6,7 +6,7 @@
 from .qmpc_server import QMPCServer
 from .share import Share
 from .utils.parse_csv import parse, parse_csv
-from .utils.restore import restore
+from .restore import restore
 
 logger = logging.getLogger(__name__)
 # qmpc.JobStatus でアクセスできるようにエイリアスを設定する

packages/client/libclient-py/quickmpc/qmpc_server.py

@@ -147,7 +147,7 @@ def __stream_result(stream: Iterable, job_uuid: str, party: int,
         is_ok: bool = True
         res_list = []
         for res in stream:
-            if path is not None:
+            if path and res.status == JobStatus.Value('COMPLETED'):
                 file_title = "dim1"
                 if res.HasField("is_dim2"):
                     file_title = "dim2"

…t/libclient-py/quickmpc/utils/restore.py …/client/libclient-py/quickmpc/restore.pypackages/client/libclient-py/quickmpc/utils/restore.py renamed to packages/client/libclient-py/quickmpc/restore.py packages/client/libclient-py/quickmpc/utils/restore.py renamed to packages/client/libclient-py/quickmpc/restore.py

@@ -1,11 +1,15 @@
 import csv
 import glob
-from decimal import Decimal
 from typing import Any
 
 import numpy as np
 from natsort import natsorted
 
+import google.protobuf.json_format
+from .proto.common_types.common_types_pb2 import Schema
+from .share import Share
+from .utils.if_present import if_present
+
 
 def get_meta(job_uuid: str, path: str):
     file_name = glob.glob(f"{path}/dim?-{job_uuid}-*")[0]
@@ -28,30 +32,45 @@ def get_result(job_uuid: str, path: str, party: int):
 
 
 def restore(job_uuid: str, path: str, party_size: int) -> Any:
-    schema = []
-    result: Any = []
-
     column_number = get_meta(job_uuid, path)
+
+    schema: Any = [None]*column_number
+    results: Any = []
+
+    is_schema = True if len(
+        glob.glob(f"{path}/schema-{job_uuid}-*")) != 0 else False
     is_dim2 = True if len(
         glob.glob(f"{path}/dim2-{job_uuid}-*")) != 0 else False
     if column_number == 0:
-        return [[]] if is_dim2 else []
+        if is_schema:
+            return {"schema": [], "table": [[]]}
+        elif is_dim2:
+            return [[]]
+        else:
+            return []
 
     for party in range(party_size):
         if party == 0:
-            for val in get_result(job_uuid, f"{path}/schema", party):
-                schema.append(val)
+            for i, val in enumerate(get_result(
+                                    job_uuid, f"{path}/schema", party)):
+                col_sch = google.protobuf.json_format.Parse(
+                    val, Schema())
+                schema[i] = col_sch
 
         itr = 0
         for val in get_result(job_uuid, f"{path}/dim?", party):
-            if itr >= len(result):
-                result.append(Decimal(val))
+            f = Share.get_pre_convert_func(schema[itr % column_number])
+            if itr >= len(results):
+                results.append(f(val))
             else:
-                result[itr] += Decimal(val)
+                results[itr] += f(val)
             itr += 1
+    results = np.array(results)\
+        .reshape(-1, column_number).tolist() if is_dim2 else results
+    if is_dim2 and len(results) == 0:
+        schema = []
+        results = [[]]
 
-    result_float = np.vectorize(float)(result)
-    result = np.array(result_float)\
-        .reshape(-1, column_number).tolist() if is_dim2 else result_float
-    result = {"schema": schema, "table": result} if len(schema) else result
-    return result
+    results = if_present(results, Share.convert_type, schema)
+    results = {"schema": schema, "table": results} if is_schema else results
+    return results

scripts/libclient/src/tests/test_get_computation_result.py

@@ -0,0 +1,108 @@
+import math
+import time
+import os
+
+import pytest
+from utils import get_result, qmpc
+
+from tests.common import data_id
+from quickmpc import JobStatus
+
+def execute_computation_param(dataIds=[data_id([[1, 2, 3], [4, 5, 6]])],
+                              src=[1, 2, 3]):
+    return (dataIds, src)
+
+
+@pytest.mark.parametrize(
+    ("param", "expected"),
+    [
+        # usually case
+        (execute_computation_param(),
+         [5, 7, 9]),
+
+        # small table size case
+        (execute_computation_param(dataIds=[data_id([[1]])],
+                                   src=[1]),
+         [1]),
+
+        # large data case
+        (execute_computation_param(dataIds=[data_id([[10**18], [10**18]])],
+                                   src=[1]),
+            [2*10**18]),
+
+        # small data case
+        (execute_computation_param(dataIds=[data_id([[10**-10], [10**-10]])],
+                                   src=[1]),
+            [2*10**-10]),
+
+        # duplicated src case
+        (execute_computation_param(src=[1, 2, 2, 3, 3, 3]),
+         [5, 7, 7, 9, 9, 9]),
+    ]
+)
+def test_get_computation_result(param: tuple, expected: list):
+    # 総和計算リクエスト送信
+    res = qmpc.sum(*param)
+    assert (res["is_ok"])
+    for i in range(10000):
+        time.sleep(i+1)
+        get_res = qmpc.get_computation_result(res["job_uuid"])
+        assert (get_res["is_ok"])
+
+        if get_res["results"] is None:
+            continue
+        # 正しく取得できたか確認
+        for x, y in zip(get_res["results"], expected):
+            assert (math.isclose(x, y, abs_tol=0.1))
+        break
+
+
+@pytest.mark.parametrize(
+    ("param", "expected"),
+    [
+        # usually case
+        (execute_computation_param(),
+         [5, 7, 9]),
+
+        # small table size case
+        (execute_computation_param(dataIds=[data_id([[1]])],
+                                   src=[1]),
+         [1]),
+
+        # large data case
+        (execute_computation_param(dataIds=[data_id([[10**18], [10**18]])],
+                                   src=[1]),
+            [2*10**18]),
+
+        # small data case
+        (execute_computation_param(dataIds=[data_id([[10**-10], [10**-10]])],
+                                   src=[1]),
+            [2*10**-10]),
+
+        # duplicated src case
+        (execute_computation_param(src=[1, 2, 2, 3, 3, 3]),
+         [5, 7, 7, 9, 9, 9]),
+    ]
+)
+def test_restore(param: tuple, expected: list):
+    path = "./result"
+    if not os.path.isdir(path):
+        os.mkdir(path)
+
+    res = qmpc.sum(*param)
+    assert (res["is_ok"])
+    for i in range(10000):
+        time.sleep(i+1)
+        get_res = qmpc.get_computation_result(res["job_uuid"], path)
+        assert (get_res["is_ok"])
+
+        all_completed = all([status == JobStatus.COMPLETED
+                                for status in get_res["statuses"]])
+
+        if not all_completed:
+            continue
+        res = qmpc.restore(res["job_uuid"], path)
+        # 正しく取得できたか確認
+        for x, y in zip(res, expected):
+            assert (math.isclose(x, y, abs_tol=0.1))
+        break

scripts/libclient/src/utils.py

@@ -1,6 +1,7 @@
 
 import time
-
+import glob
+import os
 from quickmpc import QMPC, JobStatus
 
 qmpc: QMPC = QMPC([
@@ -9,9 +10,11 @@
     "http://localhost:50003",
 ])
 
-def __try_get_computation_result(job_uuid, is_limit):
+def __try_get_computation_result(job_uuid, is_limit, path = './result'):
     try:
-        get_res = qmpc.get_computation_result(job_uuid)
+        if not os.path.isdir(path):
+            os.mkdir(path)
+        get_res = qmpc.get_computation_result(job_uuid, path)
     except:
         if is_limit:
             raise
@@ -24,6 +27,8 @@ def __try_get_computation_result(job_uuid, is_limit):
                                 for status in get_res["statuses"]])
 
     if all_completed:
+        res = qmpc.restore(job_uuid, path)
+        get_res["results"] = res
         # NOTE
         # 計算結果取得時には含まれていないが，job_uuidがmodelの取得に使われるので
         # このような処理をしている