From 70c236c3864f55992f05605cb5d77fda2312548b Mon Sep 17 00:00:00 2001
From: Alex Barghi <105237337+alexbarghi-nv@users.noreply.github.com>
Date: Mon, 1 Jul 2024 12:19:39 -0400
Subject: [PATCH] [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: https://github.com/rapidsai/cugraph/pull/4461
---
 python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py b/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py
index 82a612622a1..b299fc2a1a1 100644
--- a/python/cugraph-pyg/cugraph_pyg/examples/gcn_dist_sg.py
+++ b/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,