From d0d1632958c7d543e07afc672a8501d704e5a65f Mon Sep 17 00:00:00 2001
From: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
Date: Fri, 9 Jun 2023 16:49:02 +0200
Subject: [PATCH] Fix Pipeline CI OOM issue (#24124)

* fix

* fix

* fix

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
---
 tests/pipelines/test_pipelines_common.py      | 15 +++++++++++++++
 .../test_pipelines_conversational.py          | 19 ++++++++++++++++++-
 tests/pipelines/test_pipelines_fill_mask.py   | 11 +++++++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py
index 1a3ddace28..7909ad22d6 100644
--- a/tests/pipelines/test_pipelines_common.py
+++ b/tests/pipelines/test_pipelines_common.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import logging
 import os
 import sys
@@ -507,6 +508,10 @@ class PipelineUtilsTest(unittest.TestCase):
 
             self.check_default_pipeline(task, "pt", set_seed_fn, self.check_models_equal_pt)
 
+            # clean-up as much as possible GPU memory occupied by PyTorch
+            gc.collect()
+            torch.cuda.empty_cache()
+
     @slow
     @require_tf
     def test_load_default_pipelines_tf(self):
@@ -522,6 +527,9 @@ class PipelineUtilsTest(unittest.TestCase):
 
             self.check_default_pipeline(task, "tf", set_seed_fn, self.check_models_equal_tf)
 
+            # clean-up as much as possible GPU memory occupied by PyTorch
+            gc.collect()
+
     @slow
     @require_torch
     def test_load_default_pipelines_pt_table_qa(self):
@@ -530,6 +538,10 @@ class PipelineUtilsTest(unittest.TestCase):
         set_seed_fn = lambda: torch.manual_seed(0)  # noqa: E731
         self.check_default_pipeline("table-question-answering", "pt", set_seed_fn, self.check_models_equal_pt)
 
+        # clean-up as much as possible GPU memory occupied by PyTorch
+        gc.collect()
+        torch.cuda.empty_cache()
+
     @slow
     @require_tf
     @require_tensorflow_probability
@@ -539,6 +551,9 @@ class PipelineUtilsTest(unittest.TestCase):
         set_seed_fn = lambda: tf.random.set_seed(0)  # noqa: E731
         self.check_default_pipeline("table-question-answering", "tf", set_seed_fn, self.check_models_equal_tf)
 
+        # clean-up as much as possible GPU memory occupied by PyTorch
+        gc.collect()
+
     def check_default_pipeline(self, task, framework, set_seed_fn, check_models_equal_fn):
         from transformers.pipelines import SUPPORTED_TASKS, pipeline
 
diff --git a/tests/pipelines/test_pipelines_conversational.py b/tests/pipelines/test_pipelines_conversational.py
index 70ecbc3104..f627150d80 100644
--- a/tests/pipelines/test_pipelines_conversational.py
+++ b/tests/pipelines/test_pipelines_conversational.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import unittest
 
 from transformers import (
@@ -29,7 +30,14 @@ from transformers import (
     TFAutoModelForCausalLM,
     pipeline,
 )
-from transformers.testing_utils import is_pipeline_test, require_tf, require_torch, slow, torch_device
+from transformers.testing_utils import (
+    is_pipeline_test,
+    is_torch_available,
+    require_tf,
+    require_torch,
+    slow,
+    torch_device,
+)
 
 from .test_pipelines_common import ANY
 
@@ -39,6 +47,15 @@ DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0
 
 @is_pipeline_test
 class ConversationalPipelineTests(unittest.TestCase):
+    def tearDown(self):
+        super().tearDown()
+        # clean-up as much as possible GPU memory occupied by PyTorch
+        gc.collect()
+        if is_torch_available():
+            import torch
+
+            torch.cuda.empty_cache()
+
     model_mapping = dict(
         list(MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items())
         if MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
diff --git a/tests/pipelines/test_pipelines_fill_mask.py b/tests/pipelines/test_pipelines_fill_mask.py
index d6dec8db17..4c53a905ed 100644
--- a/tests/pipelines/test_pipelines_fill_mask.py
+++ b/tests/pipelines/test_pipelines_fill_mask.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import unittest
 
 from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
 from transformers.pipelines import PipelineException
 from transformers.testing_utils import (
     is_pipeline_test,
+    is_torch_available,
     nested_simplify,
     require_tf,
     require_torch,
@@ -33,6 +35,15 @@ class FillMaskPipelineTests(unittest.TestCase):
     model_mapping = MODEL_FOR_MASKED_LM_MAPPING
     tf_model_mapping = TF_MODEL_FOR_MASKED_LM_MAPPING
 
+    def tearDown(self):
+        super().tearDown()
+        # clean-up as much as possible GPU memory occupied by PyTorch
+        gc.collect()
+        if is_torch_available():
+            import torch
+
+            torch.cuda.empty_cache()
+
     @require_tf
     def test_small_model_tf(self):
         unmasker = pipeline(task="fill-mask", model="sshleifer/tiny-distilroberta-base", top_k=2, framework="tf")