Merge pull request #9 from shaojian-ant/main

repo-sync-2024-04-12T15:11:18+0800

Author: shaojian-ant

.bazeliskrc

@@ -1,4 +1,5 @@
 common --experimental_repo_remote_exec
+common --experimental_cc_shared_library
 
 build --incompatible_new_actions_api=false 
 build --copt=-fdiagnostics-color=always
@@ -29,3 +30,12 @@ build:linux --action_env=BAZEL_LINKLIBS=-l%:libstdc++.a:-l%:libgcc.a
 build:linux --copt=-fopenmp
 build:linux --linkopt=-fopenmp
 
+# Required by OpenXLA
+build --nocheck_visibility
+
+# default off CUDA build
+build --@rules_cuda//cuda:enable=false
+
+# Only on when asked
+build:gpu --@rules_cuda//cuda:archs=compute_80:compute_80
+build:gpu --@rules_cuda//cuda:enable=true

.bazelrc

@@ -1 +1 @@
-5.4.1
+6.5.0

.bazelversion

@@ -6,4 +6,6 @@ cmake-build-debug-remote
 bazel-*
 
 build
-.DS_Store
+.DS_Store
+
+.idea

.gitignore

@@ -4,17 +4,17 @@
 
 ## 构建
 
-interconnection-impl引用了spu仓库代码，需要根据[spu构建前提](https://github.com/secretflow/spu/blob/main/CONTRIBUTING.md#build)在编译环境上安装好依赖库
+interconnection-impl 引用了 spu 仓库代码，需要根据[ spu 构建前提](https://github.com/secretflow/spu/blob/main/CONTRIBUTING.md#build)在编译环境上安装好依赖库
 
 然后执行以下构建指令：
 
 ```shell
 bazel build ic_impl/ic_main
 ```
 
-## 运行
+## 运行 ECDH-PSI
 
-### ECDH-PSI
+### 命令行传参
 
 本地同时执行以下两条指令：
 
@@ -30,16 +30,49 @@ bazel run ic_impl/ic_main -- -rank=1 -algo=ECDH_PSI -protocol_families=ECC \
         -parties=127.0.0.1:9530,127.0.0.1:9531
 ```
 
-### SS-LR
+### 环境变量传参
 
-运行SS-LR之前，需要先启动Beaver服务。Beaver服务的代码位于SPU仓库中，需要将SPU代码克隆到本地，然后编译并启动Beaver服务：
+为满足北京金融科技产业联盟的调度层互联互通标准对算法组件接口的要求，interconnection-impl 支持 ECDH-PSI 算法从环境变量读取配置参数
+
+当某个参数在环境变量和命令行选项都被指定时，优先选择读取环境变量参数
+
+程序运行需要关闭握手过程：
+```shell
+bazel run ic_impl/ic_main -- -disable_handshake=1
+```
+
+ECDH-PSI 算法配置的环境变量如下表所示。环境变量设置可参考 [ecdh-psi-env-alice.sh](./ic_impl/env/ecdh-psi-env-alice.sh) 和 [ecdh-psi-env-bob.sh](./ic_impl/env/ecdh-psi-env-bob.sh)
+
+| 环境变量                                                   |                   参考值                    |                          描述                          |
+|:-------------------------------------------------------|:----------------------------------------:|:----------------------------------------------------:|
+| runtime.component.parameter.algo                       |                 ecdh_psi                 |                      algorithm                       |
+| runtime.component.parameter.protocol_families          |                   ecc                    |      comma-separated list of protocol families       |
+| runtime.component.parameter.curve_type                 |                curve25519                |                 elliptic curve type                  |
+| runtime.component.parameter.hash_type                  |                 sha_256                  |                      hash type                       |
+| runtime.component.parameter.hash2curve_strategy        |          direct_hash_as_point_x          |                hash to curve strategy                |
+| runtime.component.parameter.point_octet_format         |               uncompressed               |              point Octet-String format               |
+| runtime.component.parameter.bit_length_after_truncated |                    -1                    | optimization method: secondary ciphertext truncation |
+| system.storage.host.url                                |           file://path/to/root            |            root path of input/output file            |
+| runtime.component.input.train_data                     | {"namespace":"data","name":"psi_1.csv"}  |         relative path and name of input file         |
+| runtime.component.parameter.field_names                |                    id                    |                     field names                      |
+| runtime.component.parameter.result_to_rank             |                    -1                    |              which rank gets the result              |
+| runtime.component.output.train_data                    | {"namespace":"output","name":"result_a"} |        relative path and name of output file         |
+
+## 运行 SS-LR
+
+### 启动 Beaver 服务
+
+运行 SS-LR 之前，需要先启动 Beaver 服务。Beaver 服务的代码位于 SPU 仓库中，需要将 SPU 代码克隆到本地，然后编译并启动 Beaver
+服务：
 
 ```shell
 git clone git@github.com:secretflow/spu.git
 cd spu && bazel run libspu/mpc/semi2k/beaver/ttp_server:beaver_server_main -- -port=9449
 ```
 
-启动Beaver服务后，本地同时执行以下两条指令：
+### 命令行传参
+
+启动 Beaver 服务后，本地同时执行以下两条指令：
 
 ```shell
 bazel run ic_impl/ic_main -- -rank=0 -algo=SS_LR -protocol_families=SS \
@@ -55,13 +88,18 @@ bazel run ic_impl/ic_main -- -rank=1 -algo=SS_LR -protocol_families=SS \
         -parties=127.0.0.1:9530,127.0.0.1:9531
 ```
 
-## 环境变量传参
+### 环境变量传参
 
-为满足北京金融科技产业联盟的调度层互联互通标准对算法组件接口的要求，interconnection-impl支持SS-LR算法从环境变量读取配置参数
+为满足北京金融科技产业联盟的调度层互联互通标准对算法组件接口的要求，interconnection-impl 支持 SS-LR 算法从环境变量读取配置参数
 
 当某个参数在环境变量和命令行选项都被指定时，优先选择读取环境变量参数
 
-### 环境变量定义
+程序运行需要关闭握手过程：
+```shell
+bazel run ic_impl/ic_main -- -disable_handshake=1
+```
+
+SS-LR 算法配置的环境变量如下表所示。环境变量设置可参考 [ss-lr-env-alice.sh](./ic_impl/env/ss-lr-env-alice.sh) 和 [ss-lr-env-bob.sh](./ic_impl/env/ss-lr-env-bob.sh)
 
 | 环境变量                                               |                        参考值                        |                            描述                             |
 |:---------------------------------------------------|:-------------------------------------------------:|:---------------------------------------------------------:|
@@ -85,10 +123,11 @@ bazel run ic_impl/ic_main -- -rank=1 -algo=SS_LR -protocol_families=SS \
 | runtime.component.parameter.ttp_server_host        |                      ip:port                      |   remote ip:port or load-balance uri of beaver service    |
 | runtime.component.parameter.ttp_session_id         |               interconnection-root                |               session id of beaver service                |
 | runtime.component.parameter.ttp_adjust_rank        |                         0                         |      which rank do adjust rpc call to beaver service      |
-| system.storage                                     |                file://path/to/root                |             root path of input / output file              |
+| system.storage.host.url                            |                file://path/to/root                |              root path of input/output file               |
 | runtime.component.input.train_data                 | {"namespace":"data","name":"perfect_logit_a.csv"} |           relative path and name of input file            |
 | runtime.component.parameter.skip_rows=1            |                         1                         |            number of skipped rows from dataset            |
-| runtime.component.parameter.has_label              |                       true                        |       if true, label is the last column of dataset        |
+| runtime.component.parameter.label_owner            |                      host.0                       |             which party owns the label column             |
+| runtime.component.parameter.feature_nums           |            {"host.0":10, "guest.0":10}            |             feature column nums of each party             |
 | runtime.component.output.train_data                |     {"namespace":"output","name":"result_a"}      |           relative path and name of output file           |
 
 ## FAQ

README.md

@@ -20,6 +20,10 @@ protocol_deps()
 load("//bazel:repositories.bzl", "ic_impl_deps")
 ic_impl_deps()
 
+load("@psi//bazel:repositories.bzl", "psi_deps")
+
+psi_deps()
+
 # spu
 load("@spulib//bazel:repositories.bzl", "spu_deps")
 spu_deps()
@@ -28,6 +32,20 @@ spu_deps()
 load("@yacl//bazel:repositories.bzl", "yacl_deps")
 yacl_deps()
 
+load("@rules_python//python:repositories.bzl", "py_repositories")
+
+py_repositories()
+
+load("@com_github_nelhage_rules_boost//:boost/boost.bzl", "boost_deps")
+
+boost_deps()
+
+load("@rules_cuda//cuda:repositories.bzl", "register_detected_cuda_toolchains", "rules_cuda_dependencies")
+
+rules_cuda_dependencies()
+
+register_detected_cuda_toolchains()
+
 # xla
 load("@xla//:workspace4.bzl", "xla_workspace4")
 xla_workspace4()

WORKSPACE

@@ -17,9 +17,11 @@ load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 SECRETFLOW_GIT = "https://github.com/secretflow"
 
-SPU_COMMIT_ID  = "b28e086c6b0fc2b01e7be80ea438afd9ef54f3e2"
+SPU_COMMIT_ID  = "8bf3c97da503f1cffd1292c8e365ecbc30675400"
 
-IC_COMMIT_ID  = "a2ffff21a528456c383a113f6a57b98b3c9de6fe"
+PSI_COMMIT_ID  = "58f687b9949bfda4cf31a09bb6d4a3bc1f375757"
+
+IC_COMMIT_ID  = "30e4220b7444d0bb077a9040f1b428632124e31a"
 
 SPU_REPOSITORY = "SPU"
 
@@ -30,7 +32,14 @@ def ic_impl_deps():
         git_repository,
         name = "spulib",
         commit = SPU_COMMIT_ID,
-        remote = "{}/{}.git".format("https://github.com/shaojian-ant", SPU_REPOSITORY),
+        remote = "{}/{}.git".format(SECRETFLOW_GIT, SPU_REPOSITORY),
+    )
+
+    maybe(
+        git_repository,
+        name = "psi",
+        commit = PSI_COMMIT_ID,
+        remote = "{}/psi.git".format(SECRETFLOW_GIT),
     )
 
 def protocol_deps():

bazel/repositories.bzl

@@ -25,6 +25,11 @@ cc_library(
         "//ic_impl:handler",
         "@spulib//libspu/mpc:factory",
         "@com_google_absl//absl/functional:bind_front",
+        "@spulib//libspu/kernel/hal:constants",
+        "@spulib//libspu/kernel/hal:polymorphic",
+        "@spulib//libspu/kernel/hal:type_cast",
+        "@spulib//libspu/kernel/hal:shape_ops",
+        "@spulib//libspu/kernel/hal:public_helper",
     ]
 )
 
@@ -37,7 +42,6 @@ cc_library(
         "//ic_impl:context",
         "//ic_impl/op/sigmoid",
         "//ic_impl/protocol_family/ss",
-        "@spulib//libspu/kernel/hal",
     ]
 )
 

ic_impl/algo/lr/BUILD.bazel

@@ -15,6 +15,7 @@
 #include "ic_impl/algo/lr/lr_context.h"
 
 #include "gflags/gflags.h"
+#include "nlohmann/json.hpp"
 
 #include "ic_impl/op/sigmoid/sigmoid.h"
 #include "ic_impl/util.h"
@@ -30,14 +31,12 @@ DEFINE_double(l0_norm, 0.0, "l0 norm");
 DEFINE_double(l1_norm, 0.0, "l1 norm");
 DEFINE_double(l2_norm, 0.5, "l2 norm");
 
+DECLARE_bool(disable_handshake);
+
 namespace ic_impl::algo::lr {
 
 namespace {
 
-bool SuggestedHasLabel() {
-  return util::GetParamEnv("has_label", FLAGS_has_label);
-}
-
 int64_t SuggestedNumEpoch() {
   return util::GetParamEnv("num_epoch", FLAGS_num_epoch);
 }
@@ -71,13 +70,54 @@ LrHyperParam SuggestedLrHyperParam() {
   return lr_param;
 }
 
+int32_t GetLabelRank(const std::shared_ptr<IcContext>& ic_ctx) {
+  if (!FLAGS_disable_handshake) {
+    return FLAGS_has_label ? ic_ctx->lctx->Rank() : -1;
+  }
+
+  // get parameters from ENV
+  char* label_owner = util::GetParamEnv("label_owner");
+  YACL_ENFORCE(label_owner != nullptr, "label_owner not in ENV");
+
+  size_t word_size = ic_ctx->lctx->WorldSize();
+  for (size_t i = 0; i < word_size; ++i) {
+    if (ic_ctx->lctx->PartyIdByRank(i) == label_owner) {
+      return i;
+    }
+  }
+
+  YACL_THROW("get label_rank failed");
+}
+
+std::vector<int32_t> GetFeatureNums(const std::shared_ptr<IcContext>& ic_ctx) {
+  std::vector<int32_t> feature_nums;
+  if (!FLAGS_disable_handshake) {
+    return feature_nums;  // got through the handshake that follows
+  }
+
+  // get parameters from ENV
+  char* json_str = util::GetParamEnv("feature_nums");
+  YACL_ENFORCE(json_str != nullptr, "feature_nums not in ENV");
+
+  size_t word_size = ic_ctx->lctx->WorldSize();
+  feature_nums.resize(word_size);
+  auto json_object = nlohmann::json::parse(json_str);
+  for (size_t i = 0; i < word_size; ++i) {
+    feature_nums.at(i) = json_object.at(ic_ctx->lctx->PartyIdByRank(i));
+  }
+
+  return feature_nums;
+}
+
 }  // namespace
 
 std::shared_ptr<LrContext> CreateLrContext(std::shared_ptr<IcContext> ic_ctx) {
   auto ctx = std::make_shared<LrContext>();
   ctx->lr_param = SuggestedLrHyperParam();
 
-  ctx->io_param.label_rank = SuggestedHasLabel() ? ic_ctx->lctx->Rank() : -1;
+  ctx->io_param.label_rank = GetLabelRank(ic_ctx);
+
+  ctx->io_param.feature_nums = GetFeatureNums(ic_ctx);
 
   ctx->sigmoid_mode = op::sigmoid::SuggestedSigmoidMode();
 

ic_impl/algo/lr/lr_context.cc

@@ -21,17 +21,23 @@
 #include "gflags/gflags.h"
 #include "libspu/core/config.h"
 #include "libspu/core/encoding.h"
-#include "libspu/kernel/hal/hal.h"
+#include "libspu/kernel/hal/constants.h"
+#include "libspu/kernel/hal/polymorphic.h"
+#include "libspu/kernel/hal/public_helper.h"
+#include "libspu/kernel/hal/shape_ops.h"
+#include "libspu/kernel/hal/type_cast.h"
 #include "libspu/mpc/aby3/type.h"
 #include "libspu/mpc/factory.h"
 #include "libspu/mpc/semi2k/type.h"
+#include "xtensor/xarray.hpp"
 #include "xtensor/xcsv.hpp"
-#include "xtensor/xio.hpp"
+#include "xtensor/xview.hpp"
 
 DEFINE_string(dataset, "data.csv", "dataset file, only csv is supported");
 DEFINE_int32(skip_rows, 1, "skip number of rows from dataset");
 DEFINE_string(lr_output, "/tmp/sslr_result", "full path name of output file");
 DECLARE_int32(rank);
+DECLARE_bool(disable_handshake);
 
 namespace ic_impl::algo::lr {
 
@@ -196,11 +202,17 @@ bool LrHandler::PrepareDataset() {
   int32_t feature_num =
       ctx_->HasLabel() ? dataset_->shape(1) - 1 : dataset_->shape(1);
   YACL_ENFORCE(sample_size > 0);
-  YACL_ENFORCE(feature_num > 0);
 
   ctx_->io_param.sample_size = sample_size;
-  ctx_->io_param.feature_nums.resize(ctx_->ic_ctx->lctx->WorldSize());
-  ctx_->io_param.feature_nums[ctx_->ic_ctx->lctx->Rank()] = feature_num;
+  auto self_rank = ctx_->ic_ctx->lctx->Rank();
+
+  if (FLAGS_disable_handshake) {
+    YACL_ENFORCE(ctx_->io_param.feature_nums.at(self_rank) == feature_num);
+  } else {
+    YACL_ENFORCE(feature_num > 0);
+    ctx_->io_param.feature_nums.resize(ctx_->ic_ctx->lctx->WorldSize());
+    ctx_->io_param.feature_nums.at(self_rank) = feature_num;
+  }
 
   return true;
 }
@@ -676,15 +688,17 @@ float Accuracy(const xt::xarray<float>& y_true,
 
 void ProduceOutput(spu::SPUContext* sctx, const spu::Value& w) {
   // output result shares to the file
-  spu::PtType out_pt_type;
-  auto w_output = spu::decodeFromRing(
-      w.data(), w.dtype(), sctx->config().fxp_fraction_bits(), &out_pt_type);
-  YACL_ENFORCE(out_pt_type == spu::PT_F32);
+  spu::PtType pt_type = getDecodeType(w.dtype());
+  spu::NdArrayRef dst(makePtType(pt_type), w.shape());
+  spu::PtBufferView pv(static_cast<void*>(dst.data()), pt_type, dst.shape(),
+                       dst.strides());
+  spu::decodeFromRing(w.data(), w.dtype(), sctx->config().fxp_fraction_bits(),
+                      &pv);
 
   std::string out_file_name = GetLrOutputFileName();
   std::ofstream of(out_file_name);
   YACL_ENFORCE(of, "open file={} failed", out_file_name);
-  for (auto it = w_output.begin(); it != w_output.end(); ++it) {
+  for (auto it = dst.begin(); it != dst.end(); ++it) {
     auto* item = reinterpret_cast<float*>(it.getRawPtr());
     of << *item << '\n';
   }
@@ -756,10 +770,17 @@ std::unique_ptr<spu::SPUContext> LrHandler::MakeSpuContext() {
 
 spu::Value LrHandler::EncodingDataset(spu::PtBufferView dataset) {
   // encode to ring.
+  auto array = convertToNdArray(dataset);
+  SPU_ENFORCE(array.eltype().isa<spu::PtTy>(), "expect PtType, got={}",
+              array.eltype());
+
+  const spu::PtType pt_type = array.eltype().as<spu::PtTy>()->pt_type();
+  spu::PtBufferView pv(static_cast<const void*>(array.data()), pt_type,
+                       array.shape(), array.strides());
+
   spu::DataType dtype;
   spu::NdArrayRef encoded =
-      encodeToRing(convertToNdArray(dataset),
-                   static_cast<spu::FieldType>(ctx_->ss_param.field_type),
+      encodeToRing(pv, static_cast<spu::FieldType>(ctx_->ss_param.field_type),
                    ctx_->ss_param.fxp_bits, &dtype);
 
   return spu::Value(encoded, dtype);