From 302633cd70494cc993fbde42667fe07c68fa0469 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 15:46:19 -0400 Subject: [PATCH 1/9] Update rp_upload.py --- runpod/serverless/utils/rp_upload.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/runpod/serverless/utils/rp_upload.py b/runpod/serverless/utils/rp_upload.py index f56484ae..9f4c05e1 100644 --- a/runpod/serverless/utils/rp_upload.py +++ b/runpod/serverless/utils/rp_upload.py @@ -25,6 +25,9 @@ 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: @@ -32,10 +35,9 @@ def extract_region_from_url(endpoint_url): # DigitalOcean Spaces-like URL elif 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 +74,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, From ea83a32c4e678bd7b1a58a66216a9f96de349faa Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 15:53:53 -0400 Subject: [PATCH 2/9] Update __init__.py --- runpod/serverless/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index f72d33bc..a7cbe659 100644 --- a/runpod/serverless/__init__.py +++ b/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) From 719baa03acccac5e6e6a7b1f4396f412707ac28c Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 15:55:08 -0400 Subject: [PATCH 3/9] Update rp_upload.py --- runpod/serverless/utils/rp_upload.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/runpod/serverless/utils/rp_upload.py b/runpod/serverless/utils/rp_upload.py index 9f4c05e1..1d9b82ec 100644 --- a/runpod/serverless/utils/rp_upload.py +++ b/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,7 +19,7 @@ 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" @@ -32,8 +33,9 @@ def extract_region_from_url(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] return None From d4ca6821a04d057950aaf9d0580854984cbd5847 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 16:13:39 -0400 Subject: [PATCH 4/9] Update test_worker.py --- tests/test_serverless/test_worker.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index 31f0b9f3..0a68b081 100644 --- a/tests/test_serverless/test_worker.py +++ b/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,16 @@ def test_local_api(self): assert mock_fastapi.WorkerAPI.called + @patch('runpod.serverless.RunPodLogger') + @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 """ From 6b3e22d3de98dc8ff6b6347482d76990678973a5 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 16:17:30 -0400 Subject: [PATCH 5/9] Update test_worker.py --- tests/test_serverless/test_worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index 0a68b081..a94a67e5 100644 --- a/tests/test_serverless/test_worker.py +++ b/tests/test_serverless/test_worker.py @@ -93,12 +93,13 @@ def test_local_api(self): assert mock_fastapi.WorkerAPI.called - @patch('runpod.serverless.RunPodLogger') + @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 From 58bcf6c64f4d816ad40e75fcef47a9c80c6953db Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 16:39:50 -0400 Subject: [PATCH 6/9] Update test_upload.py --- .../test_serverless/test_utils/test_upload.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/test_serverless/test_utils/test_upload.py b/tests/test_serverless/test_utils/test_upload.py index ef3b0e28..84bc7606 100644 --- a/tests/test_serverless/test_utils/test_upload.py +++ b/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_once_with( + 's3', + endpoint_url=bucket_creds['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_once_with( + 's3', + endpoint_url=bucket_creds['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 From 4b610a93dd56739d42d48a8d42d94af6533cef37 Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 16:43:37 -0400 Subject: [PATCH 7/9] Update test_upload.py --- tests/test_serverless/test_utils/test_upload.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_serverless/test_utils/test_upload.py b/tests/test_serverless/test_utils/test_upload.py index 84bc7606..79af8263 100644 --- a/tests/test_serverless/test_utils/test_upload.py +++ b/tests/test_serverless/test_utils/test_upload.py @@ -67,7 +67,7 @@ def test_get_boto_client(self): boto_client, transfer_config = get_boto_client(creds_s3) - mock_session.return_value.client.assert_called_once_with( + mock_session.return_value.client.assert_called_with( 's3', endpoint_url=bucket_creds['endpointUrl'], aws_access_key_id=bucket_creds['accessId'], @@ -81,7 +81,7 @@ def test_get_boto_client(self): boto_client, transfer_config = get_boto_client(creds_do) - mock_session.return_value.client.assert_called_once_with( + mock_session.return_value.client.assert_called_with( 's3', endpoint_url=bucket_creds['endpointUrl'], aws_access_key_id=bucket_creds['accessId'], From d8f6e1a6c079582bd7e3efb638ceeb9b810a3bcc Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Tue, 19 Sep 2023 16:46:37 -0400 Subject: [PATCH 8/9] Update test_upload.py --- tests/test_serverless/test_utils/test_upload.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_serverless/test_utils/test_upload.py b/tests/test_serverless/test_utils/test_upload.py index 79af8263..619be95f 100644 --- a/tests/test_serverless/test_utils/test_upload.py +++ b/tests/test_serverless/test_utils/test_upload.py @@ -69,7 +69,7 @@ def test_get_boto_client(self): mock_session.return_value.client.assert_called_with( 's3', - endpoint_url=bucket_creds['endpointUrl'], + endpoint_url=creds_s3['endpointUrl'], aws_access_key_id=bucket_creds['accessId'], aws_secret_access_key=bucket_creds['accessSecret'], config=unittest.mock.ANY, @@ -83,7 +83,7 @@ def test_get_boto_client(self): mock_session.return_value.client.assert_called_with( 's3', - endpoint_url=bucket_creds['endpointUrl'], + endpoint_url=creds_do['endpointUrl'], aws_access_key_id=bucket_creds['accessId'], aws_secret_access_key=bucket_creds['accessSecret'], config=unittest.mock.ANY, From 786faf4d982296eb8a7ce13af8b0ae9ea5d4d7ab Mon Sep 17 00:00:00 2001 From: Justin Merrell Date: Wed, 20 Sep 2023 13:22:55 -0400 Subject: [PATCH 9/9] Update CI-e2e.yml --- .github/workflows/CI-e2e.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/CI-e2e.yml b/.github/workflows/CI-e2e.yml index 3b6cbbd8..3c813bfd 100644 --- a/.github/workflows/CI-e2e.yml +++ b/.github/workflows/CI-e2e.yml @@ -64,7 +64,7 @@ jobs: steps: - uses: actions/checkout@v4 - name: Run Tests - uses: direlines/runpod-test-runner@v1.6 + uses: direlines/runpod-test-runner@v1.7 with: image-tag: ${{ vars.DOCKERHUB_REPO }}/${{ vars.DOCKERHUB_IMG }}:${{ needs.e2e-build.outputs.docker_tag }} runpod-api-key: ${{ secrets.RUNPOD_API_KEY }}