Merge pull request #263 from runpod/update-async-runner

Update async runner

Author: justinmerrell

runpod/endpoint/asyncio/asyncio_runner.py

@@ -1,13 +1,12 @@
-'''
-Author: Oleg Rybalko
-Github: https://github.com/SkullMag
-Date: 2023-03-27
-'''
+""" Module for running endpoints asynchronously. """
 # pylint: disable=too-few-public-methods,R0801
 
+from typing import Any, Dict
 import asyncio
 import aiohttp
 
+from runpod.endpoint.helpers import FINAL_STATES, is_completed
+
 
 class Job:
     """Class representing a job for an asynchronous endpoint"""
@@ -17,45 +16,84 @@ def __init__(self, endpoint_id: str, job_id: str, session: aiohttp.ClientSession
 
         self.endpoint_id = endpoint_id
         self.job_id = job_id
-        self.status_url = f"{endpoint_url_base}/{self.endpoint_id}/status/{self.job_id}"
-        self.cancel_url = f"{endpoint_url_base}/{self.endpoint_id}/cancel/{self.job_id}"
         self.headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {api_key}"
         }
         self.session = session
+        self.endpoint_url_base = endpoint_url_base
+
+        self.job_status = None
+        self.job_output = None
+
+    async def _fetch_job(self, source: str = "status") -> Dict[str, Any]:
+        """ Returns the raw json of the status, reaises an exception if invalid.
+
+        Args:
+            source: The URL source path of the job status.
+        """
+        status_url = f"{self.endpoint_url_base}/{self.endpoint_id}/{source}/{self.job_id}"
+        job_state = await self.session.get(status_url, headers=self.headers)
+        job_state = await job_state.json()
+
+        if is_completed(job_state["status"]):
+            self.job_status = job_state["status"]
+            self.job_output = job_state.get("output", None)
+
+        return job_state
 
     async def status(self) -> str:
         """Gets jobs' status
 
         Returns:
             COMPLETED, FAILED or IN_PROGRESS
         """
-        async with self.session.get(self.status_url, headers=self.headers) as resp:
-            return (await resp.json())["status"]
+        if self.job_status is not None:
+            return self.job_status
+
+        job_state = await self._fetch_job()
+        return job_state["status"]
+
+    async def _wait_for_completion(self):
+        while not is_completed(await self.status()):
+            await asyncio.sleep(1)
 
-    async def output(self) -> any:
+    async def output(self, timeout: int = 0) -> Any:
         """Waits for serverless API job to complete or fail
 
         Returns:
             Output of job
         Raises:
             KeyError if job Failed
         """
-        while await self.status() not in ["COMPLETED", "FAILED"]:
-            await asyncio.sleep(1)
+        if self.job_output is not None:
+            return self.job_output
+
+        try:
+            await asyncio.wait_for(self._wait_for_completion(), timeout)
+        except asyncio.TimeoutError as exc:
+            raise TimeoutError("Job timed out.") from exc
 
-        async with self.session.get(self.status_url, headers=self.headers) as resp:
-            return (await resp.json())["output"]
+        job_data = await self._fetch_job()
+        return job_data.get("output", None)
+
+    async def stream(self) -> Any:
+        """ Returns a generator that yields the output of the job request. """
+        while True:
+            await asyncio.sleep(1)
+            stream_partial = await self._fetch_job(source="stream")
+            if stream_partial["status"] not in FINAL_STATES:
+                for chunk in stream_partial.get("stream", []):
+                    yield chunk["output"]
 
     async def cancel(self) -> dict:
         """Cancels current job
 
         Returns:
             Output of cancel operation
         """
-
-        async with self.session.post(self.cancel_url, headers=self.headers) as resp:
+        cancel_url = f"{self.endpoint_url_base}/{self.endpoint_id}/cancel/{self.job_id}"
+        async with self.session.post(cancel_url, headers=self.headers) as resp:
             return await resp.json()
 
 
@@ -88,3 +126,21 @@ async def run(self, endpoint_input: dict) -> Job:
             json_resp = await resp.json()
 
         return Job(self.endpoint_id, json_resp["id"], self.session)
+
+    async def health(self) -> dict:
+        """Checks health of endpoint
+
+        Returns:
+            Health of endpoint
+        """
+        async with self.session.get(f"{self.endpoint_id}/health", headers=self.headers) as resp:
+            return await resp.json()
+
+    async def purge_queue(self) -> dict:
+        """Purges queue of endpoint
+
+        Returns:
+            Purge status
+        """
+        async with self.session.post(f"{self.endpoint_id}/purge", headers=self.headers) as resp:
+            return await resp.json()

