Merge pull request #298 from runpod/fastai-clarity

Fastai clarity

Author: justinmerrell

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Change Log
 
+## Release 1.6.2 (2/12/24)
+
+### Fixed
+
+- Reorder FastAPI endpoints for clarity.
+- Truncate long logs to avoid overloading the server.
+
+---
+
 ## Release 1.6.1 (2/11/24)
 
 ### Added

runpod/serverless/modules/rp_fastapi.py

@@ -163,7 +163,7 @@ def _send_webhook(url: str, payload: Dict[str, Any]) -> bool:
             response.raise_for_status()  # Raises exception for 4xx/5xx responses
             return True
         except requests.RequestException as err:
-            print(f"Request to {url} failed: {err}")
+            print(f"WEBHOOK | Request to {url} failed: {err}")
             return False
 
 
@@ -186,17 +186,17 @@ def __init__(self, config: Dict[str, Any]):
         self.config = config
 
         tags_metadata = [
+            {
+                "name": "Synchronously Submit Request & Get Job Results",
+                "description": "Endpoints for submitting job requests and getting the results."
+            },
             {
                 "name": "Submit Job Requests",
                 "description": "Endpoints for submitting job requests."
             },
             {
                 "name": "Check Job Results",
                 "description": "Endpoints for checking the status of a job and getting the results."
-            },
-            {
-                "name": "Synchronously Submit Request & Get Job Results",
-                "description": "Endpoints for submitting job requests and getting the results."
             }
         ]
 

runpod/serverless/modules/rp_job.py

