Add batch timeout override

src/promptflow-evals/CHANGELOG.md

@@ -3,6 +3,7 @@
 Please insert change log into "Next Release" ONLY.
 
 ## Next release
+- Exposed batch evaluation run timeout variable via "PF_BATCH_TIMEOUT_SEC" environment variable
 
 ## 0.0.1
 - Introduced package

src/promptflow-evals/promptflow/evals/_constants.py

@@ -23,4 +23,5 @@ class Prefixes:
 
 CONTENT_SAFETY_DEFECT_RATE_THRESHOLD_DEFAULT = 4
 
-BATCH_RUN_TIMEOUT = 3600
+PF_BATCH_TIMEOUT_DEFAULT = 3600
+PF_BATCH_TIMEOUT_SEC = "PF_BATCH_TIMEOUT_SEC"

src/promptflow-evals/promptflow/evals/evaluate/_batch_run_client/code_client.py

@@ -7,11 +7,12 @@
 
 import pandas as pd
 
+from promptflow._utils.utils import get_int_env_var
 from promptflow.contracts.types import AttrDict
 from promptflow.evals.evaluate._utils import _apply_column_mapping, _has_aggregator, load_jsonl
 from promptflow.tracing import ThreadPoolExecutorWithContext as ThreadPoolExecutor
 
-from ..._constants import BATCH_RUN_TIMEOUT
+from ..._constants import PF_BATCH_TIMEOUT_DEFAULT, PF_BATCH_TIMEOUT_SEC
 
 LOGGER = logging.getLogger(__name__)
 
@@ -24,15 +25,17 @@ def __init__(self, run, input_data, evaluator_name=None, aggregated_metrics=None
         self.aggregated_metrics = aggregated_metrics
 
     def get_result_df(self, exclude_inputs=False):
-        result_df = self.run.result(timeout=BATCH_RUN_TIMEOUT)
+        batch_run_timeout = get_int_env_var(PF_BATCH_TIMEOUT_SEC, PF_BATCH_TIMEOUT_DEFAULT)
+        result_df = self.run.result(timeout=batch_run_timeout)
         if exclude_inputs:
             result_df = result_df.drop(columns=[col for col in result_df.columns if col.startswith("inputs.")])
         return result_df
 
     def get_aggregated_metrics(self):
         try:
+            batch_run_timeout = get_int_env_var(PF_BATCH_TIMEOUT_SEC, PF_BATCH_TIMEOUT_DEFAULT)
             aggregated_metrics = (
-                self.aggregated_metrics.result(timeout=BATCH_RUN_TIMEOUT)
+                self.aggregated_metrics.result(timeout=batch_run_timeout)
                 if self.aggregated_metrics is not None
                 else None
             )

src/promptflow-evals/promptflow/evals/evaluate/_batch_run_client/proxy_client.py

@@ -5,10 +5,11 @@
 
 import numpy as np
 
+from promptflow._utils.utils import get_int_env_var
 from promptflow.client import PFClient
 from promptflow.tracing import ThreadPoolExecutorWithContext as ThreadPoolExecutor
 
-from ..._constants import BATCH_RUN_TIMEOUT
+from ..._constants import PF_BATCH_TIMEOUT_DEFAULT, PF_BATCH_TIMEOUT_SEC
 
 LOGGER = logging.getLogger(__name__)
 
@@ -30,11 +31,13 @@ def run(self, flow, data, column_mapping=None, **kwargs):
         return ProxyRun(run=eval_future)
 
     def get_details(self, proxy_run, all_results=False):
-        run = proxy_run.run.result(timeout=BATCH_RUN_TIMEOUT)
+        batch_run_timeout = get_int_env_var(PF_BATCH_TIMEOUT_SEC, PF_BATCH_TIMEOUT_DEFAULT)
+        run = proxy_run.run.result(timeout=batch_run_timeout)
         result_df = self._pf_client.get_details(run, all_results=all_results)
         result_df.replace("(Failed)", np.nan, inplace=True)
         return result_df
 
     def get_metrics(self, proxy_run):
-        run = proxy_run.run.result(timeout=BATCH_RUN_TIMEOUT)
+        batch_run_timeout = get_int_env_var(PF_BATCH_TIMEOUT_SEC, PF_BATCH_TIMEOUT_DEFAULT)
+        run = proxy_run.run.result(timeout=batch_run_timeout)
         return self._pf_client.get_metrics(run)

src/promptflow-evals/tests/evals/unittests/test_batch_run_context.py

@@ -1,10 +1,13 @@
+import os
 from unittest.mock import MagicMock
 
 import pytest
 
 from promptflow.client import PFClient
+from promptflow.evals._constants import PF_BATCH_TIMEOUT_DEFAULT, PF_BATCH_TIMEOUT_SEC
 from promptflow.evals._user_agent import USER_AGENT
 from promptflow.evals.evaluate._batch_run_client import BatchRunContext, CodeClient
+from promptflow.evals.evaluate._batch_run_client.code_client import CodeRun
 
 
 @pytest.fixture
@@ -17,6 +20,11 @@ def pf_client_mock():
     return MagicMock(spec=PFClient)
 
 
+@pytest.fixture
+def code_run_mock():
+    return MagicMock()
+
+
 @pytest.mark.unittest
 class TestBatchRunContext:
     def test_with_codeclient(self, mocker, code_client_mock):
@@ -50,3 +58,14 @@ def test_with_pfclient(self, mocker, pf_client_mock):
             pass
 
         mock_recover_openai_api.assert_not_called()
+
+    def test_get_result_timeout(self, code_run_mock):
+        code_run_instance = CodeRun(run=code_run_mock, input_data={})
+        code_run_instance.get_result_df()
+
+        code_run_mock.result.assert_called_once_with(timeout=PF_BATCH_TIMEOUT_DEFAULT)
+
+        custom_timeout = "100000"
+        os.environ[PF_BATCH_TIMEOUT_SEC] = custom_timeout
+        code_run_instance.get_result_df()
+        code_run_mock.result.assert_called_with(timeout=int(custom_timeout))