From f61ed4730b36ec2272a00776c0fc02180266e7c1 Mon Sep 17 00:00:00 2001
From: baptiste <collebaptiste@gmail.com>
Date: Mon, 14 Oct 2024 10:03:58 +0000
Subject: [PATCH] fix bugs

---
 .github/workflows/test_cli_cuda_pytorch.yaml | 21 +++++++++---------
 .github/workflows/test_cli_rocm_pytorch.yaml | 23 ++++++++++----------
 examples/tei_bge.yaml                        |  4 ----
 tests/test_examples.py                       | 10 +++++++++
 4 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/.github/workflows/test_cli_cuda_pytorch.yaml b/.github/workflows/test_cli_cuda_pytorch.yaml
index 588e72f7..4cb0060a 100644
--- a/.github/workflows/test_cli_cuda_pytorch.yaml
+++ b/.github/workflows/test_cli_cuda_pytorch.yaml
@@ -47,10 +47,14 @@ jobs:
         run: |
           pip install -e .[testing,diffusers,timm,peft,bitsandbytes,autoawq,auto-gptq]
 
-      - name: Run tests
+      - name: Run cli tests
         run: |
           pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and not (dp or ddp or device_map or deepspeed)"
 
+      - name: Run tests from example folder
+        run: |
+          pytest tests/test_examples.py -x -s -k "cuda and pytorch"
+
   run_cli_cuda_pytorch_multi_gpu_tests:
     if: ${{
       (github.event_name == 'push') ||
@@ -77,14 +81,11 @@ jobs:
         run: |
           pip install -e .[testing,diffusers,timm,peft,deepspeed]
 
-      # - name: Run cli tests (parallel)
-      #   run: |
-      #     pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (dp or ddp or device_map)"
-
-      # - name: Run cli tests (sequential)
-      #   run: |
-      #     FORCE_SEQUENTIAL=1 pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (deepspeed_inference)"
+      - name: Run cli tests (parallel)
+        run: |
+          pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (dp or ddp or device_map)"
 
-      - name: Run test from example folder
+      - name: Run cli tests (sequential)
         run: |
-          pytest tests/test_examples.py -x -s -k "cuda and pytorch"
\ No newline at end of file
+          FORCE_SEQUENTIAL=1 pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (deepspeed_inference)"
+
diff --git a/.github/workflows/test_cli_rocm_pytorch.yaml b/.github/workflows/test_cli_rocm_pytorch.yaml
index bc0dca66..8da670d7 100644
--- a/.github/workflows/test_cli_rocm_pytorch.yaml
+++ b/.github/workflows/test_cli_rocm_pytorch.yaml
@@ -50,9 +50,13 @@ jobs:
         run: |
           pip install -e .[testing,diffusers,timm,peft,autoawq,auto-gptq]
 
-      - name: Run tests
+      # - name: Run cli tests
+      #   run: |
+      #     pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and not (dp or ddp or device_map or deepspeed) and not bnb"
+
+      - name: Run tests from example folder
         run: |
-          pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and not (dp or ddp or device_map or deepspeed) and not bnb"
+          pytest tests/test_examples.py -x -s -k "rocm and pytorch"
 
   run_cli_rocm_pytorch_multi_gpu_tests:
     if: ${{
@@ -84,14 +88,11 @@ jobs:
         run: |
           pip install -e .[testing,diffusers,timm,deepspeed,peft,autoawq,auto-gptq] "deepspeed<0.15"
 
-      # - name: Run cli tests (parallel)
-      #   run: |
-      #     pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (dp or ddp or device_map)"
-
-      # - name: Run cli tests (sequential)
-      #   run: |
-      #     FORCE_SEQUENTIAL=1 pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (deepspeed_inference)"
+      - name: Run cli tests (parallel)
+        run: |
+          pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (dp or ddp or device_map)"
 
-      - name: Run tests from example folder
+      - name: Run cli tests (sequential)
         run: |
-          pytest tests/test_examples.py -x -s -k "cuda and pytorch"
\ No newline at end of file
+          FORCE_SEQUENTIAL=1 pytest tests/test_cli.py -x -s -k "cli and cuda and pytorch and (deepspeed_inference)"
+
diff --git a/examples/tei_bge.yaml b/examples/tei_bge.yaml
index dbbab7d5..bbae52d6 100644
--- a/examples/tei_bge.yaml
+++ b/examples/tei_bge.yaml
@@ -7,10 +7,6 @@ defaults:
 
 name: tei_bert
 
-launcher:
-  device_isolation: true
-  device_isolation_action: warn
-
 backend:
   device: cpu
   model: BAAI/bge-base-en-v1.5
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 0f3fc7a8..0194a06d 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -207,3 +207,13 @@ def test_cuda_vllm_configs(config_name):
 @pytest.mark.parametrize("script_name", CUDA_PYTORCH_SCRIPTS)
 def test_cuda_pytorch_scripts(script_name):
     execute_python_script(script_name)
+
+ROCM_PYTORCH_CONFIGS = [config for config in CUDA_PYTORCH_CONFIGS if config != "energy_star.yaml"]
+
+@pytest.mark.parametrize("config_name", ROCM_PYTORCH_CONFIGS)
+def test_rocm_pytorch_configs(config_name):
+    test_yaml_config(config_name)
+
+@pytest.mark.parametrize("script_name", CUDA_PYTORCH_SCRIPTS)
+def test_rocm_pytorch_scripts(script_name):
+    execute_python_script(script_name)