[CI] [GHA] Use OV_CACHE in the WWB tests

tools/who_what_benchmark/tests/constants.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+import os
+import tempfile
+
+
+WWB_CACHE_PATH = Path(os.path.join(os.environ.get('OV_CACHE', tempfile.TemporaryDirectory()), 'wwb_cache'))
+SHOULD_CLEANUP = bool(os.environ.get('CLEANUP_CACHE', None))

tools/who_what_benchmark/tests/test_cli_image.py

@@ -6,12 +6,13 @@
 import logging
 import tempfile
 import re
+from constants import WWB_CACHE_PATH, SHOULD_CLEANUP
 
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-MODEL_CACHE = tempfile.mkdtemp()
+MODEL_CACHE = WWB_CACHE_PATH
 OV_IMAGE_MODELS = ["echarlaix/tiny-random-stable-diffusion-xl",
                    "yujiepan/stable-diffusion-3-tiny-random",
                    "katuni4ka/tiny-random-flux",
@@ -42,8 +43,9 @@ def setup_module():
 
 
 def teardown_module():
-    logger.info("Remove models")
-    shutil.rmtree(MODEL_CACHE)
+    if SHOULD_CLEANUP:
+        logger.info("Removing models")
+        shutil.rmtree(MODEL_CACHE)
 
 
 def get_similarity(output: str) -> float:
@@ -67,11 +69,12 @@ def get_similarity(output: str) -> float:
     ],
 )
 def test_image_model_types(model_id, model_type, backend, tmp_path):
+    MODEL_PATH = os.path.join(MODEL_CACHE, model_id.replace("/", "--"))
     wwb_args = [
         "--base-model",
-        model_id,
+        MODEL_PATH,
         "--target-model",
-        model_id,
+        MODEL_PATH,
         "--num-samples",
         "1",
         "--gt-data",
@@ -195,9 +198,10 @@ def test_image_model_genai(model_id, model_type, tmp_path):
 )
 def test_image_custom_dataset(model_id, model_type, backend, tmp_path):
     GT_FILE = tmp_path / "test_sd.csv"
+    MODEL_PATH = os.path.join(MODEL_CACHE, model_id.replace("/", "--"))
     wwb_args = [
         "--base-model",
-        model_id,
+        MODEL_PATH,
         "--num-samples",
         "1",
         "--gt-data",

tools/who_what_benchmark/tests/test_cli_text.py

@@ -1,12 +1,13 @@
 import os
 import shutil
-import tempfile
 import pandas as pd
 import pytest
 import logging
 import json
 import sys
 
+from constants import WWB_CACHE_PATH, SHOULD_CLEANUP
+
 from transformers import AutoTokenizer
 from optimum.intel.openvino import OVModelForCausalLM, OVWeightQuantizationConfig
 
@@ -18,9 +19,9 @@
 
 
 model_id = "facebook/opt-125m"
-tmp_dir = tempfile.mkdtemp()
-base_model_path = os.path.join(tmp_dir, "opt125m")
-target_model_path = os.path.join(tmp_dir, "opt125m_int8")
+cache_dir = WWB_CACHE_PATH
+base_model_path = os.path.join(cache_dir, "opt125m")
+target_model_path = os.path.join(cache_dir, "opt125m_int8")
 
 gptq_model_id = "ybelkada/opt-125m-gptq-4bit"
 awq_model_id = "TitanML/tiny-mixtral-AWQ-4bit"
@@ -30,23 +31,26 @@ def setup_module():
     from optimum.exporters.openvino.convert import export_tokenizer
 
     logger.info("Create models")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    base_model = OVModelForCausalLM.from_pretrained(model_id)
-    base_model.save_pretrained(base_model_path)
-    tokenizer.save_pretrained(base_model_path)
-    export_tokenizer(tokenizer, base_model_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=WWB_CACHE_PATH)
+    base_model = OVModelForCausalLM.from_pretrained(model_id, cache_dir=WWB_CACHE_PATH)
+    if not os.path.exists(base_model_path):
+        base_model.save_pretrained(base_model_path)
+        tokenizer.save_pretrained(base_model_path)
+        export_tokenizer(tokenizer, base_model_path)
 
     target_model = OVModelForCausalLM.from_pretrained(
-        model_id, quantization_config=OVWeightQuantizationConfig(bits=8)
+        model_id, quantization_config=OVWeightQuantizationConfig(bits=8), cache_dir=WWB_CACHE_PATH
     )
-    target_model.save_pretrained(target_model_path)
-    tokenizer.save_pretrained(target_model_path)
-    export_tokenizer(tokenizer, target_model_path)
+    if not os.path.exists(target_model_path):
+        target_model.save_pretrained(target_model_path)
+        tokenizer.save_pretrained(target_model_path)
+        export_tokenizer(tokenizer, target_model_path)
 
 
 def teardown_module():
-    logger.info("Remove models")
-    shutil.rmtree(tmp_dir)
+    if SHOULD_CLEANUP:
+        logger.info("Removing models")
+        shutil.rmtree(cache_dir)
 
 
 def test_text_target_model():
@@ -138,9 +142,10 @@ def test_text_verbose():
 
 def test_text_language(tmp_path):
     temp_file_name = tmp_path / "gt.csv"
+    MODEL_PATH = os.path.join(WWB_CACHE_PATH, 'Qwen/Qwen2-0.5B')
     run_wwb([
         "--base-model",
-        "Qwen/Qwen2-0.5B",
+        MODEL_PATH,
         "--gt-data",
         temp_file_name,
         "--num-samples",
@@ -170,9 +175,10 @@ def test_text_language(tmp_path):
 )
 def test_text_hf_model(model_id, tmp_path):
     temp_file_name = tmp_path / "gt.csv"
+    MODEL_PATH = os.path.join(WWB_CACHE_PATH, model_id.replace("/", "--"))
     run_wwb([
         "--base-model",
-        model_id,
+        MODEL_PATH,
         "--gt-data",
         temp_file_name,
         "--num-samples",

tools/who_what_benchmark/tests/test_cli_vlm.py

@@ -2,6 +2,7 @@
 import pytest
 import logging
 from test_cli_image import run_wwb
+from constants import WWB_CACHE_PATH
 
 
 logging.basicConfig(level=logging.INFO)
@@ -14,9 +15,9 @@
         ("katuni4ka/tiny-random-llava", "visual-text"),
     ],
 )
-def test_vlm_basic(model_id, model_type, tmp_path):
-    GT_FILE = tmp_path / "gt.csv"
-    MODEL_PATH = tmp_path / model_id.replace("/", "--")
+def test_vlm_basic(model_id, model_type):
+    GT_FILE = WWB_CACHE_PATH / "gt.csv"
+    MODEL_PATH = WWB_CACHE_PATH / model_id.replace("/", "--")
 
     result = subprocess.run(["optimum-cli", "export",
                              "openvino", "-m", model_id,
@@ -71,13 +72,13 @@ def test_vlm_basic(model_id, model_type, tmp_path):
         model_type,
         "--genai",
         "--output",
-        tmp_path,
+        WWB_CACHE_PATH,
     ])
 
     # test w/o models
     run_wwb([
         "--target-data",
-        tmp_path / "target.csv",
+        WWB_CACHE_PATH / "target.csv",
         "--num-samples",
         "1",
         "--gt-data",