SafeMulRound2ThreadDemo.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SafeMul 完整链路 demo：

1. CSP 按 block 执行第一轮，向 DO 发送 payload
2. DO 基于自己的模型参数，使用多进程执行第二轮
3. CSP 汇总第三轮结果

这里只模拟效率实验，不依赖真实 CNN/ResNet 训练：
- 模型参数向量随机生成，长度可取 61706 / 272474
- CSP 的向量组随机生成，先不强制正交
- 重点统计第二轮 DO 提交阶段的多进程耗时
"""

from __future__ import annotations

import argparse
import os
import random
import statistics
import sys
import time
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Sequence, Tuple

sys.path.append(os.path.abspath(os.path.dirname(__file__)))

from utils.SafeMul import SafeInnerProduct


DEFAULT_CASES: Sequence[Tuple[int, int]] = (
    (61706, 512),
    (61706, 1024),
    (272474, 512),
    (272474, 1024),
)


@dataclass(frozen=True)
class BenchmarkCase:
    """单个实验配置。"""

    model_size: int
    vector_count: int

    @property
    def label(self) -> str:
        return f"1x{self.model_size} dot {self.model_size}x{self.vector_count}"


def _build_case(case_text: str) -> BenchmarkCase:
    text = case_text.lower().replace(" ", "")
    if "x" not in text:
        raise ValueError(f"非法 case: {case_text}，格式应为 61706x512")
    left, right = text.split("x", 1)
    return BenchmarkCase(model_size=int(left), vector_count=int(right))


def _iter_cases(case_args: Sequence[str]) -> Iterable[BenchmarkCase]:
    if case_args:
        for item in case_args:
            yield _build_case(item)
        return
    for model_size, vector_count in DEFAULT_CASES:
        yield BenchmarkCase(model_size=model_size, vector_count=vector_count)


def _sample_safe_mul_noise(precision: int, p: int, alpha: int) -> float:
    """对齐 DO._sample_safe_mul_noise 的逻辑。"""
    denom = 2 * (precision ** 2) * (alpha ** 2)
    if denom <= 0:
        return 0.0
    r_max = (p // denom) - 1
    if r_max <= 0:
        return 0.0
    r_cap = 10 ** 6
    r_max = min(r_max, r_cap)
    if r_max <= 0:
        return 0.0
    rnd = random.SystemRandom()
    return float(rnd.randrange(1, r_max + 1))


def _generate_random_vector(length: int, rng: random.Random) -> List[float]:
    """生成随机向量，数值范围与现有实验风格保持接近。"""
    return [rng.uniform(-1.0, 1.0) for _ in range(length)]


def _dot_plain(vec_a: Sequence[float], vec_b: Sequence[float]) -> float:
    """纯 Python 点积，和 DO.safe_mul_round2_process_block 对齐。"""
    total = 0.0
    for x, y in zip(vec_a, vec_b):
        total += x * y
    return total


def _build_random_vector_group(
    model_size: int,
    vector_count: int,
    seed: int,
) -> List[List[float]]:
    """生成 CSP 侧向量组。当前按你的要求，先随机生成，不强制正交。"""
    rng = random.Random(seed)
    return [_generate_random_vector(model_size, rng) for _ in range(vector_count)]


def _split_into_blocks(total_count: int, block_size: int) -> List[Tuple[int, int]]:
    """返回 block 的 [start, end) 切分。"""
    blocks: List[Tuple[int, int]] = []
    for start in range(0, total_count, block_size):
        end = min(start + block_size, total_count)
        blocks.append((start, end))
    return blocks


def _csp_round1_prepare_block(
    vector_group: Sequence[Sequence[float]],
    start: int,
    end: int,
    precision: int,
) -> Dict[str, Any]:
    """模拟 CSP.safe_mul_prepare_payload_block。"""
    sip = SafeInnerProduct(precision_factor=precision)
    vectors = vector_group[start:end]
    extended_vectors = [list(vec) + [1.0] for vec in vectors]
    p, alpha, c_all, _, s_inv = sip.round1_setup_and_encrypt(extended_vectors)
    return {
        "start": start,
        "end": end,
        "p": p,
        "alpha": alpha,
        "C_all": c_all,
        "s_inv": s_inv,
    }


def _do_round2_process_block(
    payload: Dict[str, Any],
    model_vector: Sequence[float],
    vector_group_block: Sequence[Sequence[float]],
    precision: int,
) -> Dict[str, Any]:
    """模拟 DO.safe_mul_round2_process_block。"""
    sip = SafeInnerProduct(precision_factor=precision)
    p = int(payload["p"])
    alpha = int(payload["alpha"])
    c_all = payload["C_all"]

    block_start_time = time.perf_counter()
    r = _sample_safe_mul_noise(precision, p, alpha)
    b_vector_ext = list(model_vector) + [r]
    d_sums = sip.round2_client_process(b_vector_ext, c_all, alpha, p)

    do_part: List[float] = []
    for vec in vector_group_block:
        do_part.append(_dot_plain(vec, model_vector) - r)

    block_elapsed = time.perf_counter() - block_start_time
    return {
        "start": int(payload["start"]),
        "end": int(payload["end"]),
        "d_sums": d_sums,
        "do_part": do_part,
        "elapsed": block_elapsed,
    }


def _csp_round3_finalize_block(
    payload: Dict[str, Any],
    d_sums: Sequence[int],
    do_part: Sequence[float],
    precision: int,
) -> List[float]:
    """模拟 CSP.safe_mul_finalize_block。"""
    sip = SafeInnerProduct(precision_factor=precision)
    p = int(payload["p"])
    alpha = int(payload["alpha"])
    s_inv = int(payload["s_inv"])
    csp_part = sip.round3_decrypt(list(d_sums), s_inv, alpha, p)
    return [csp_part[i] + do_part[i] for i in range(len(csp_part))]


def run_case(
    case: BenchmarkCase,
    block_size: int,
    processes: int,
    precision: int,
    seed: int,
) -> Dict[str, Any]:
    """执行单个 case：CSP 第一步，DO 第二步多进程，CSP 第三步。"""
    vector_seed = seed + case.model_size * 17 + case.vector_count * 19
    model_seed = seed + case.model_size * 31 + case.vector_count * 7

    model_rng = random.Random(model_seed)
    model_vector = _generate_random_vector(case.model_size, model_rng)

    build_start = time.perf_counter()
    vector_group = _build_random_vector_group(
        model_size=case.model_size,
        vector_count=case.vector_count,
        seed=vector_seed,
    )
    build_elapsed = time.perf_counter() - build_start

    blocks = _split_into_blocks(case.vector_count, block_size)

    round1_start = time.perf_counter()
    payloads: List[Dict[str, Any]] = []
    for start, end in blocks:
        payloads.append(
            _csp_round1_prepare_block(
                vector_group=vector_group,
                start=start,
                end=end,
                precision=precision,
            )
        )
    round1_elapsed = time.perf_counter() - round1_start

    round2_start = time.perf_counter()
    worker_count = max(1, min(processes, len(blocks)))
    with ProcessPoolExecutor(max_workers=worker_count) as executor:
        futures = []
        for payload in payloads:
            start = int(payload["start"])
            end = int(payload["end"])
            futures.append(
                executor.submit(
                    _do_round2_process_block,
                    payload,
                    model_vector,
                    vector_group[start:end],
                    precision,
                )
            )
        round2_results = [future.result() for future in futures]
    round2_elapsed = time.perf_counter() - round2_start

    round3_start = time.perf_counter()
    round2_results.sort(key=lambda item: int(item["start"]))
    final_projection: List[float] = []
    block_elapsed_list: List[float] = []
    for payload, result in zip(payloads, round2_results):
        projection_block = _csp_round3_finalize_block(
            payload=payload,
            d_sums=result["d_sums"],
            do_part=result["do_part"],
            precision=precision,
        )
        final_projection.extend(projection_block)
        block_elapsed_list.append(float(result["elapsed"]))
    round3_elapsed = time.perf_counter() - round3_start

    if len(final_projection) != case.vector_count:
        raise ValueError(
            f"最终投影长度异常: expected={case.vector_count}, got={len(final_projection)}"
        )

    checksum = float(sum(final_projection))
    return {
        "build_seconds": build_elapsed,
        "round1_seconds": round1_elapsed,
        "round2_seconds": round2_elapsed,
        "round3_seconds": round3_elapsed,
        "total_seconds": round1_elapsed + round2_elapsed + round3_elapsed,
        "round2_block_avg_seconds": statistics.mean(block_elapsed_list) if block_elapsed_list else 0.0,
        "round2_block_max_seconds": max(block_elapsed_list) if block_elapsed_list else 0.0,
        "block_count": len(blocks),
        "processes": worker_count,
        "checksum": checksum,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="SafeMul 第二轮 DO 多进程提交耗时测试")
    parser.add_argument(
        "--case",
        action="append",
        default=[],
        help="测试用例，格式如 61706x512；可重复传入多次",
    )
    parser.add_argument(
        "--block-size",
        type=int,
        default=26,
        help="CSP 第一轮切分给 DO 的 block 大小",
    )
    parser.add_argument(
        "--processes",
        type=int,
        default=max(1, min(os.cpu_count() or 4, 8)),
        help="DO 第二轮使用的进程数",
    )
    parser.add_argument(
        "--precision",
        type=int,
        default=10 ** 6,
        help="SafeMul 精度缩放因子",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=2025,
        help="随机种子",
    )
    args = parser.parse_args()

    cases = list(_iter_cases(args.case))
    if not cases:
        raise SystemExit("没有可执行的测试用例")

    print("===== SafeMul 多进程效率测试 =====")
    print(f"block_size = {args.block_size}")
    print(f"processes = {args.processes}")
    print(f"precision = {args.precision}")
    print("")

    summary_rows: List[Tuple[str, float]] = []

    for case in cases:
        print(f"--- Case: {case.label} ---")
        result = run_case(
            case=case,
            block_size=max(1, args.block_size),
            processes=max(1, args.processes),
            precision=args.precision,
            seed=args.seed,
        )
        print(f"随机数据生成时间: {result['build_seconds']:.4f}s")
        print(f"CSP 第一轮总时间: {result['round1_seconds']:.4f}s")
        print(f"DO 第二轮多进程总时间: {result['round2_seconds']:.4f}s")
        print(f"CSP 第三轮总时间: {result['round3_seconds']:.4f}s")
        print(f"协议总时间(1+2+3): {result['total_seconds']:.4f}s")
        print(f"第二轮 block 平均时间: {result['round2_block_avg_seconds']:.4f}s")
        print(f"第二轮 block 最大时间: {result['round2_block_max_seconds']:.4f}s")
        print(f"block 数量: {result['block_count']}")
        print(f"进程数量: {result['processes']}")
        print(f"checksum: {result['checksum']:.6f}")
        print("")
        summary_rows.append((case.label, float(result["round2_seconds"])))

    print("===== 四组核心时间（DO 第二轮多进程） =====")
    for label, round2_seconds in summary_rows:
        print(f"{label}: {round2_seconds:.4f}s")


if __name__ == "__main__":
    main()