From abe9ce211b2676f8a6a2623f5767aa5f20e406cd Mon Sep 17 00:00:00 2001
From: Harutaka Kawamura <hkawamura0130@gmail.com>
Date: Fri, 15 Mar 2024 15:28:23 +0900
Subject: [PATCH] Add pre-test step in cross version tests (#11430)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
---
 .github/workflows/cross-version-tests.yml | 6 +++---
 dev/set_matrix.py                         | 6 +++++-
 mlflow/ml-package-versions.yml            | 2 ++
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/cross-version-tests.yml b/.github/workflows/cross-version-tests.yml
index 37b5e7b05f79e..ebdaaeb5eb337 100644
--- a/.github/workflows/cross-version-tests.yml
+++ b/.github/workflows/cross-version-tests.yml
@@ -157,10 +157,10 @@ jobs:
       - name: Check package versions
         run: |
           python dev/show_package_release_dates.py
-      - name: Prefetch Transformer models
-        if: matrix.package == 'transformers' && matrix.category == 'models'
+      - name: Pre-test
+        if: matrix.pre_test
         run: |
-          python tests/transformers/helper.py
+          ${{ matrix.pre_test }}
       - uses: ./.github/actions/pipdeptree
       - name: Run tests
         env:
diff --git a/dev/set_matrix.py b/dev/set_matrix.py
index 0b57b4af407ba..7c6885c5e195d 100644
--- a/dev/set_matrix.py
+++ b/dev/set_matrix.py
@@ -73,6 +73,7 @@ class TestConfig(BaseModel):
     java: t.Optional[t.Dict[str, str]] = None
     run: str
     allow_unreleased_max_version: t.Optional[bool] = None
+    pre_test: t.Optional[str] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -102,6 +103,7 @@ class MatrixItem(BaseModel):
     java: str
     supported: bool
     free_disk_space: bool
+    pre_test: t.Optional[str] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -379,6 +381,7 @@ def expand_config(config):
                         java=java,
                         supported=ver <= cfg.maximum,
                         free_disk_space=free_disk_space,
+                        pre_test=cfg.pre_test,
                     )
                 )
 
@@ -406,6 +409,7 @@ def expand_config(config):
                         java=java,
                         supported=False,
                         free_disk_space=free_disk_space,
+                        pre_test=cfg.pre_test,
                     )
                 )
     return matrix
@@ -461,7 +465,7 @@ def generate_matrix(args):
 class CustomEncoder(json.JSONEncoder):
     def default(self, o):
         if isinstance(o, MatrixItem):
-            return dict(o)
+            return o.model_dump(exclude_none=True)
         elif isinstance(o, Version):
             return str(o)
         return super().default(o)
diff --git a/mlflow/ml-package-versions.yml b/mlflow/ml-package-versions.yml
index 257a3ff97a140..6f31fa349a003 100644
--- a/mlflow/ml-package-versions.yml
+++ b/mlflow/ml-package-versions.yml
@@ -553,6 +553,8 @@ transformers:
       # prior to 4.30.0 are incompatible with this breaking change in Keras.
       # See: https://github.com/huggingface/transformers/issues/23352 for discussion
       "< 4.30.0": ["tensorflow<2.13.0"]
+    pre_test: |
+      python tests/transformers/helper.py
     run: |
       # Run all Transformers tests except autologging
       pytest tests/transformers --ignore tests/transformers/test_transformers_autolog.py