test(distributed): fix termination of tests meant to run on distributed systems to avoid ungraceful exit.

compiler/hpx.ini

@@ -3,7 +3,7 @@ location = ${HPX_LOCATION:$[system.prefix]}
 component_path = $[hpx.location]/lib/hpx:$[system.executable_prefix]/lib/hpx:$[system.executable_prefix]/../lib/hpx
 master_ini_path = $[hpx.location]/share/hpx-<version>:$[system.executable_prefix]/share/hpx-<version>:$[system.executable_prefix]/../share/hpx-<version>
 ini_path = $[hpx.master_ini_path]/ini
-os_threads = 1
+os_threads = 2
 localities = 1
 program_name =
 cmd_line =

compiler/tests/end_to_end_tests/end_to_end_jit_auto_parallelization.cc

@@ -6,6 +6,7 @@
 #include <type_traits>
 
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 ///////////////////////////////////////////////////////////////////////////////
 // Auto-parallelize independent FHE ops /////////////////////////////////////
@@ -73,10 +74,10 @@ func.func @main(%arg0: !FHE.eint<7>, %arg1: !FHE.eint<7>, %arg2: !FHE.eint<7>, %
     ASSERT_EXPECTED_VALUE(res_3, 60);
     ASSERT_EXPECTED_VALUE(res_4, 28);
   } else {
-    ASSERT_EXPECTED_FAILURE(lambda(1_u64, 2_u64, 3_u64, 4_u64));
-    ASSERT_EXPECTED_FAILURE(lambda(4_u64, 5_u64, 6_u64, 7_u64));
-    ASSERT_EXPECTED_FAILURE(lambda(1_u64, 1_u64, 1_u64, 1_u64));
-    ASSERT_EXPECTED_FAILURE(lambda(5_u64, 7_u64, 11_u64, 13_u64));
+    ASSERT_EXPECTED_FAILURE(lambda());
+    ASSERT_EXPECTED_FAILURE(lambda());
+    ASSERT_EXPECTED_FAILURE(lambda());
+    ASSERT_EXPECTED_FAILURE(lambda());
   }
 }
 
