Merge pull request #135 from runpod/main

updating branch

Author: justinmerrell

runpod/serverless/__init__.py

@@ -8,6 +8,7 @@
 import sys
 import json
 import time
+import signal
 import argparse
 from typing import Dict, Any
 
@@ -82,6 +83,13 @@ def _get_realtime_concurrency() -> int:
     """
     return int(os.environ.get("RUNPOD_REALTIME_CONCURRENCY", "1"))
 
+def _signal_handler(sig, frame):
+    """
+    Handles the SIGINT signal.
+    """
+    del sig, frame
+    log.info("SIGINT received. Shutting down.")
+    sys.exit(0)
 
 # ---------------------------------------------------------------------------- #
 #                            Start Serverless Worker                           #
@@ -100,6 +108,8 @@ def start(config: Dict[str, Any]):
     from runpod import __version__ as runpod_version # pylint: disable=import-outside-toplevel,cyclic-import
     print(f"--- Starting Serverless Worker |  Version {runpod_version} ---")
 
+    signal.signal(signal.SIGINT, _signal_handler)
+
     config["reference_counter_start"] = time.perf_counter()
     config = _set_config_args(config)
 

runpod/serverless/utils/rp_upload.py

@@ -10,6 +10,7 @@
 import threading
 import multiprocessing
 from io import BytesIO
+from urllib.parse import urlparse
 from typing import Optional, Tuple
 
 import boto3
@@ -18,24 +19,27 @@
 from boto3.s3.transfer import TransferConfig
 from botocore.config import Config
 from tqdm_loggable.auto import tqdm
-from urllib.parse import urlparse
+
 
 logger = logging.getLogger("runpod upload utility")
 FMT = "%(filename)-20s:%(lineno)-4d %(asctime)s %(message)s"
 logging.basicConfig(level=logging.INFO, format=FMT, handlers=[logging.StreamHandler()])
 
 def extract_region_from_url(endpoint_url):
+    """
+    Extracts the region from the endpoint URL.
+    """
     parsed_url = urlparse(endpoint_url)
     # AWS/backblaze S3-like URL
     if '.s3.' in endpoint_url:
         return endpoint_url.split('.s3.')[1].split('.')[0]
+
     # DigitalOcean Spaces-like URL
-    elif parsed_url.netloc.endswith('.digitaloceanspaces.com'):
+    if parsed_url.netloc.endswith('.digitaloceanspaces.com'):
         return endpoint_url.split('.')[1].split('.digitaloceanspaces.com')[0]
-    else:
-        # Additional cases can be added here
-        return None
-  
+
+    return None
+
 
 # --------------------------- S3 Bucket Connection --------------------------- #
 def get_boto_client(
@@ -72,7 +76,7 @@ def get_boto_client(
     if endpoint_url and access_key_id and secret_access_key:
         # Extract region from the endpoint URL
         region = extract_region_from_url(endpoint_url)
-        
+
         boto_client = bucket_session.client(
             's3',
             endpoint_url=endpoint_url,

tests/test_serverless/test_utils/test_upload.py

@@ -62,6 +62,36 @@ def test_get_boto_client(self):
                 region_name=None
             )
 
+            creds_s3 = bucket_creds.copy()
+            creds_s3['endpointUrl'] = "https://bucket-name.s3.region-code.amazonaws.com/key-name"
+
+            boto_client, transfer_config = get_boto_client(creds_s3)
+
+            mock_session.return_value.client.assert_called_with(
+                's3',
+                endpoint_url=creds_s3['endpointUrl'],
+                aws_access_key_id=bucket_creds['accessId'],
+                aws_secret_access_key=bucket_creds['accessSecret'],
+                config=unittest.mock.ANY,
+                region_name="region-code"
+            )
+
+            creds_do = bucket_creds.copy()
+            creds_do['endpointUrl'] = "https://name.region-code.digitaloceanspaces.com/key-name"
+
+            boto_client, transfer_config = get_boto_client(creds_do)
+
+            mock_session.return_value.client.assert_called_with(
+                's3',
+                endpoint_url=creds_do['endpointUrl'],
+                aws_access_key_id=bucket_creds['accessId'],
+                aws_secret_access_key=bucket_creds['accessSecret'],
+                config=unittest.mock.ANY,
+                region_name="region-code"
+            )
+
+
+
     def test_get_boto_client_environ(self):
         '''
         Tests get_boto_client with environment variables

tests/test_serverless/test_worker.py

@@ -9,7 +9,7 @@
 
 import runpod
 from runpod.serverless.modules.rp_logger import RunPodLogger
-
+from runpod.serverless import _signal_handler
 
 nest_asyncio.apply()
 
@@ -93,6 +93,17 @@ def test_local_api(self):
 
             assert mock_fastapi.WorkerAPI.called
 
+    @patch('runpod.serverless.log')
+    @patch('runpod.serverless.sys.exit')
+    def test_signal_handler(self, mock_exit, mock_logger):
+        '''
+        Test signal handler.
+        '''
+
+        _signal_handler(None, None)
+
+        assert mock_exit.called
+        assert mock_logger.info.called
 
 class TestWorkerTestInput(IsolatedAsyncioTestCase):
     """ Tests for runpod | serverless| worker """