(feat) Configure request method for progress tracking requests (#213)

Signed-off-by: Wojciech Prazuch <[email protected]>
Signed-off-by: Anna Warno <[email protected]>

Author: wprazuch

packages/nemo-evaluator/src/nemo_evaluator/adapters/adapter_config.py

@@ -537,6 +537,7 @@ def from_legacy_config(
                 "progress_tracking_interval": legacy_config[
                     "progress_tracking_interval"
                 ],
+                "request_method": "POST",  # Legacy method uses POST
                 "output_dir": cls._get_default_output_dir(legacy_config, run_config),
             }
             if legacy_config["progress_tracking_url"] is not None:

packages/nemo-evaluator/src/nemo_evaluator/adapters/interceptors/progress_tracking_interceptor.py

@@ -15,6 +15,7 @@
 
 """Progress tracking interceptor that tracks number of samples processed via webhook."""
 
+import os
 import pathlib
 import threading
 from typing import Optional, final
@@ -45,12 +46,16 @@ class Params(BaseLoggingParams):
 
         progress_tracking_url: Optional[str] = Field(
             default="http://localhost:8000",
-            description="URL to post the number of processed samples to.",
+            description="URL to post the number of processed samples to. Supports expansion of shell variables if present.",
         )
         progress_tracking_interval: int = Field(
             default=1,
             description="How often (every how many samples) to send a progress information.",
         )
+        request_method: str = Field(
+            default="PATCH",
+            description="Request method to use for updating the evaluation progress.",
+        )
         output_dir: Optional[str] = Field(
             default=None,
             description="Evaluation output directory. If provided, the progress tracking will be saved to a file in this directory.",
@@ -59,6 +64,7 @@ class Params(BaseLoggingParams):
     progress_tracking_url: Optional[str]
     progress_tracking_interval: int
     progress_filepath: Optional[pathlib.Path]
+    request_method: str
 
     def __init__(self, params: Params):
         """
@@ -67,8 +73,9 @@ def __init__(self, params: Params):
         Args:
             params: Configuration parameters
         """
-        self.progress_tracking_url = params.progress_tracking_url
+        self.progress_tracking_url = os.path.expandvars(params.progress_tracking_url)
         self.progress_tracking_interval = params.progress_tracking_interval
+        self.request_method = params.request_method
         if params.output_dir is not None:
             output_dir = pathlib.Path(params.output_dir)
             output_dir.mkdir(parents=True, exist_ok=True)
@@ -111,20 +118,32 @@ def _write_progress(self, num_samples: int):
                 samples_processed=num_samples,
             )
 
-    def _send_progress(self, num_samples: int):
+    def _send_progress(self, num_samples: int) -> requests.Response:
         self.logger.debug(
             "Sending progress to tracking server",
             url=self.progress_tracking_url,
+            method=self.request_method,
             samples_processed=num_samples,
         )
+        body = {"samples_processed": num_samples}
         try:
-            requests.post(
+            resp = requests.request(
+                self.request_method,
                 self.progress_tracking_url,
-                json={"samples_processed": num_samples},
-            )
-            self.logger.debug(
-                "Progress sent successfully", samples_processed=num_samples
+                json=body,
             )
+            if resp.status_code >= 200 and resp.status_code < 300:
+                self.logger.debug(
+                    "Progress sent successfully", samples_processed=num_samples
+                )
+            else:
+                self.logger.warning(
+                    "Failed to update job progress",
+                    body=body,
+                    status_code=resp.status_code,
+                    response_text=resp.text,
+                )
+            return resp
         except requests.exceptions.RequestException as e:
             self.logger.error(
                 "Failed to communicate with progress tracking server",

packages/nemo-evaluator/tests/unit_tests/adapters/interceptors/test_progress_tracking_interceptor.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import threading
 import time
 from typing import List
@@ -34,13 +35,13 @@
 class FakeProgressTrackingServer:
     """Test server to receive progress tracking webhooks."""
 
-    def __init__(self, port: int = 8000):
+    def __init__(self, port: int = 8000, request_method="PATCH"):
         self.port = port
         self.app = Flask(__name__)
         self.received_updates: List[dict] = []
         self.lock = threading.Lock()
 
-        @self.app.route("/", methods=["POST"])
+        @self.app.route("/", methods=[request_method])
         def progress_webhook():
             """Receive progress updates."""
             data = request.get_json()
@@ -93,6 +94,18 @@ def test_init_with_custom_params(self):
         assert interceptor.progress_tracking_url == "http://test-server:9000"
         assert interceptor.progress_tracking_interval == 5
 
+    @patch.dict(os.environ, {"NEMO_JOB_ID": "job-1234"})
+    def test_init_url_with_env_expansion(self):
+        """Test initialization with URL with env variable is expanded."""
+        params = ProgressTrackingInterceptor.Params(
+            progress_tracking_url="http://test-server:8000/jobs/${NEMO_JOB_ID}/status-details"
+        )
+        interceptor = ProgressTrackingInterceptor(params)
+        assert (
+            interceptor.progress_tracking_url
+            == "http://test-server:8000/jobs/job-1234/status-details"
+        )
+
     def test_intercept_response_sends_progress(self):
         """Test that intercept_response sends progress updates."""
         # Start test server
@@ -213,11 +226,11 @@ def process_samples():
         finally:
             server.stop()
 
-    @patch("requests.post")
-    def test_network_error_handling(self, mock_post):
+    @patch("requests.request")
+    def test_network_error_handling(self, mock_request):
         """Test that network errors are handled gracefully."""
-        # Mock requests.post to raise an exception
-        mock_post.side_effect = requests.exceptions.RequestException(
+        # Mock requests.patch to raise an exception
+        mock_request.side_effect = requests.exceptions.RequestException(
             "Connection failed"
         )
 
@@ -240,7 +253,7 @@ def test_network_error_handling(self, mock_post):
         assert result == mock_response
 
         # Verify that the request was attempted
-        mock_post.assert_called_once()
+        mock_request.assert_called_once()
 
     def test_interval_configuration(self):
         """Test different interval configurations."""
@@ -308,6 +321,52 @@ def test_json_payload_format(self):
         finally:
             server.stop()
 
+    def test_configured_method(self):
+        """Test that the JSON payload has the correct format."""
+        # Start test server
+        server = FakeProgressTrackingServer(port=8006, request_method="POST")
+        server.start()
+
+        try:
+            # Create interceptor
+            params = ProgressTrackingInterceptor.Params(
+                progress_tracking_url="http://localhost:8006",
+                progress_tracking_interval=1,
+                request_method="POST",
+            )
+            interceptor = ProgressTrackingInterceptor(params)
+
+            mock_response = AdapterResponse(
+                r=requests.Response(),
+                rctx=AdapterRequestContext(),
+            )
+            context = AdapterGlobalContext(output_dir="/tmp", url="http://test")
+
+            # Process one sample
+            interceptor.intercept_response(mock_response, context)
+
+            # Check the payload format
+            updates = server.get_updates()
+            assert len(updates) == 1
+            assert "samples_processed" in updates[0]
+            assert updates[0]["samples_processed"] == 1
+            assert isinstance(updates[0]["samples_processed"], int)
+
+            # Verify PATCH does not update the server
+            params = ProgressTrackingInterceptor.Params(
+                progress_tracking_url="http://localhost:8006",
+                progress_tracking_interval=1,
+                request_method="PATCH",
+            )
+            interceptor = ProgressTrackingInterceptor(params)
+            interceptor.intercept_response(mock_response, context)
+            assert updates == server.get_updates(), (
+                "server should not update with misconfigured method"
+            )
+
+        finally:
+            server.stop()
+
 
 if __name__ == "__main__":
     # Simple test runner for manual testing