runpod/endpoint/helpers.py

@@ -0,0 +1,14 @@
+""" Helper functions for the RunPod Endpoint API. """
+
+FINAL_STATES = ["COMPLETED", "FAILED", "TIMED_OUT"]
+
+# Exception Messages
+UNAUTHORIZED_MSG = "401 Unauthorized | Make sure Runpod API key is set and valid."
+API_KEY_NOT_SET_MSG = ("Expected `run_pod.api_key` to be initialized. "
+                       "You can solve this by setting `run_pod.api_key = 'your-key'. "
+                       "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"]

runpod/endpoint/runner.py

@@ -6,19 +6,9 @@
 import requests
 from requests.adapters import HTTPAdapter, Retry
 
-FINAL_STATES = ["COMPLETED", "FAILED", "TIMED_OUT"]
-
-# Exception Messages
-UNAUTHORIZED_MSG = "401 Unauthorized | Make sure Runpod API key is set and valid."
-API_KEY_NOT_SET_MSG = ("Expected `run_pod.api_key` to be initialized. "
-                       "You can solve this by setting `run_pod.api_key = 'your-key'. "
-                       "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"]
+from runpod.endpoint.helpers import (
+    FINAL_STATES, UNAUTHORIZED_MSG, API_KEY_NOT_SET_MSG, is_completed
+)
 
 
 # ---------------------------------------------------------------------------- #
@@ -146,16 +136,6 @@ 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)
-
     def stream(self) -> Any:
         """ Returns a generator that yields the output of the job request. """
         while True:
@@ -167,6 +147,16 @@ def stream(self) -> Any:
             elif stream_partial["status"] in FINAL_STATES:
                 break
 
+    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                                   #

tests/test_endpoint/test_asyncio_runner.py

@@ -4,7 +4,7 @@
 import tracemalloc
 import asyncio
 import unittest
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch, MagicMock, AsyncMock
 from unittest import IsolatedAsyncioTestCase
 
 from runpod.endpoint.asyncio.asyncio_runner import Job, Endpoint
@@ -19,37 +19,91 @@ async def test_status(self):
         '''
         Tests Job.status
         '''
-        with patch("aiohttp.ClientSession") as mock_session:
-            mock_resp = MagicMock()
-            mock_resp.json = MagicMock(return_value=asyncio.Future())
-            mock_resp.json.return_value.set_result({"status": "COMPLETED"})
-            mock_session.get.return_value.__aenter__.return_value = mock_resp
+        with patch("aiohttp.ClientSession", new_callable=AsyncMock) as mock_session_class:
+            mock_session = mock_session_class.return_value
+            mock_get = mock_session.get
+            mock_resp = AsyncMock()
+
+            mock_resp.json.return_value = {"status": "COMPLETED"}
+            mock_get.return_value = mock_resp
 
             job = Job("endpoint_id", "job_id", mock_session)
             status = await job.status()
             assert status == "COMPLETED"
+            assert await job.status() == "COMPLETED"
 
     async def test_output(self):
         '''
         Tests Job.output
         '''
         with patch("runpod.endpoint.asyncio.asyncio_runner.asyncio.sleep") as mock_sleep, \
-             patch("aiohttp.ClientSession") as mock_session:
-            mock_resp = MagicMock()
+             patch("aiohttp.ClientSession", new_callable=AsyncMock) as mock_session_class:
+            mock_session = mock_session_class.return_value
+            mock_get = mock_session.get
+            mock_resp = AsyncMock()
 
             async def json_side_effect():
                 if mock_sleep.call_count == 0:
                     return {"status": "IN_PROGRESS"}
                 return {"output": "OUTPUT", "status": "COMPLETED"}
 
-            mock_resp.json = json_side_effect
-            mock_session.get.return_value.__aenter__.return_value = mock_resp
+            mock_resp.json.side_effect = json_side_effect
+            mock_get.return_value = mock_resp
 
             job = Job("endpoint_id", "job_id", mock_session)
-            output_task = asyncio.create_task(job.output())
+            output_task = asyncio.create_task(job.output(timeout=3))
 
             output = await output_task
             assert output == "OUTPUT"
