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(is_3rd: bool) -> kcal.LinkDesc:
nodes = [
{"party": "alice", "address": "127.0.0.1:41929"},
{"party": "bob", "address": "127.0.0.1:56815"},
{"party": "carol", "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="kcal_arith_test",
connect_retry_times=3,
connect_retry_interval_ms=1000,
recv_timeout_ms=60000
)
def create_context(args):
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)
return kcal.Context.create_with_link_config(config, yacl_link_desc, args.rank, args.mode == "file")
def test_basic_arithmetic(context, share1: kcal.MpcShare, share2: kcal.MpcShare, type: str) -> kcal.MpcShare:
print("\n=== Testing Basic Arithmetic Operations ===")
out_share = kcal.MpcShare()
add_op = kcal.create_mpc(context, kcal.AlgorithmsType.ADD)
sub_op = kcal.create_mpc(context, kcal.AlgorithmsType.SUB)
div_op = kcal.create_mpc(context, kcal.AlgorithmsType.DIV)
sum_op = kcal.create_mpc(context, kcal.AlgorithmsType.SUM)
less_op = kcal.create_mpc(context, kcal.AlgorithmsType.LESS)
less_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.LESS_EQUAL)
greater_op = kcal.create_mpc(context, kcal.AlgorithmsType.GREATER)
greater_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.GREATER_EQUAL)
equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.EQUAL)
no_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.NO_EQUAL)
mul_op = kcal.create_mpc(context, kcal.AlgorithmsType.MUL)
avg_op = kcal.create_mpc(context, kcal.AlgorithmsType.AVG)
max_op = kcal.create_mpc(context, kcal.AlgorithmsType.MAX)
min_op = kcal.create_mpc(context, kcal.AlgorithmsType.MIN)
ascend_sort_op = kcal.create_mpc(context, kcal.AlgorithmsType.ASCEND_SORT)
descend_sort_op = kcal.create_mpc(context, kcal.AlgorithmsType.DESCEND_SORT)
if type == 'ADD':
add_op.run([share1, share2], out_share)
elif type == 'SUB':
sub_op.run([share1, share2], out_share)
elif type == 'DIV':
div_op.run([share1, share2], out_share)
elif type == 'MUL':
mul_op.run([share1, share2], out_share)
elif type == 'LESS':
less_op.run([share1, share2], out_share)
elif type == 'LESS_EQUAL':
less_equal_op.run([share1, share2], out_share)
elif type == 'GREATER':
greater_op.run([share1, share2], out_share)
elif type == 'GREATER_EQUAL':
greater_equal_op.run([share1, share2], out_share)
elif type == 'EQUAL':
equal_op.run([share1, share2], out_share)
elif type == 'NO_EQUAL':
no_equal_op.run([share1, share2], out_share)
elif type == 'SUM':
sum_op.run([share1], out_share)
elif type == 'AVG':
avg_op.run([share1], out_share)
elif type == 'MAX':
max_op.run([share1], out_share)
elif type == 'MIN':
min_op.run([share1], out_share)
elif type == 'ASCEND_SORT':
ascend_sort_op.run([share1], out_share)
elif type == 'DESCEND_SORT':
descend_sort_op.run([share1], out_share)
else:
raise ValueError(f"Unkown operation type: {type}")
return out_share
def test_basic_arithmetic_file(context, rank: int, share_file0: str, share_file1: str, type: str, file_path: str) -> str:
print("\n=== Testing Basic Arithmetic File Operations ===")
add_op = kcal.create_mpc(context, kcal.AlgorithmsType.ADD)
sub_op = kcal.create_mpc(context, kcal.AlgorithmsType.SUB)
div_op = kcal.create_mpc(context, kcal.AlgorithmsType.DIV)
mul_op = kcal.create_mpc(context, kcal.AlgorithmsType.MUL)
less_op = kcal.create_mpc(context, kcal.AlgorithmsType.LESS)
less_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.LESS_EQUAL)
greater_op = kcal.create_mpc(context, kcal.AlgorithmsType.GREATER)
greater_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.GREATER_EQUAL)
equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.EQUAL)
no_equal_op = kcal.create_mpc(context, kcal.AlgorithmsType.NO_EQUAL)
sum_op = kcal.create_mpc(context, kcal.AlgorithmsType.SUM)
avg_op = kcal.create_mpc(context, kcal.AlgorithmsType.AVG)
max_op = kcal.create_mpc(context, kcal.AlgorithmsType.MAX)
min_op = kcal.create_mpc(context, kcal.AlgorithmsType.MIN)
ascend_sort_op = kcal.create_mpc(context, kcal.AlgorithmsType.ASCEND_SORT)
descend_sort_op = kcal.create_mpc(context, kcal.AlgorithmsType.DESCEND_SORT)
calculate_file = os.path.join(file_path, f"arithmetic_calculate_result_{str(rank)}{str(time.time())}.csv")
if type == 'ADD':
add_op.run([share_file0, share_file1], calculate_file)
elif type == 'SUB':
sub_op.run([share_file0, share_file1], calculate_file)
elif type == 'MUL':
mul_op.run([share_file0, share_file1], calculate_file)
elif type == 'DIV':
div_op.run([share_file0, share_file1], calculate_file)
elif type == 'LESS':
less_op.run([share_file0, share_file1], calculate_file)
elif type == 'LESS_EQUAL':
less_equal_op.run([share_file0, share_file1], calculate_file)
elif type == 'GREATER':
greater_op.run([share_file0, share_file1], calculate_file)
elif type == 'GREATER_EQUAL':
greater_equal_op.run([share_file0, share_file1], calculate_file)
elif type == 'EQUAL':
equal_op.run([share_file0, share_file1], calculate_file)
elif type == 'NO_EQUAL':
no_equal_op.run([share_file0, share_file1], calculate_file)
elif type == 'SUM':
sum_op.run([share_file0], calculate_file)
elif type == 'AVG':
avg_op.run([share_file0], calculate_file)
elif type == 'MAX':
max_op.run([share_file0], calculate_file)
elif type == 'MIN':
min_op.run([share_file0], calculate_file)
elif type == 'ASCEND_SORT':
ascend_sort_op.run([share_file0], calculate_file)
elif type == 'DESCEND_SORT':
descend_sort_op.run([share_file0], calculate_file)
else:
raise ValueError(f"Unkown operation type: {type}")
return calculate_file
def create_parser():
parser = argparse.ArgumentParser(description="KCAL python wrapper demo.")
try:
parser.add_argument("--mode", type=str, required=True, choices=['memory', 'file'])
parser.add_argument("--work_dir", type=str, default="./data")
parser.add_argument("--operate_type_1st", type=str, required=False)
parser.add_argument("--operate_type_2nd", type=str, required=False)
parser.add_argument("--use_sm_alg", default=False, action="store_true")
parser.add_argument("--n_part", required=False, action="store_true")
parser.add_argument("--rank", type=int, required=True)
return parser
except argparse.ArgumentParser:
print("\nParam error.")
sys.exit(1)
def do_memory_func(context, rank: int, file_path: str, type1: str, type2: str):
input_data_1 = []
input_file_1 = os.path.join(file_path, f"arith_{str(rank)}_1st.csv")
with open(input_file_1, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
input_data_1.append(float(row[0]))
input_data_2 = []
input_file_2 = os.path.join(file_path, f"arith_{str(rank)}_2nd.csv")
with open(input_file_2, 'r', newline = '', encoding = 'utf-8') as f1:
reader = csv.reader(f1)
for row in reader:
input_data_2.append(float(row[0]))
start_time = time.time()
share1 = kcal.MpcShare()
share2 = kcal.MpcShare()
share3 = kcal.MpcShare()
share4 = kcal.MpcShare()
output_share1 = do_memory_combination_calculate(context, rank, input_data_1, share1, share2, type1)
output_share2 = do_memory_combination_calculate(context, rank, input_data_2, share3, share4, type1)
reveal_share_op = kcal.create_reveal_share(context)
output_share3 = test_basic_arithmetic(context, output_share1, output_share2, type2)
output = []
reveal_share_op.run(output_share3, output)
print(f"final result: {output}")
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"Basics arithmetic test completed in: {duration_ms:.2f} ms")
def do_memory_combination_calculate(context, rank: int, input_data, share1: kcal.MpcShare, share2: kcal.MpcShare, type: str)-> kcal.MpcShare:
make_share_op = kcal.create_make_share(context)
if rank == 0:
print("\n rank 0: Processing arithmetic operations...")
make_share_op.run(input_data, 0, share1)
make_share_op.run([], 1, share2)
elif rank == 1:
print("\n rank 1: Processing arithmetic operations...")
make_share_op.run([], 1, share1)
make_share_op.run(input_data, 0, share2)
else:
print("\n rank 2: Processing arithmetic operations...")
make_share_op.run([], 1, share1)
make_share_op.run([], 1, share2)
return test_basic_arithmetic(context, share1, share2, type)
def do_file_func(context, rank: int, file_path: str, type1: str, type2: str):
input_file_1 = os.path.join(file_path, f"arith_{str(rank)}_1st.csv")
input_file_2 = os.path.join(file_path, f"arith_{str(rank)}_2nd.csv")
share_file1 = os.path.join(file_path, f"arithmetic_{str(rank)}1_result_1st.csv")
share_file2 = os.path.join(file_path, f"arithmetic_{str(rank)}2_result_1st.csv")
share_file3 = os.path.join(file_path, f"arithmetic_{str(rank)}1_result_2nd.csv")
share_file4 = os.path.join(file_path, f"arithmetic_{str(rank)}2_result_2nd.csv")
start_time = time.time()
calculate_file1 = do_file_combination_calculate(context, rank, input_file_1, share_file1, share_file2, type1, file_path)
calculate_file2 = do_file_combination_calculate(context, rank, input_file_2, share_file3, share_file4, type1, file_path)
final_calculate_file = test_basic_arithmetic_file(context, rank, calculate_file1, calculate_file2, type2, file_path)
reveal_share_op = kcal.create_reveal_share(context)
reveal_file = os.path.join(file_path, f"arithmetic_reveal_result_{str(rank)}.csv")
reveal_share_op.run(final_calculate_file, reveal_file)
end_time = time.time()
duration_ms = (end_time - start_time) * 1000
print(f"Basics arithmetic test completed in: {duration_ms:.2f} ms")
def do_file_combination_calculate(context, rank: int, input_file, share_file0: str, share_file1: str, type: str, file_path: str)-> kcal.MpcShare:
make_share_op = kcal.create_make_share(context)
if rank == 0:
print("\n rank 0: Processing arithmetic operations...")
make_share_op.run(input_file, 0, share_file0)
make_share_op.run(input_file, 1, share_file1)
elif rank == 1:
print("\n rank 1: Processing arithmetic operations...")
make_share_op.run(input_file, 1, share_file0)
make_share_op.run(input_file, 0, share_file1)
else:
print("\n rank 0: Processing arithmetic operations...")
make_share_op.run(input_file, 1, share_file0)
make_share_op.run(input_file, 1, share_file1)
return test_basic_arithmetic_file(context, rank, share_file0, share_file1, type, file_path)
def main(argv=None):
parser = create_parser()
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)
context = create_context(args)
if args.mode == "memory":
do_memory_func(context, args.rank, args.work_dir, args.operate_type_1st, args.operate_type_2nd)
else:
do_file_func(context, args.rank, args.work_dir, args.operate_type_1st, args.operate_type_2nd)
if __name__ == "__main__":
main()