From 6ce6cdb835e14fafee714c0276045f5aa4e47270 Mon Sep 17 00:00:00 2001
From: YuanTingHsieh <yuantingh@nvidia.com>
Date: Thu, 3 Oct 2024 17:59:36 -0700
Subject: [PATCH] Update integration test script and upgrade tenseal and psi
 version

---
 ci/run_integration.sh | 10 ++++++++--
 setup.cfg             |  6 ++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/ci/run_integration.sh b/ci/run_integration.sh
index aa4f79002c..4c8bae5bb9 100755
--- a/ci/run_integration.sh
+++ b/ci/run_integration.sh
@@ -45,8 +45,14 @@ remove_pipenv() {
 
 integration_test_tf() {
     echo "Run TF integration test..."
-    # not using pipenv because we need tensorflow package from the container
-    python -m pip install -e .[dev]
+    # since running directly in container, point python to python3.12
+    ln -sfn /usr/bin/python3.12 /usr/bin/python
+    ln -sfn /usr/bin/python3.12 /usr/bin/python3
+    # somehow the base container has blinker which should be removed
+    apt remove -y python3-blinker python-blinker-doc || true
+    # pipenv does not work with TensorFlow so using pip
+    python3.12 -m pip install -e .[dev]
+    python3.12 -m pip install tensorflow[and-cuda]
     export PYTHONPATH=$PWD
     testFolder="tests/integration_test"
     clean_up_snapshot_and_job
diff --git a/setup.cfg b/setup.cfg
index b8da6a86d9..759494a153 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -38,6 +38,10 @@ install_requires =
     pyhocon
 
 [options.extras_require]
+HE =
+    tenseal==0.3.15
+PSI =
+    openmined.psi==2.0.5
 PT =
     torch
     torchvision
@@ -51,6 +55,8 @@ TRACKING =
 CONFIG =
     omegaconf
 app_opt =
+    %(HE)s
+    %(PSI)s
     %(PT)s
     %(SKLEARN)s
     %(TRACKING)s