From 950c831d4562e8c835bf76051282bff09b2b5d08 Mon Sep 17 00:00:00 2001
From: Matthijs Douze <matthijs@fb.com>
Date: Sun, 17 Jan 2021 13:40:10 -0800
Subject: [PATCH] Fix number of threads in test_ivfpq_codec when running in
 sandcastle

Summary: When running in a heavily parallelized env, the test becomes very slow and causes timeouts. Here we reduce the nb of threads.

Reviewed By: wickedfoo, Ben0mega

Differential Revision: D25921771

fbshipit-source-id: 1e0aacbb3e4f6e8f33ec893984b343eb5a610424
---
 tests/test_ivfpq_codec.cpp | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/test_ivfpq_codec.cpp b/tests/test_ivfpq_codec.cpp
index 2732b6787..b74f4fd98 100644
--- a/tests/test_ivfpq_codec.cpp
+++ b/tests/test_ivfpq_codec.cpp
@@ -9,6 +9,8 @@
 #include <cstdlib>
 #include <random>
 
+#include <omp.h>
+
 #include <gtest/gtest.h>
 
 #include <faiss/IndexIVFPQ.h>
@@ -49,6 +51,22 @@ double eval_codec_error (long ncentroids, long m, const std::vector<float> &v)
 }  // namespace
 
 
+
+bool runs_on_sandcastle() {
+    // see discussion here https://fburl.com/qc5kpdo2
+    const char * sandcastle = getenv("SANDCASTLE");
+    if (sandcastle && !strcmp(sandcastle, "1")) {
+        return true;
+    }
+    const char * tw_job_user = getenv("TW_JOB_USER");
+    if (tw_job_user && !strcmp(tw_job_user, "sandcastle")) {
+        return true;
+    }
+
+    return false;
+}
+
+
 TEST(IVFPQ, codec) {
 
     std::vector <float> database (nb * d);
@@ -58,6 +76,11 @@ TEST(IVFPQ, codec) {
         database[i] = distrib(rng);
     }
 
+    // limit number of threads when running on heavily parallelized test environment
+    if (runs_on_sandcastle()) {
+        omp_set_num_threads(2);
+    }
+
     double err0 = eval_codec_error(16, 8, database);
 
     // should be more accurate as there are more coarse centroids