Add api status and use mock in test

Signed-off-by: Junyi Xu <219237550+JunyiXu-nv@users.noreply.github.com>

Author: JunyiXu-nv

tensorrt_llm/llmapi/llm.py

@@ -766,7 +766,13 @@ def shutdown(self) -> None:
             self.mpi_session.shutdown()
             self.mpi_session = None
 
+    @set_api_status("prototype")
     def check_health(self) -> bool:
+        """Check if the LLM is healthy.
+
+        Returns:
+            bool: True if the executor is running and not shutdown, False otherwise.
+        """
         if hasattr(self, "_executor") and self._executor is not None:
             return not self._executor.is_shutdown()
 

tests/unittest/api_stability/references/llm.yaml

@@ -271,6 +271,10 @@ methods:
     parameters: {}
     return_annotation: None
     status: beta
+  check_health:
+    parameters: {}
+    return_annotation: bool
+    status: prototype
 properties:
   llm_id:
     annotation: str

tests/unittest/llmapi/apps/_test_openai_metrics.py

@@ -1,5 +1,7 @@
 """Test the metrics endpoint when using OpenAI API to send requests"""
 
+from unittest.mock import patch
+
 import pytest
 from fastapi.testclient import TestClient
 
@@ -40,9 +42,12 @@ def client(llm):
                                                       (False, 503)])
 def test_health(client, llm, is_healthy, response_code):
     if not is_healthy:
-        llm.shutdown()
-    response = client.get("/health")
-    assert response.status_code == response_code
+        with patch.object(llm._executor, 'is_shutdown', return_value=True):
+            response = client.get("/health")
+            assert response.status_code == response_code
+    else:
+        response = client.get("/health")
+        assert response.status_code == response_code
 
 
 def test_version(client):