@@ -184,18 +184,25 @@ async def run_job_generator(
     Run generator job used to stream output.
     Yields output partials from the generator.
     '''
+    is_async_gen = inspect.isasyncgenfunction(handler)
+    log.debug('Using Async Generator' if is_async_gen else 'Using Standard Generator', job["id"])
+
     try:
         job_output = handler(job)
-        if inspect.isasyncgenfunction(handler):
-            log.debug('Async generator', job["id"])
+
+        if is_async_gen:
             async for output_partial in job_output:
+                log.debug(f"Async Generator output: {output_partial}", job["id"])
                 yield {"output": output_partial}
         else:
-            log.debug('Generator', job["id"])
             for output_partial in job_output:
+                log.debug(f"Generator output: {output_partial}", job["id"])
                 yield {"output": output_partial}
+
     except Exception as err:    # pylint: disable=broad-except
         log.error(err, job["id"])
-        yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}
+        yield {
+            "error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"
+        }
     finally:
         log.info('Finished running generator.', job["id"])

runpod/serverless/modules/rp_logger.py

@@ -15,6 +15,7 @@
 from typing import Optional
 
 
+MAX_MESSAGE_LENGTH = 4096
 LOG_LEVELS = ['NOTSET', 'DEBUG', 'INFO', 'WARN', 'ERROR']
 
 
@@ -72,6 +73,14 @@ def log(self, message, message_level='INFO', job_id=None):
         if level_index > LOG_LEVELS.index(message_level) and message_level != 'TIP':
             return
 
+        message = str(message)
+        # Truncate message over 10MB, remove chunk from the middle
+        if len(message) > MAX_MESSAGE_LENGTH:
+            half_max_length = MAX_MESSAGE_LENGTH // 2
+            truncated_amount = len(message) - MAX_MESSAGE_LENGTH
+            truncation_note = f'\n...TRUNCATED {truncated_amount} CHARACTERS...\n'
+            message = message[:half_max_length] + truncation_note + message[-half_max_length:]
+
         if os.environ.get('RUNPOD_ENDPOINT_ID'):
             log_json = {
                 'requestId': job_id,

runpod/serverless/worker.py

@@ -50,6 +50,7 @@ async def _process_job(job, session, job_scaler, config):
 
         job_result = {'output': []}
         async for stream_output in generator_output:
+            log.debug(f"Stream output: {stream_output}", job['id'])
             if 'error' in stream_output:
                 job_result = stream_output
                 break

tests/test_serverless/test_modules/test_logger.py

@@ -127,9 +127,27 @@ def test_log_job_id(self):
             # Test with endpoint id set
             os.environ["RUNPOD_ENDPOINT_ID"] = "test_endpoint_id"
             logger.log("test_message", "INFO", job_id)
+            os.environ.pop("RUNPOD_ENDPOINT_ID")
 
             mock_print.assert_called_with(
                 '{"requestId": "test_job_id", "message": "test_message", "level": "INFO"}',
                 flush=True
             )
-            os.environ.pop("RUNPOD_ENDPOINT_ID")
+
+    def test_log_truncate(self):
+        """Tests that the log method truncates """
+        logger = rp_logger.RunPodLogger()
+        job_id = "test_job_id"
+        long_message = "a" * (rp_logger.MAX_MESSAGE_LENGTH + 100)
+        expected_start = "a" * (rp_logger.MAX_MESSAGE_LENGTH // 2)
+        expected_end = "a" * (rp_logger.MAX_MESSAGE_LENGTH // 2)
+        truncated_amount = len(long_message) - rp_logger.MAX_MESSAGE_LENGTH
+        truncation_note = f'\n...TRUNCATED {truncated_amount} CHARACTERS...\n'
+        truncated_message = expected_start + truncation_note + expected_end
+
+        with patch("builtins.print") as mock_print:
+            logger.log(long_message, "INFO", job_id)
+
+            expected_log_output = f'INFO   | {job_id} | {truncated_message}'
+
+            mock_print.assert_called_once_with(expected_log_output, flush=True)

tests/test_serverless/test_worker.py

@@ -160,9 +160,27 @@ def generator_handler_exception(job):
     '''
     print(job)
     yield "test1"
+    print("Raise exception")
     raise Exception()  # pylint: disable=broad-exception-raised
 
 
+def test_generator_handler_exception():
+    """ Test generator_handler_exception """
+    job = {"id": "test_job"}
+    gen = generator_handler_exception(job)
+
+    # Process the first yielded value
+    output = next(gen)
+    assert output == "test1", "First output should be 'test1'"
+
+    # Attempt to get the next value, expecting an exception
+    try:
+        next(gen)
+        assert False, "Expected an exception to be raised"
+    except Exception:  # pylint: disable=broad-except
+        assert True, "Exception was caught as expected"
+
+
 class TestRunWorker(IsolatedAsyncioTestCase):
     """ Tests for runpod | serverless| worker """
 
@@ -254,27 +272,27 @@ async def test_run_worker_generator_handler_exception(
         '''
         Test run_worker with generator handler.
 
-        Args:
-            mock_stream_result (_type_): _description_
-            mock_run_job_generator (_type_): _description_
-            mock_run_job (_type_): _description_
-            mock_get_job (_type_): _description_
+        This test verifies that:
+        - `stream_result` is called exactly once before an exception occurs.
+        - `run_job` is never called since `handler` is a generator function.
+        - An error is correctly reported back via `send_result`.
         '''
-        # Define the mock behaviors
-        mock_get_job.return_value = {
-            "id": "generator-123-exception", "input": {"number": 1}}
+        RunPodLogger().set_level("DEBUG")
 
-        # Test generator handler
-        generator_config = {
-            "handler": generator_handler_exception, "refresh_worker": True}
-        runpod.serverless.start(generator_config)
+        # Setup: Mock `get_job` to return a predefined job.
+        mock_get_job.return_value = {"id": "generator-123-exception", "input": {"number": 1}}
+
+        runpod.serverless.start({
+            "handler": generator_handler_exception,
+            "refresh_worker": True
+        })
 
         assert mock_stream_result.call_count == 1
         assert not mock_run_job.called
 
         # Since return_aggregate_stream is NOT activated, we should not submit any outputs.
         _, args, _ = mock_send_result.mock_calls[0]
-        assert 'error' in args[1]
+        assert 'error' in args[1], "Expected the error to be reported in the results."
 
     @patch("runpod.serverless.modules.rp_scale.get_job")
     @patch("runpod.serverless.worker.run_job")