Merge pull request #117 from runpod/logging-transparency

Logging transparency

Author: justinmerrell

runpod/serverless/modules/rp_http.py

@@ -49,7 +49,7 @@ async def _handle_result(session, job_data, job, url_template, log_message):
         log.debug(f"{job['id']} | {log_message}")
 
     except aiohttp.ClientError as err:
-        log.error(f"{job['id']} | Client response error while transmitting job. | {err}")
+        log.error(f"{job['id']} | Failed to return job results. | {err}")
 
     except (TypeError, RuntimeError) as err:
         log.error(f"Error while returning job result {job['id']}: {err}")

runpod/serverless/modules/rp_job.py

@@ -1,7 +1,9 @@
 """
 Job related helpers.
 """
+# pylint: disable=too-many-branches
 
+import sys
 import inspect
 from typing import Any, Callable, Dict, Generator, Optional, Union
 
@@ -55,7 +57,11 @@ async def get_job(session: ClientSession, retry=True) -> Optional[Dict[str, Any]
                         return None
                     continue
 
-                if response.status not in [200, 400]:
+                if response.status == 400:
+                    log.info("Received 400 status, stopping worker.")
+                    sys.exit(0)
+
+                if response.status != 200:
                     log.error(f"Failed to get job, status code: {response.status}")
                     if not retry:
                         return None

tests/test_serverless/test_modules/test_job.py

@@ -62,6 +62,24 @@ async def test_get_job_204(self):
             assert job is None
             assert mock_session_204.get.call_count == 1
 
+    async def test_get_job_400(self):
+        '''
+        Test the get_job function with a 400 response
+        '''
+        # 400 Mock
+        response_400 = Mock(ClientResponse)
+        response_400.status = 400
+
+        with patch("runpod.serverless.modules.rp_job.sys.exit") as mock_exit, \
+             patch("aiohttp.ClientSession") as mock_session_400, \
+             patch("runpod.serverless.modules.rp_job.JOB_GET_URL", "http://mock.url"):
+
+            mock_session_400.get.return_value.__aenter__.return_value = response_400
+            job = await rp_job.get_job(mock_session_400, retry=False)
+
+            assert job is None
+            assert mock_exit.call_count == 1
+
     async def test_get_job_500(self):
         '''
         Tests the get_job function with a 500 response