Merge pull request #227 from DireLines/alignSdk

Align sdk skeleton

Author: justinmerrell

runpod/endpoint/runner.py

@@ -13,7 +13,9 @@
                        "An API key can be generated at "
                        "https://runpod.io/console/user/settings")
 
-
+def is_completed(status:str)->bool:
+    """Returns true if status is one of the possible final states for a serverless request."""
+    return status in ["COMPLETED", "FAILED", "TIMED_OUT", "CANCELLED"]
 # ---------------------------------------------------------------------------- #
 #                                    Client                                    #
 # ---------------------------------------------------------------------------- #
@@ -107,7 +109,7 @@ def _fetch_job(self):
         status_url = f"{self.endpoint_id}/status/{self.job_id}"
         job_state = self.rp_client.get(endpoint=status_url)
 
-        if job_state["status"] in ["COMPLETED", "FAILED", "TIMEOUT"]:
+        if is_completed(job_state["status"]):
             self.job_status = job_state["status"]
             self.job_output = job_state.get("output", None)
 
@@ -128,7 +130,7 @@ def output(self, timeout: int = 0) -> Any:
             timeout: The number of seconds to wait for the server to send data before giving up.
         """
         if timeout > 0:
-            while self.status() not in ["COMPLETED", "FAILED", "TIMEOUT"]:
+            while not is_completed(self.status()):
                 time.sleep(1)
                 timeout -= 1
                 if timeout <= 0:
@@ -139,6 +141,17 @@ def output(self, timeout: int = 0) -> Any:
 
         return self._fetch_job().get("output", None)
 
+    def cancel(self, timeout: int = 3) -> Any:
+        """
+        Cancels the job and returns the result of the cancellation request.
+
+        Args:
+            timeout: The number of seconds to wait for the server to respond before giving up.
+        """
+        return self.rp_client.post(f"{self.endpoint_id}/cancel/{self.job_id}",
+                                   data=None,timeout=timeout)
+
+
 
 # ---------------------------------------------------------------------------- #
 #                                   Endpoint                                   #
@@ -195,3 +208,20 @@ def run_sync(self, request_input: Dict[str, Any], timeout: int = 86400) -> Dict[
             return job_request.get("output", None)
 
         return Job(self.endpoint_id, job_request["id"], self.rp_client).output(timeout=timeout)
+
+    def health(self,timeout: int = 3) -> Dict[str, Any]:
+        """
+        Check the health of the endpoint (number/state of workers, number/state of requests).
+
+        Args:
+            timeout: The number of seconds to wait for the server to respond before giving up.
+        """
+        return self.rp_client.get(f"{self.endpoint_id}/health",timeout=timeout)
+    def purge_queue(self,timeout: int = 3) -> Dict[str, Any]:
+        """
+        Purges the endpoint's job queue and returns the result of the purge request.
+
+        Args:
+            timeout: The number of seconds to wait for the server to respond before giving up.
+        """
+        return self.rp_client.post(f"{self.endpoint_id}/purge-queue",data=None,timeout=timeout)

tests/test_endpoint/test_runner.py

@@ -132,6 +132,23 @@ def test_endpoint_run_sync(self, mock_client_request):
             {'input': {'YOUR_MODEL_INPUT_JSON': 'YOUR_MODEL_INPUT_VALUE'}}, 86400
         )
 
+    @patch('runpod.endpoint.runner.RunPodClient._request')
+    def test_endpoint_health(self, mock_client_request):
+        ''' Test the health method of Endpoint '''
+        self.endpoint.health()
+
+        mock_client_request.assert_called_once_with('GET', f"{self.ENDPOINT_ID}/health", timeout=3)
+
+    @patch('runpod.endpoint.runner.RunPodClient._request')
+    def test_endpoint_purge_queue(self, mock_client_request):
+        ''' Test the health method of Endpoint '''
+        self.endpoint.purge_queue()
+
+        mock_client_request.assert_called_once_with(
+            'POST', f"{self.ENDPOINT_ID}/purge-queue",
+            None, 3
+        )
+
     def test_missing_api_key(self):
         '''
         Tests Endpoint.run without api_key
@@ -213,6 +230,7 @@ def test_run_sync_with_timeout(self, mock_client_request):
 
 class TestJob(unittest.TestCase):
     ''' Tests for Job '''
+    MODEL_OUTPUT = {"result": "YOUR_MODEL_OUTPUT_VALUE"}
 
     @patch('runpod.endpoint.runner.RunPodClient')
     def test_status(self, mock_client):
@@ -274,6 +292,17 @@ def test_output_timeout(self, mock_client):
         with self.assertRaises(TimeoutError):
             job.output(timeout=1)
 
+    @patch('runpod.endpoint.runner.RunPodClient')
+    def test_cancel(self, mock_client):
+        ''' Test the cancel method of Job with a successful job initiation. '''        
+        job = runner.Job("endpoint_id", "job_id", mock_client)
+
+        job.cancel()
+
+        mock_client.post.assert_called_with("endpoint_id/cancel/job_id",
+                                            data=None, timeout=3)
+
+
     @patch('runpod.endpoint.runner.RunPodClient')
     def test_job_status(self, mock_client):
         '''