[IMP] Limit the Test Data Size when Running CI in gcn_dist_sg.py (#…

…4461)

Running `gcn_dist_sg.py` can take 15-25 minutes in CI because of the size of the test dataset.  This PR cuts the size of that test dataset for CI runs, which results in nearly identical accuracy anyways and a runtime of less than 5 minutes.

Authors:
  - Alex Barghi (https://github.com/alexbarghi-nv)
  - Naim (https://github.com/naimnv)

Approvers:
  - Brad Rees (https://github.com/BradReesWork)

URL: #4461

python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py

@@ -15,6 +15,7 @@
 import argparse
 import tempfile
 import os
+import warnings
 
 from typing import Optional, Tuple, Dict
 
@@ -159,6 +160,10 @@ def parse_args():
         args.dataset, args.dataset_root
     )
 
+    if "CI_RUN" in os.environ and os.environ["CI_RUN"] == "1":
+        warnings.warn("Pruning test dataset for CI run.")
+        split_idx["test"] = split_idx["test"][:1000]
+
     with tempfile.TemporaryDirectory(dir=args.tempdir_root) as samples_dir:
         loader_kwargs = {
             "data": data,