From 751a778762536c523d03b98ecacd0c6a03be5f04 Mon Sep 17 00:00:00 2001 From: kevin Date: Thu, 27 Nov 2025 18:35:22 +0800 Subject: [PATCH] add aksk check --- .../engine/sched/resource_manager_v1.py | 9 ++- fastdeploy/utils.py | 15 ++++- tests/utils/test_download.py | 66 ++++++++++++++++++- 3 files changed, 87 insertions(+), 3 deletions(-) diff --git a/fastdeploy/engine/sched/resource_manager_v1.py b/fastdeploy/engine/sched/resource_manager_v1.py index 9f4873112d1..6fca0323f63 100644 --- a/fastdeploy/engine/sched/resource_manager_v1.py +++ b/fastdeploy/engine/sched/resource_manager_v1.py @@ -821,7 +821,14 @@ def download_bos_features(bos_client, features_urls): return None if self.bos_client is None: - self.bos_client = init_bos_client() + try: + self.bos_client = init_bos_client() + except Exception as e: + error_msg = f"request {request.request_id} init bos client error: {str(e)}" + llm_logger.error(error_msg) + request.error_message = error_msg + request.error_code = 540 + return None inputs = request.multimodal_inputs if inputs.get("video_feature_urls") is not None and len(inputs["video_feature_urls"]) > 0: diff --git a/fastdeploy/utils.py b/fastdeploy/utils.py index 771e7165258..fcb5648c115 100644 --- a/fastdeploy/utils.py +++ b/fastdeploy/utils.py @@ -969,13 +969,26 @@ def check_download_links(bos_client, links, timeout=1): def init_bos_client(): from baidubce.auth.bce_credentials import BceCredentials from baidubce.bce_client_configuration import BceClientConfiguration + from baidubce.exception import BceHttpClientError, BceServerError from baidubce.services.bos.bos_client import BosClient cfg = BceClientConfiguration( credentials=BceCredentials(envs.ENCODE_FEATURE_BOS_AK, envs.ENCODE_FEATURE_BOS_SK), endpoint=envs.ENCODE_FEATURE_ENDPOINT, ) - return BosClient(cfg) + + try: + client = BosClient(cfg) + client.list_buckets(max_keys=1) + except BceServerError as e: + if e.status_code == 403: + raise Exception("BOS authentication failed: Invalid AK/SK") from e + raise Exception(f"BOS connection failed: {str(e)}") from e + except BceHttpClientError as e: + raise Exception(f"Invalid BOS endpoint configuration: {str(e)}") from e + except Exception as e: + raise Exception(f"BOS client validation error: {str(e)}") from e + return client def download_from_bos(bos_client, bos_links, retry: int = 0): diff --git a/tests/utils/test_download.py b/tests/utils/test_download.py index 52e364f2cc7..50b8e99c07d 100644 --- a/tests/utils/test_download.py +++ b/tests/utils/test_download.py @@ -17,7 +17,7 @@ import os import unittest -from fastdeploy.utils import retrive_model_from_server +from fastdeploy.utils import init_bos_client, retrive_model_from_server class TestAistudioDownload(unittest.TestCase): @@ -67,5 +67,69 @@ def test_retrive_model_from_huggingface_server_model_not_exist(self): os.environ.clear() +class TestInitBosClient(unittest.TestCase): + """ + Test cases for initializing Baidu Object Storage (BOS) client using FastDeploy utilities. + """ + + def test_init_bos_client_success(self): + """ + Test case for successful initialization of BOS client with valid environment variables. + """ + from unittest.mock import MagicMock, patch + + # Mock BosClient and related dependencies to skip validation + with patch("baidubce.auth.bce_credentials.BceCredentials") as mock_credentials: + with patch("baidubce.bce_client_configuration.BceClientConfiguration") as mock_config: + with patch("baidubce.services.bos.bos_client.BosClient") as mock_bos_client: + # Additional mock to make list_buckets call succeed + mock_bos_client_instance = MagicMock() + # Mock the list_buckets call to return empty buckets list + mock_bos_client_instance.list_buckets = MagicMock(return_value=MagicMock()) + mock_bos_client.return_value = mock_bos_client_instance + + # Mock the credentials and config + mock_credentials_instance = MagicMock() + mock_credentials.return_value = mock_credentials_instance + + mock_config_instance = MagicMock() + mock_config.return_value = mock_config_instance + + # Set up valid environment variables + os.environ["ENCODE_FEATURE_BOS_AK"] = "test_access_key" + os.environ["ENCODE_FEATURE_BOS_SK"] = "test_secret_key" + os.environ["ENCODE_FEATURE_ENDPOINT"] = "http://test.endpoint.com" + + try: + # Call the function + client = init_bos_client() + + # Verify that BosClient was created with correct arguments + mock_credentials.assert_called_once_with("test_access_key", "test_secret_key") + mock_config.assert_called_once_with( + credentials=mock_credentials_instance, endpoint="http://test.endpoint.com" + ) + mock_bos_client.assert_called_once_with(mock_config_instance) + + # Verify that the returned client is the mock instance + self.assertEqual(client, mock_bos_client_instance) + finally: + os.environ.clear() + + def test_init_bos_client_missing_envs(self): + """ + Test case for initializing BOS client when necessary environment variables are missing. + """ + # Test with empty environment variables + os.environ["ENCODE_FEATURE_BOS_AK"] = "" + os.environ["ENCODE_FEATURE_BOS_SK"] = "" + os.environ["ENCODE_FEATURE_ENDPOINT"] = "" + + with self.assertRaises(Exception) as context: + init_bos_client() + self.assertIn("BOS client validation error", str(context.exception)) + os.environ.clear() + + if __name__ == "__main__": unittest.main()