@@ -119,12 +120,13 @@ TEST(ParallelizeAndRunFHE, nn_small_parallel) {
     ASSERT_EQ(res->size(), dim0 * dim2);
     parallel_results = *res;
   } else {
-    ASSERT_EXPECTED_FAILURE(lambda.operator()<std::vector<uint64_t>>({&arg}));
+    ASSERT_EXPECTED_FAILURE(lambda.operator()<std::vector<uint64_t>>());
   }
 }
 
 TEST(ParallelizeAndRunFHE, nn_small_sequential) {
-  checkedJit(lambda, R"XXX(
+  if (mlir::concretelang::dfr::_dfr_is_root_node()) {
+    checkedJit(lambda, R"XXX(
   func @main(%arg0: tensor<4x5x!FHE.eint<5>>) -> tensor<4x7x!FHE.eint<5>> {
     %cst = arith.constant dense<[[0, 0, 1, 0, 1, 1, 0], [1, 1, 1, 0, 1, 0, 0], [1, 1, 0, 0, 0, 0, 0], [0, 0, 0, 0, 1, 1, 1]]> : tensor<4x7xi6>
     %cst_0 = arith.constant dense<[[1, 0, 1, 1, 0, 1, 1], [0, 1, 0, 0, 0, 0, 1], [0, 1, 1, 1, 1, 0, 0], [0, 1, 1, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 1]]> : tensor<5x7xi6>
@@ -135,30 +137,32 @@ TEST(ParallelizeAndRunFHE, nn_small_sequential) {
     return %2 : tensor<4x7x!FHE.eint<5>>
   }
 )XXX",
-             "main", false, false, false);
+               "main", false, false, false);
 
-  const size_t numDim = 2;
-  const size_t dim0 = 4;
-  const size_t dim1 = 5;
-  const size_t dim2 = 7;
-  const int64_t dims[numDim]{dim0, dim1};
-  const llvm::ArrayRef<int64_t> shape2D(dims, numDim);
-  std::vector<uint8_t> input;
-  input.reserve(dim0 * dim1);
+    const size_t numDim = 2;
+    const size_t dim0 = 4;
+    const size_t dim1 = 5;
+    const size_t dim2 = 7;
+    const int64_t dims[numDim]{dim0, dim1};
+    const llvm::ArrayRef<int64_t> shape2D(dims, numDim);
+    std::vector<uint8_t> input;
+    input.reserve(dim0 * dim1);
 
-  for (int i = 0; i < dim0 * dim1; ++i)
-    input.push_back(i % 17 % 4);
+    for (int i = 0; i < dim0 * dim1; ++i)
+      input.push_back(i % 17 % 4);
 
-  mlir::concretelang::TensorLambdaArgument<
-      mlir::concretelang::IntLambdaArgument<uint8_t>>
-      arg(input, shape2D);
+    mlir::concretelang::TensorLambdaArgument<
+        mlir::concretelang::IntLambdaArgument<uint8_t>>
+        arg(input, shape2D);
 
-  // This is sequential: only execute on root node.
-  if (mlir::concretelang::dfr::_dfr_is_root_node()) {
-    llvm::Expected<std::vector<uint64_t>> res =
-        lambda.operator()<std::vector<uint64_t>>({&arg});
-    ASSERT_EXPECTED_SUCCESS(res);
-    for (size_t i = 0; i < dim0 * dim2; i++)
-      EXPECT_EQ(parallel_results[i], (*res)[i]) << "result differ at pos " << i;
+    // This is sequential: only execute on root node.
+    if (mlir::concretelang::dfr::_dfr_is_root_node()) {
+      llvm::Expected<std::vector<uint64_t>> res =
+          lambda.operator()<std::vector<uint64_t>>({&arg});
+      ASSERT_EXPECTED_SUCCESS(res);
+      for (size_t i = 0; i < dim0 * dim2; i++)
+        EXPECT_EQ(parallel_results[i], (*res)[i])
+            << "result differ at pos " << i;
+    }
   }
 }

compiler/tests/end_to_end_tests/end_to_end_jit_clear_tensor.cc

@@ -1,4 +1,5 @@
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 ///////////////////////////////////////////////////////////////////////////////
 // 1D tensor //////////////////////////////////////////////////////////////////

compiler/tests/end_to_end_tests/end_to_end_jit_distributed.cc

@@ -5,6 +5,7 @@
 #include <type_traits>
 
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 ///////////////////////////////////////////////////////////////////////////////
 // Auto-parallelize independent FHE ops /////////////////////////////////////

compiler/tests/end_to_end_tests/end_to_end_jit_distributed.sh

@@ -4,7 +4,7 @@
 #SBATCH --mail-user=antoniu.pop@zama.ai
 #SBATCH --nodes=4
 #SBATCH --cpus-per-task=8
-#SBATCH --time=00:20:00
+#SBATCH --time=00:45:00
 #SBATCH --output=end_to_end_jit_distributed_%j.log
 
 echo "Date              = $(date)"

compiler/tests/end_to_end_tests/end_to_end_jit_encrypted_tensor.cc

@@ -1,4 +1,5 @@
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 TEST(End2EndJit_EncryptedTensor_2D, extract_slice_parametric_2x2) {
   checkedJit(lambda, R"XXX(

compiler/tests/end_to_end_tests/end_to_end_jit_fhe.cc

@@ -7,6 +7,7 @@
 #include "concretelang/Support/LibrarySupport.h"
 #include "end_to_end_fixture/EndToEndFixture.h"
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 #include "tests_tools/keySetCache.h"
 
 template <typename LambdaSupport>

compiler/tests/end_to_end_tests/end_to_end_jit_fhelinalg.cc

@@ -1,4 +1,5 @@
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 namespace Z = mlir::concretelang;
 template <class Elmt>

compiler/tests/end_to_end_tests/end_to_end_jit_lambda.cc

@@ -1,6 +1,7 @@
 #include <gtest/gtest.h>
 
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 TEST(Lambda_check_param, int_to_void_missing_param) {
   checkedJit(lambda, R"XXX(

compiler/tests/end_to_end_tests/end_to_end_jit_test.cc

@@ -4,6 +4,7 @@
 #include <type_traits>
 
 #include "end_to_end_jit_test.h"
+#include "tests_tools/GtestEnvironment.h"
 
 TEST(CompileAndRunClear, add_u64) {
   checkedJit(lambda, R"XXX(

compiler/tests/end_to_end_tests/globals.cc

@@ -1,3 +1,7 @@
 #include "globals.h"
+#include "tests_tools/GtestEnvironment.h"
 
 const mlir::concretelang::V0FHEConstraint defaultV0Constraints{10, 7};
+
+testing::Environment *const dfr_env =
+    testing::AddGlobalTestEnvironment(new DFREnvironment);

compiler/tests/tests_tools/GtestEnvironment.h

@@ -0,0 +1,19 @@
+#ifndef GTEST_ENVIRONMENT_H
+#define GTEST_ENVIRONMENT_H
+
+#include "concretelang/Runtime/DFRuntime.hpp"
+#include <gtest/gtest.h>
+
+class DFREnvironment : public ::testing::Environment {
+public:
+  ~DFREnvironment() override {}
+
+  // Override this to define how to set up the environment.
+  void SetUp() override {}
+
+  // Override this to define how to tear down the environment.
+  void TearDown() override { _dfr_terminate(); }
+};
+
+extern testing::Environment *const dfr_env;
+#endif

compiler/tests/unit_tests/concretelang/TestLib/testlib_unit_test.cpp

@@ -11,11 +11,15 @@
 #include "concretelang/Support/CompilerEngine.h"
 #include "concretelang/TestLib/TestTypedLambda.h"
 
+#include "tests_tools/GtestEnvironment.h"
 #include "tests_tools/assert.h"
 #include "tests_tools/keySetCache.h"
 
 #include "call_2t_1s_with_header-client.h.generated"
 
+testing::Environment *const dfr_env =
+    testing::AddGlobalTestEnvironment(new DFREnvironment);
+
 const std::string FUNCNAME = "main";
 
 using namespace concretelang::testlib;