+            assert await job.output() == "OUTPUT"
+
+    async def test_output_timeout(self):
+        '''
+        Tests Job.output with a timeout
+        '''
+        with patch("aiohttp.ClientSession", new_callable=AsyncMock) as mock_session_class:
+            mock_session = mock_session_class.return_value
+            mock_get = mock_session.get
+            mock_resp = AsyncMock()
+
+            mock_resp.json.return_value = {"status": "IN_PROGRESS"}
+            mock_get.return_value = mock_resp
+
+            job = Job("endpoint_id", "job_id", mock_session)
+            output_task = asyncio.create_task(job.output(timeout=1))
+
+            with self.assertRaises(TimeoutError):
+                await output_task
+
+    async def test_stream(self):
+        '''
+        Tests Job.stream
+        '''
+        with patch("aiohttp.ClientSession", new_callable=AsyncMock) as mock_session_class:
+            mock_session = mock_session_class.return_value
+            mock_get = mock_session.get
+            mock_resp = AsyncMock()
+
+            responses = [
+                {"stream": [{"output": "OUTPUT1"}], "status": "IN_PROGRESS"},
+                {"stream": [{"output": "OUTPUT2"}], "status": "IN_PROGRESS"},
+            ]
+
+            async def json_side_effect():
+                return responses.pop(0) if responses else {"stream": [], "status": "COMPLETED"}
+
+            mock_resp.json.side_effect = json_side_effect
+            mock_get.return_value = mock_resp
+
+            job = Job("endpoint_id", "job_id", mock_session)
+
+            outputs = []
+            async for stream_output in job.stream():
+                outputs.append(stream_output)
+                if not responses:  # Break the loop when responses are exhausted
+                    break
+
+            assert outputs == ["OUTPUT1", "OUTPUT2"]
 
     async def test_cancel(self):
         '''
@@ -67,27 +121,25 @@ async def test_cancel(self):
 
     async def test_output_in_progress_then_completed(self):
         '''Tests Job.output when status is initially IN_PROGRESS and then changes to COMPLETED'''
-        with patch("runpod.endpoint.asyncio.asyncio_runner.asyncio.sleep") as mock_sleep, \
-             patch("aiohttp.ClientSession") as mock_session:
-            mock_resp = MagicMock()
+        with patch("aiohttp.ClientSession", new_callable=AsyncMock) as mock_session_class:
+            mock_session = mock_session_class.return_value
+            mock_get = mock_session.get
+            mock_resp = AsyncMock()
+
             responses = [
                 {"status": "IN_PROGRESS"},
-                {"status": "COMPLETED"},
-                {"output": "OUTPUT"}
+                {"status": "COMPLETED", "output": "OUTPUT"}
             ]
 
             async def json_side_effect():
-                if responses:
-                    return responses.pop(0)
-                return {"status": "IN_PROGRESS"}
+                return responses.pop(0) if responses else {"status": "COMPLETED", "output": "OUTPUT"} # pylint: disable=line-too-long
 
-            mock_resp.json = json_side_effect
-            mock_session.get.return_value.__aenter__.return_value = mock_resp
+            mock_resp.json.side_effect = json_side_effect
+            mock_get.return_value = mock_resp
 
             job = Job("endpoint_id", "job_id", mock_session)
-            output = await job.output()
+            output = await job.output(timeout=3)
             assert output == "OUTPUT"
-            mock_sleep.assert_called_once_with(1)
 
 class TestEndpoint(IsolatedAsyncioTestCase):
     ''' Unit tests for the Endpoint class. '''
@@ -106,6 +158,34 @@ async def test_run(self):
             job = await endpoint.run({"input": "INPUT"})
             assert job.job_id == "job_id"
 
+    async def test_health(self):
+        '''
+        Tests Endpoint.health
+        '''
+        with patch("aiohttp.ClientSession") as mock_session:
+            mock_resp = MagicMock()
+            mock_resp.json = MagicMock(return_value=asyncio.Future())
+            mock_resp.json.return_value.set_result({"status": "HEALTHY"})
+            mock_session.get.return_value.__aenter__.return_value = mock_resp
+
+            endpoint = Endpoint("endpoint_id", mock_session)
+            health = await endpoint.health()
+            assert health == {"status": "HEALTHY"}
+
+    async def test_purge_queue(self):
+        '''
+        Tests Endpoint.purge_queue
+        '''
+        with patch("aiohttp.ClientSession") as mock_session:
+            mock_resp = MagicMock()
+            mock_resp.json = MagicMock(return_value=asyncio.Future())
+            mock_resp.json.return_value.set_result({"result": "PURGED"})
+            mock_session.post.return_value.__aenter__.return_value = mock_resp
+
+            endpoint = Endpoint("endpoint_id", mock_session)
+            purge_result = await endpoint.purge_queue()
+            assert purge_result == {"result": "PURGED"}
+
 class TestEndpointInitialization(unittest.TestCase):
     '''Tests for the Endpoint class initialization.'''