From 9b00c92182e07013c47d5052960f3f6dc586221f Mon Sep 17 00:00:00 2001
From: Arthur Zucker <arthur.zucker@gmail.com>
Date: Mon, 29 Apr 2024 13:21:32 +0200
Subject: [PATCH] [push-ci-image]

---
 .circleci/parse_test_outputs.py   | 6 +++---
 docker/examples-torch.dockerfile  | 8 ++++++--
 docker/pipeline-torch.dockerfile  | 8 ++++++--
 docker/torch-jax-light.dockerfile | 8 ++++++--
 docker/torch-light.dockerfile     | 8 ++++++--
 5 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/.circleci/parse_test_outputs.py b/.circleci/parse_test_outputs.py
index 758572bf310956..9d24a986742f06 100644
--- a/.circleci/parse_test_outputs.py
+++ b/.circleci/parse_test_outputs.py
@@ -11,9 +11,9 @@ def parse_pytest_output(file_path):
                 skipped_count += 1
                 test_file, test_line, reason = match.groups()
                 skipped_tests[reason] = skipped_tests.get(reason, []) + [(test_file, test_line)]
-    print("Number of skipped tests:", skipped_count)
     for k,v in sorted(skipped_tests.items(), key=lambda x:len(x[1])):
         print(f"{len(v):4} skipped because: {k}")
+    print("Number of skipped tests:", skipped_count)
 
 def parse_pytest_failure_output(file_path):
     failed_tests = {}
@@ -25,9 +25,9 @@ def parse_pytest_failure_output(file_path):
                 failed_count += 1
                 _, error, reason = match.groups()
                 failed_tests[reason] = failed_tests.get(reason, []) + [error]
-    print("Number of failed tests:", failed_count)
     for k,v in sorted(failed_tests.items(), key=lambda x:len(x[1])):
         print(f"{len(v):4} failed because `{v[0]}` -> {k}")
+    print("Number of failed tests:", failed_count)
     if failed_count>0:
         exit(1)
 
@@ -42,9 +42,9 @@ def parse_pytest_errors_output(file_path):
                 error_count += 1
                 _, test_error, reason = match.groups()
                 error_tests[reason] = error_tests.get(reason, []) + [test_error]
-    print("Number of errors:", error_count)
     for k,v in sorted(error_tests.items(), key=lambda x:len(x[1])):
         print(f"{len(v):4} errored out because of `{v[0]}` -> {k}")
+    print("Number of errors:", error_count)
     if error_count>0:
         exit(1)
 
diff --git a/docker/examples-torch.dockerfile b/docker/examples-torch.dockerfile
index f2ba6e01de281e..752edfee3a2d9d 100644
--- a/docker/examples-torch.dockerfile
+++ b/docker/examples-torch.dockerfile
@@ -14,6 +14,10 @@ RUN uv pip install --no-cache-dir "transformers[sklearn,torch,testing,sentencepi
 
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-RUN apt-get autoremove  --purge -y cmake
 RUN pip cache remove "nvidia-*"
-RUN pip cache remove triton
\ No newline at end of file
+RUN pip uninstall -y `pip freeze | grep "nvidia-*"` || true
+RUN pip cache remove triton
+RUN apt-get --purge remove "*nvidia*" || true
+RUN apt-get --purge remove "*cublas*" "cuda*" "nsight*"  || true
+RUN apt-get autoremove
+RUN apt-get autoclean
\ No newline at end of file
diff --git a/docker/pipeline-torch.dockerfile b/docker/pipeline-torch.dockerfile
index f2ba6e01de281e..752edfee3a2d9d 100644
--- a/docker/pipeline-torch.dockerfile
+++ b/docker/pipeline-torch.dockerfile
@@ -14,6 +14,10 @@ RUN uv pip install --no-cache-dir "transformers[sklearn,torch,testing,sentencepi
 
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-RUN apt-get autoremove  --purge -y cmake
 RUN pip cache remove "nvidia-*"
-RUN pip cache remove triton
\ No newline at end of file
+RUN pip uninstall -y `pip freeze | grep "nvidia-*"` || true
+RUN pip cache remove triton
+RUN apt-get --purge remove "*nvidia*" || true
+RUN apt-get --purge remove "*cublas*" "cuda*" "nsight*"  || true
+RUN apt-get autoremove
+RUN apt-get autoclean
\ No newline at end of file
diff --git a/docker/torch-jax-light.dockerfile b/docker/torch-jax-light.dockerfile
index 088c2f529926c3..cfa56ca4f98b2c 100644
--- a/docker/torch-jax-light.dockerfile
+++ b/docker/torch-jax-light.dockerfile
@@ -15,6 +15,10 @@ RUN uv pip install --no-cache-dir "transformers[sklearn,flax,torch,testing,sente
 
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-RUN apt-get autoremove  --purge -y cmake
 RUN pip cache remove "nvidia-*"
-RUN pip cache remove triton
\ No newline at end of file
+RUN pip uninstall -y `pip freeze | grep "nvidia-*"` || true
+RUN pip cache remove triton
+RUN apt-get --purge remove "*nvidia*" || true
+RUN apt-get --purge remove "*cublas*" "cuda*" "nsight*"  || true
+RUN apt-get autoremove
+RUN apt-get autoclean
\ No newline at end of file
diff --git a/docker/torch-light.dockerfile b/docker/torch-light.dockerfile
index cc72e5ad60f4ff..63b5bf4c552b60 100644
--- a/docker/torch-light.dockerfile
+++ b/docker/torch-light.dockerfile
@@ -16,6 +16,10 @@ RUN uv pip install --no-cache-dir "transformers[sklearn,sentencepiece,vision,tim
 
 RUN pip uninstall -y transformers
 RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-RUN apt-get autoremove  --purge -y cmake
 RUN pip cache remove "nvidia-*"
-RUN pip cache remove triton
\ No newline at end of file
+RUN pip uninstall -y `pip freeze | grep "nvidia-*"` || true
+RUN pip cache remove triton
+RUN apt-get --purge remove "*nvidia*" || true
+RUN apt-get --purge remove "*cublas*" "cuda*" "nsight*"  || true
+RUN apt-get autoremove
+RUN apt-get autoclean
\ No newline at end of file