minor change to avoid timeout

rapidsai · Dec 29, 2023 · c2e8520 · c2e8520
1 parent 7222cba
commit c2e8520
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/benchmarks/cugraph/standalone/bulk_sampling/bench_cugraph_training.py b/benchmarks/cugraph/standalone/bulk_sampling/bench_cugraph_training.py
@@ -11,12 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import os
-os.environ['RAPIDS_NO_INITIALIZE'] = 1
-os.environ['CUDF_SPILL'] = 1
+os.environ['RAPIDS_NO_INITIALIZE'] = "1"
+os.environ['CUDF_SPILL'] = "1"
 os.environ['LIBCUDF_CUFILE_POLICY'] = "KVIKIO"
-os.environ['KVIKIO_NTHREADS'] = 64
+os.environ['KVIKIO_NTHREADS'] = "64"
 
 import argparse
 import json

diff --git a/benchmarks/cugraph/standalone/bulk_sampling/run_train_job.sh b/benchmarks/cugraph/standalone/bulk_sampling/run_train_job.sh
@@ -4,7 +4,7 @@
 #SBATCH -p luna
 #SBATCH -J datascience_rapids_cugraphgnn-papers:bulkSamplingPyG
 #SBATCH -N 1
-#SBATCH -t 00:22:00 
+#SBATCH -t 00:25:00 
 
 CONTAINER_IMAGE="/lustre/fsw/rapids/abarghi/dlfw_patched.squash"
 SCRIPTS_DIR=$(pwd)