"""
Use the yacl communication layer to perform MPC operations, eliminating the need to implement socket callbacks at the Python layer.
The entire communication process is completed at the C++ layer through yacl. At the Python layer, only LinkDesc needs to be configured.
Operation mode (running on two terminals separately):
# terminal 1 (rank 0)
python test/yacl_psi_demo.py --rank 0
# terminal 2 (rank 1)
python test/yacl_psi_demo.py --rank 1
"""
from __future__ import annotations
import sys
import os
import csv
import argparse
import time
import kcal
from kcal import link_config
def create_link_desc(is_3rd: bool) -> kcal.LinkDesc:
nodes = [
{"party": "alice", "address": "127.0.0.1:41929"},
{"party": "bob", "address": "127.0.0.1:56815"},
{"party": "alice", "address": "127.0.0.1:56816"}
] if is_3rd else [
{"party": "alice", "address": "127.0.0.1:41929"},
{"party": "bob", "address": "127.0.0.1:56815"}
]
return link_config.create_link_from_nodes(
nodes,
link_id="kcals_psi_test",
connect_retry_times=3,
connect_retry_interval_ms=1000,
recv_timeout_ms=60000
)
def psi_demo(context, rank: int, file_path: str):
op = kcal.create_psi(context)
input_data = []
input_file = os.path.join(file_path, f"{str(rank)}.csv")
with open(input_file, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
input_data.append(row[0])
output_data = []
start_time = time.time()
ret = op.run(input_data, output_data, kcal.TeeMode.OUTPUT_STRING)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"Psi run result:{output_data}")
print(f"Psi run result:{ret}, run cost: {duration_ms:.2f} ms")
def main(argv=None):
parser = argparse.ArgumentParser(description="KCAL python wrapper demo.")
parser.add_argument("--mode", type=str, required=True, choices=['memory', 'file'])
parser.add_argument("--work_dir", type=str, default="./data")
parser.add_argument("--use_sm_alg", default=False, action="store_true")
parser.add_argument("--n_part", default=False, action="store_true")
parser.add_argument("--rank", type=int, required=True)
args = parser.parse_args(argv)
if args.rank not in [0, 1, 2]:
print("Error: --rank must be 0 or 1 or 2", file=sys.stderr)
sys.exit(1)
try:
config = kcal.Config()
config.nodeId = args.rank
config.worldSize = 3 if args.n_part else 2
config.fixBits = 3
config.threadCount = 32
config.useSMAlg = args.use_sm_alg
yacl_link_desc = create_link_desc(args.n_part)
print(f"[Rank {args.rank}] Creating context with yacl...")
context = kcal.Context.create_with_link_config(config, yacl_link_desc, args.rank, args.mode == "file")
psi_demo(context, args.rank, args.work_dir)
except Exception as e:
print(f"Error: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()