Feat: Improve LMEval assertions

Add additional validation for LMEval Job pods where pod logs are checked and better exception handling and logging

Author: sheltoncyril

tests/model_explainability/lm_eval/test_lm_eval.py

@@ -5,7 +5,7 @@
 
 from tests.model_explainability.utils import validate_tai_component_images
 
-from tests.model_explainability.lm_eval.utils import get_lmeval_tasks
+from tests.model_explainability.lm_eval.utils import get_lmeval_tasks, validate_lmeval_job_pod_and_logs
 
 LMEVALJOB_COMPLETE_STATE: str = "Complete"
 
@@ -64,7 +64,7 @@
 def test_lmeval_huggingface_model(admin_client, model_namespace, lmevaljob_hf_pod):
     """Tests that verify running common evaluations (and a custom one) on a model pulled directly from HuggingFace.
     On each test we run a different evaluation task, limiting it to 0.5% of the questions on each eval."""
-    lmevaljob_hf_pod.wait_for_status(status=lmevaljob_hf_pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_40MIN)
+    validate_lmeval_job_pod_and_logs(lmevaljob_hf_pod)
 
 
 @pytest.mark.parametrize(
@@ -89,9 +89,7 @@ def test_lmeval_local_offline_builtin_tasks_flan_arceasy(
     lmevaljob_local_offline_pod,
 ):
     """Test that verifies that LMEval can run successfully in local, offline mode using builtin tasks"""
-    lmevaljob_local_offline_pod.wait_for_status(
-        status=lmevaljob_local_offline_pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_20MIN
-    )
+    validate_lmeval_job_pod_and_logs(lmevaljob_local_offline_pod)
 
 
 @pytest.mark.parametrize(
@@ -124,9 +122,7 @@ def test_lmeval_local_offline_unitxt_tasks_flan_20newsgroups(
     lmevaljob_local_offline_pod,
 ):
     """Test that verifies that LMEval can run successfully in local, offline mode using unitxt"""
-    lmevaljob_local_offline_pod.wait_for_status(
-        status=lmevaljob_local_offline_pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_20MIN
-    )
+    validate_lmeval_job_pod_and_logs(lmevaljob_local_offline_pod)
 
 
 @pytest.mark.parametrize(
@@ -140,9 +136,7 @@ def test_lmeval_local_offline_unitxt_tasks_flan_20newsgroups(
 )
 def test_lmeval_vllm_emulator(admin_client, model_namespace, lmevaljob_vllm_emulator_pod):
     """Basic test that verifies LMEval works with vLLM using a vLLM emulator for more efficient evaluation"""
-    lmevaljob_vllm_emulator_pod.wait_for_status(
-        status=lmevaljob_vllm_emulator_pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_20MIN
-    )
+    validate_lmeval_job_pod_and_logs(lmevaljob_vllm_emulator_pod)
 
 
 @pytest.mark.parametrize(
@@ -161,9 +155,7 @@ def test_lmeval_s3_storage(
     lmevaljob_s3_offline_pod,
 ):
     """Test to verify that LMEval works with a model stored in a S3 bucket"""
-    lmevaljob_s3_offline_pod.wait_for_status(
-        status=lmevaljob_s3_offline_pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_20MIN
-    )
+    validate_lmeval_job_pod_and_logs(lmevaljob_s3_offline_pod)
 
 
 @pytest.mark.parametrize(

tests/model_explainability/lm_eval/utils.py

@@ -1,14 +1,17 @@
 from typing import List
-
+import re
 from kubernetes.dynamic import DynamicClient
 from ocp_resources.lm_eval_job import LMEvalJob
 from ocp_resources.pod import Pod
 
+from tests.model_explainability.utils import log_pod_failure_logs
 from utilities.constants import Timeout
 from simple_logger.logger import get_logger
+from timeout_sampler import TimeoutExpiredError
 
 import pandas as pd
 
+from utilities.exceptions import PodLogMissMatchError, UnexpectedFailureError
 
 LOGGER = get_logger(name=__name__)
 
@@ -84,3 +87,28 @@ def get_lmeval_tasks(min_downloads: int | float, max_downloads: int | float | No
     LOGGER.info(f"Number of unique LMEval tasks with more than {min_downloads} downloads: {len(unique_tasks)}")
 
     return unique_tasks
+
+def validate_lmeval_job_pod_and_logs(lmevaljob_pod: Pod) -> None:
+    """Validate LMEval job pod success and presence of corresponding logs.
+
+    Args:
+        lmevaljob_pod: The LMEvalJob pod.
+
+    Returns: None
+    """
+    pod_success_log_regex = r'INFO\sdriver\supdate status: job completed\s\{\"state\":\s\{\"state\"'\
+                            r':\"Complete\",\"reason\":\"Succeeded\",\"message\":\"job completed\"'
+    try:
+        lmevaljob_pod.wait_for_status(
+        status=lmevaljob_pod.Status.RUNNING, timeout=Timeout.TIMEOUT_5MIN
+    )
+    except TimeoutExpiredError as e:
+        raise UnexpectedFailureError(f"LMEval job pod did not reach a running state. Status: {lmevaljob_pod.status}") from e
+    try:
+        lmevaljob_pod.wait_for_status(Pod.Status.SUCCEEDED, timeout=Timeout.TIMEOUT_10MIN)
+    except TimeoutExpiredError as e:
+        log_pod_failure_logs(LOGGER, lmevaljob_pod)
+        raise UnexpectedFailureError("LMEval job pod failed from a running state.") from e
+    if not bool(re.search(pod_success_log_regex, lmevaljob_pod.log())):
+        log_pod_failure_logs(LOGGER, lmevaljob_pod)
+        raise PodLogMissMatchError("LMEval job pod failed.")

tests/model_explainability/utils.py

@@ -1,3 +1,4 @@
+from logging import Logger
 import re
 from ocp_resources.config_map import ConfigMap
 from ocp_resources.pod import Pod
@@ -32,3 +33,26 @@ def validate_tai_component_images(
         assert container.image in tai_configmap_values, (
             f"{container.name} : {container.image} not present in TrustyAI operator configmap."
         )
+
+def log_pod_failure_logs(logger: Logger, pod: Pod) -> None:
+    """Log pod info and logs to logger.
+
+    Args:
+        logger: logging.Logger
+         The logger to output to.
+        pod: Pod
+         The relevant pod to gain information on.
+
+    Returns: None
+    """
+    logger.error("--------------------------------- FAILED POD INFO -----------------------------------")
+    logger.error(f"Failed Pod Name: {pod.name}")
+    logger.error(f"Failed Pod Status: {pod.status}")
+    logger.error(f"Failed Pod IP: {pod.ip}")
+    logger.error(f"Failed Pod Labels: {pod.labels}")
+    logger.error(f"Failed Pod Namespace: {pod.namespace}")
+    logger.error(f"Failed Pod Spec: {pod.instance.spec.to_dict()}")
+    logger.error("--------------------------------- FAILED POD LOGS ----------------------------------")
+    logger.error(f"{pod.log()}")
+    logger.error("--------------------------------- END OF POD LOGS ----------------------------------")
+