from __future__ import annotations
import sys
import os
import kcal
import csv
import argparse
import time
from kcal import link_config
def create_link_desc() -> kcal.LinkDesc:
nodes = [
{"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="kcal_pir_test",
connect_retry_times=3,
connect_retry_interval_ms=1000,
recv_timeout_ms=60000
)
def pir_demo(context, rank: int, file_path: str):
op = kcal.create_pir(context)
key = []
value = []
query = []
output = []
if rank == 0:
kv_file = os.path.join(file_path, "pir_kv.csv")
with open(kv_file, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
if not row:
continue
if len(row) >= 2:
key.append(row[0])
value.append(row[1])
start_time = time.time()
op.ServerPreProcess(key, value)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ServerPreProcess run cost: {duration_ms:.2f} ms")
start_time = time.time()
op.ServerAnswer()
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ServerAnswer run cost: {duration_ms:.2f} ms")
else:
query_file = os.path.join(file_path, "pir_query.csv")
with open(query_file, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
query.append(row[0])
start_time = time.time()
op.ClientQuery(query, output, kcal.DummyMode.NORMAL)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ClientQuery run cost: {duration_ms:.2f} ms")
print("query result:", output)
def pir_file_demo(context, rank: int, file_path: str):
op = kcal.create_pir(context)
os.makedirs(file_path, exist_ok=True)
output = []
if rank == 0:
kv_file = os.path.join(file_path, "pir_kv.csv")
start_time = time.time()
op.ServerPreProcess(kv_file, file_path)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ServerPreProcess run cost: {duration_ms:.2f} ms")
start_time = time.time()
op.ServerAnswer(file_path, 0)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ServerAnswer run cost: {duration_ms:.2f} ms")
else:
query_file = os.path.join(file_path, "pir_query.csv")
query = []
with open(query_file, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
query.append(row[0])
start_time = time.time()
op.ClientQueryByFile(query, output, kcal.DummyMode.NORMAL)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"ClientQuery run cost: {duration_ms:.2f} ms")
print("query result:", output)
print("query result:", output)
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("--rank", type=int, required=True)
args = parser.parse_args(argv)
if args.rank not in [0, 1]:
print("Error: --rank must be 0 or 1", file=sys.stderr)
sys.exit(1)
config = kcal.Config()
config.nodeId = args.rank
config.worldSize = 2
config.fixBits = 3
config.threadCount = 32
config.useSMAlg = args.use_sm_alg
yacl_link_desc = create_link_desc()
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")
if args.mode == "memory":
pir_demo(context, args.rank, args.work_dir)
else:
pir_file_demo(context, args.rank, args.work_dir)
if __name__ == "__main__":
main()