update shell scripts for dgl

rapidsai · Jan 9, 2024 · 3df7f82 · 3df7f82
1 parent 379f498
commit 3df7f82
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/benchmarks/cugraph/standalone/bulk_sampling/run_sampling.sh b/benchmarks/cugraph/standalone/bulk_sampling/run_sampling.sh
@@ -21,6 +21,7 @@ FANOUT=$2
 REPLICATION_FACTOR=$3
 SCRIPTS_DIR=$4
 NUM_EPOCHS=$5
+SAMPLING_FRAMEWORK=$6
 
 SAMPLES_DIR=/samples
 DATASET_DIR=/datasets
@@ -78,7 +79,8 @@ if [[ $SLURM_NODEID == 0 ]]; then
         --batch_sizes $BATCH_SIZE \
         --seeds_per_call_opts "524288" \
         --num_epochs $NUM_EPOCHS \
-        --random_seed 42 
+        --random_seed 42 \
+        --sampling_target_framework $SAMPLING_FRAMEWORK
 
     echo "DONE" > ${SAMPLES_DIR}/status.txt
 fi

diff --git a/benchmarks/cugraph/standalone/bulk_sampling/run_train_job.sh b/benchmarks/cugraph/standalone/bulk_sampling/run_train_job.sh
@@ -33,7 +33,7 @@ FANOUT="10_10_10"
 NUM_EPOCHS=1
 REPLICATION_FACTOR=1
 
-# options: PyG or cuGraphPyG
+# options: PyG, cuGraphPyG, or cuGraphDGL
 FRAMEWORK="cuGraphPyG"
 GPUS_PER_NODE=8
 
@@ -59,7 +59,12 @@ if [[ "$FRAMEWORK" == "cuGraphPyG" ]]; then
     srun \
         --container-image $CONTAINER_IMAGE \
         --container-mounts=${LOGS_DIR}":/logs",${SAMPLES_DIR}":/samples",${SCRIPTS_DIR}":/scripts",${DATASETS_DIR}":/datasets" \
-        bash /scripts/run_sampling.sh $BATCH_SIZE $FANOUT $REPLICATION_FACTOR "/scripts" $NUM_EPOCHS
+        bash /scripts/run_sampling.sh $BATCH_SIZE $FANOUT $REPLICATION_FACTOR "/scripts" $NUM_EPOCHS "cugraph_pyg"
+elif [[ "$FRAMEWORK" == "cuGraphDGL" ]]; then
+    srun \
+        --container-image $CONTAINER_IMAGE \
+        --container-mounts=${LOGS_DIR}":/logs",${SAMPLES_DIR}":/samples",${SCRIPTS_DIR}":/scripts",${DATASETS_DIR}":/datasets" \
+        bash /scripts/run_sampling.sh $BATCH_SIZE $FANOUT $REPLICATION_FACTOR "/scripts" $NUM_EPOCHS "cugraph_dgl_csr"
 fi
 
 # Train