Repository: pyg-team/pytorch_geometric
Pull request #6760 (2023-02-21 13:26:51)

Capture `stdout` in CI tests


Patch:
diff --git a/torch_geometric/contrib/explain/pgm_explainer.py b/torch_geometric/contrib/explain/pgm_explainer.py
index dfb23d772..ef0993b83 100644
--- a/torch_geometric/contrib/explain/pgm_explainer.py
+++ b/torch_geometric/contrib/explain/pgm_explainer.py
@@ -278,8 +278,6 @@ class PGMExplainer(ExplainerAlgorithm):
         import pandas as pd
         from pgmpy.estimators.CITests import chi_square
 
-        logging.info(f'Explaining node: {index}')
-
         neighbors, _, _, _ = k_hop_subgraph(
             node_idx=index,
             num_hops=get_num_hops(model),
diff --git a/torch_geometric/data/dataset.py b/torch_geometric/data/dataset.py
index ec06e1f73..5567c072c 100644
--- a/torch_geometric/data/dataset.py
+++ b/torch_geometric/data/dataset.py
@@ -217,7 +217,7 @@ class Dataset(torch.utils.data.Dataset, ABC):
         if files_exist(self.processed_paths):  # pragma: no cover
             return
 
-        if self.log:
+        if self.log and 'pytest' not in sys.modules:
             print('Processing...', file=sys.stderr)
 
         makedirs(self.processed_dir)
@@ -228,7 +228,7 @@ class Dataset(torch.utils.data.Dataset, ABC):
         path = osp.join(self.processed_dir, 'pre_filter.pt')
         torch.save(_repr(self.pre_filter), path)
 
-        if self.log:
+        if self.log and 'pytest' not in sys.modules:
             print('Done!', file=sys.stderr)
 
     def __len__(self) -> int:
diff --git a/torch_geometric/data/download.py b/torch_geometric/data/download.py
index 7284ba6d4..a350f4aa3 100644
--- a/torch_geometric/data/download.py
+++ b/torch_geometric/data/download.py
@@ -25,11 +25,11 @@ def download_url(url: str, folder: str, log: bool = True,
     path = osp.join(folder, filename)
 
     if osp.exists(path):  # pragma: no cover
-        if log:
+        if log and 'pytest' not in sys.modules:
             print(f'Using existing file {filename}', file=sys.stderr)
         return path
 
-    if log:
+    if log and 'pytest' not in sys.modules:
         print(f'Downloading {url}', file=sys.stderr)
 
     makedirs(folder)
diff --git a/torch_geometric/data/extract.py b/torch_geometric/data/extract.py
index 6be5652e0..da2c1d718 100644
--- a/torch_geometric/data/extract.py
+++ b/torch_geometric/data/extract.py
@@ -7,7 +7,7 @@ import zipfile
 
 
 def maybe_log(path, log=True):
-    if log:
+    if log and 'pytest' not in sys.modules:
         print(f'Extracting {path}', file=sys.stderr)
 
 
diff --git a/torch_geometric/loader/mixin.py b/torch_geometric/loader/mixin.py
index 948691839..9fc02046d 100644
--- a/torch_geometric/loader/mixin.py
+++ b/torch_geometric/loader/mixin.py
@@ -154,8 +154,8 @@ class AffinityMixin:
         try:
             # Set CPU affinity for dataloader:
             self.worker_init_fn = init_fn
-            logging.info(f"{self.num_workers} data loader workers are "
-                         f"assigned to CPUs {loader_cores}")
+            logging.debug(f"{self.num_workers} data loader workers are "
+                          f"assigned to CPUs {loader_cores}")
             yield
         finally:
             # Restore omp_num_threads and cpu affinity:
diff --git a/torch_geometric/nn/kge/base.py b/torch_geometric/nn/kge/base.py
index 3ffa73c38..9aff243ec 100644
--- a/torch_geometric/nn/kge/base.py
+++ b/torch_geometric/nn/kge/base.py
@@ -79,8 +79,15 @@ class KGEModel(torch.nn.Module):
         return KGTripletLoader(head_index, rel_type, tail_index, **kwargs)
 
     @torch.no_grad()
-    def test(self, head_index: Tensor, rel_type: Tensor, tail_index: Tensor,
-             batch_size: int, k: int = 10) -> Tuple[float, float]:
+    def test(
+        self,
+        head_index: Tensor,
+        rel_type: Tensor,
+        tail_index: Tensor,
+        batch_size: int,
+        k: int = 10,
+        log: bool = True,
+    ) -> Tuple[float, float]:
         r"""Evaluates the model quality by computing Mean Rank and
         Hits @ :math:`k` across all possible tail entities.
 
@@ -91,9 +98,14 @@ class KGEModel(torch.nn.Module):
             batch_size (int): The batch size to use for evaluating.
             k (int, optional): The :math:`k` in Hits @ :math:`k`.
                 (default: :obj:`10`)
+            log (bool, optional): If set to :obj:`False`, will not print a
+                progress bar to the console. (default: :obj:`True`)
         """
+        arange = range(head_index.numel())
+        arange = tqdm(arange) if log else arange
+
         mean_ranks, hits_at_k = [], []
-        for i in tqdm(range(head_index.numel())):
+        for i in arange:
             h, r, t = head_index[i], rel_type[i], tail_index[i]
 
             scores = []
diff --git a/torch_geometric/nn/model_hub.py b/torch_geometric/nn/model_hub.py
index 221e657c3..c98b2c8cf 100644
--- a/torch_geometric/nn/model_hub.py
+++ b/torch_geometric/nn/model_hub.py
@@ -148,7 +148,6 @@ class PyGModelHubMixin(ModelHubMixin):
         map_location = torch.device(map_location)
 
         if os.path.isdir(model_id):
-            print("Loading weights from local directory")
             model_file = os.path.join(model_id, MODEL_WEIGHTS_NAME)
         else:
             model_file = hf_hub_download(


Test patch:
diff --git a/test/contrib/explain/test_graphmask_explainer.py b/test/contrib/explain/test_graphmask_explainer.py
index 47b6631ec..ba5f2b622 100644
--- a/test/contrib/explain/test_graphmask_explainer.py
+++ b/test/contrib/explain/test_graphmask_explainer.py
@@ -54,7 +54,7 @@ def test_graph_mask_explainer(task_level):
 
     explainer = Explainer(
         model=model,
-        algorithm=GraphMaskExplainer(2, epochs=5),
+        algorithm=GraphMaskExplainer(2, epochs=5, log=False),
         explanation_type='model',
         node_mask_type='attributes',
         edge_mask_type='object',
diff --git a/test/graphgym/test_graphgym.py b/test/graphgym/test_graphgym.py
index 79e631bc6..d8bffc78b 100644
--- a/test/graphgym/test_graphgym.py
+++ b/test/graphgym/test_graphgym.py
@@ -49,7 +49,8 @@ def trivial_metric(true, pred, task_type):
 @pytest.mark.parametrize('auto_resume', [True, False])
 @pytest.mark.parametrize('skip_train_eval', [True, False])
 @pytest.mark.parametrize('use_trivial_metric', [True, False])
-def test_run_single_graphgym(auto_resume, skip_train_eval, use_trivial_metric):
+def test_run_single_graphgym(capfd, auto_resume, skip_train_eval,
+                             use_trivial_metric):
     warnings.filterwarnings('ignore', ".*does not have many workers.*")
     warnings.filterwarnings('ignore', ".*lower value for log_every_n_steps.*")
 
@@ -110,6 +111,12 @@ def test_run_single_graphgym(auto_resume, skip_train_eval, use_trivial_metric):
 
     shutil.rmtree(cfg.out_dir)
 
+    out, _ = capfd.readouterr()
+    assert "train: {'epoch': 0," in out
+    assert "val: {'epoch': 0," in out
+    assert "train: {'epoch': 5," in out
+    assert "val: {'epoch': 5," in out
+
 
 @withPackage('yacs')
 @withPackage('pytorch_lightning')
@@ -169,7 +176,7 @@ def test_graphgym_module(tmpdir):
 
 @withPackage('yacs')
 @withPackage('pytorch_lightning')
-def test_train(tmpdir):
+def test_train(capfd, tmpdir):
     warnings.filterwarnings('ignore', ".*does not have many workers.*")
 
     import pytorch_lightning as pl
@@ -197,3 +204,7 @@ def test_train(tmpdir):
     trainer.fit(model, train_loader, val_loader)
 
     shutil.rmtree(cfg.out_dir)
+
+    out, err = capfd.readouterr()
+    assert 'Sanity Checking' in out
+    assert 'Epoch 0:' in out
diff --git a/test/nn/kge/test_transe.py b/test/nn/kge/test_transe.py
index 88b7f9c97..69c97faa8 100644
--- a/test/nn/kge/test_transe.py
+++ b/test/nn/kge/test_transe.py
@@ -19,6 +19,6 @@ def test_transe():
         loss = model.loss(h, r, t)
         assert loss >= 0.
 
-        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5)
+        mean_rank, hits_at_10 = model.test(h, r, t, batch_size=5, log=False)
         assert mean_rank <= 10
         assert hits_at_10 == 1.0
diff --git a/test/nn/pool/test_decimation.py b/test/nn/pool/test_decimation.py
index e5f3537ed..7e2809a1a 100644
--- a/test/nn/pool/test_decimation.py
+++ b/test/nn/pool/test_decimation.py
@@ -9,7 +9,6 @@ def test_decimation_basic():
     ptr = torch.tensor([0, N_1, N_1 + N_2])
 
     idx_decim, ptr_decim = decimation_indices(ptr, decimation_factor)
-    print(idx_decim, ptr_decim)
 
     expected_size = (N_1 // decimation_factor) + (N_2 // decimation_factor)
     assert idx_decim.size(0) == expected_size
diff --git a/test/profile/test_benchmark.py b/test/profile/test_benchmark.py
index 3d68f60e9..977364f53 100644
--- a/test/profile/test_benchmark.py
+++ b/test/profile/test_benchmark.py
@@ -5,7 +5,7 @@ from torch_geometric.testing import withPackage
 
 
 @withPackage('tabulate')
-def test_benchmark():
+def test_benchmark(capfd):
     def add(x, y):
         return x + y
 
@@ -16,3 +16,7 @@ def test_benchmark():
         num_warmups=1,
         backward=True,
     )
+
+    out, _ = capfd.readouterr()
+    assert '| Name   | Forward   | Backward   | Total   |' in out
+    assert '| add    |' in out
diff --git a/test/profile/test_profile.py b/test/profile/test_profile.py
index 8613afe15..7d81a6f9b 100644
--- a/test/profile/test_profile.py
+++ b/test/profile/test_profile.py
@@ -77,7 +77,7 @@ def test_profileit(get_dataset):
 
 
 @withCUDA
-def test_torch_profile(get_dataset, device):
+def test_torch_profile(capfd, get_dataset, device):
     dataset = get_dataset(name='Cora')
     data = dataset[0].to(device)
     model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
@@ -86,6 +86,11 @@ def test_torch_profile(get_dataset, device):
     with torch_profile():
         model(data.x, data.edge_index)
 
+    out, _ = capfd.readouterr()
+    assert 'Self CPU time total' in out
+    if data.x.is_cuda:
+        assert 'Self CUDA time total' in out
+
     rename_profile_file('test_profile')
     assert os.path.exists('profile-test_profile.json')
     os.remove('profile-test_profile.json')
diff --git a/test/profile/test_profiler.py b/test/profile/test_profiler.py
index 22431c436..61acb3d4d 100644
--- a/test/profile/test_profiler.py
+++ b/test/profile/test_profiler.py
@@ -6,7 +6,7 @@ from torch_geometric.testing import withCUDA
 
 
 @withCUDA
-def test_profiler(get_dataset, device):
+def test_profiler(capfd, get_dataset, device):
     x = torch.randn(10, 16, device=device)
     edge_index = torch.tensor([
         [0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9],
@@ -18,6 +18,9 @@ def test_profiler(get_dataset, device):
     with Profiler(model, profile_memory=True, use_cuda=x.is_cuda) as prof:
         model(x, edge_index)
 
+    _, err = capfd.readouterr()
+    assert 'Completed Stage' in err
+
     _, heading_list, raw_results, layer_names, layer_stats = prof.get_trace()
     assert 'Self CPU total' in heading_list
     assert 'aten::relu' in raw_results