From 2a6cd0d34ff5e76a94c6d5acd943d496f4501561 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 8 Jul 2024 16:05:00 -0700 Subject: [PATCH 01/16] use token-based credential instead of a storage key in Azure file share service --- .../services/remote/azure/azure_fileshare.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index af09f4c7230..2f8a1a3a3f5 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -16,6 +16,7 @@ from mlos_bench.services.base_service import Service from mlos_bench.services.base_fileshare import FileShareService +from mlos_bench.services.types.authenticator_type import SupportsAuth from mlos_bench.util import check_required_params _LOG = logging.getLogger(__name__) @@ -58,7 +59,6 @@ def __init__(self, self.config, { "storageAccountName", "storageFileShareName", - "storageAccountKey", } ) @@ -67,9 +67,18 @@ def __init__(self, account_name=self.config["storageAccountName"], fs_name=self.config["storageFileShareName"], ), - credential=self.config["storageAccountKey"], + credential=self._get_access_token(), + token_intent="backup", ) + def _get_access_token(self) -> str: + """ + Get the access token for Azure file share. + """ + assert self._parent is not None and isinstance(self._parent, SupportsAuth), \ + "Authorization service not provided. Include service-auth.jsonc?" + return self._parent.get_access_token() + def download(self, params: dict, remote_path: str, local_path: str, recursive: bool = True) -> None: super().download(params, remote_path, local_path, recursive) dir_client = self._share_client.get_directory_client(remote_path) From 951fdbad4fca82c5051921b1bf53a0ccdce1f98a Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 8 Jul 2024 16:31:46 -0700 Subject: [PATCH 02/16] initialize auth service before azure fileshare in unit tests --- mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py index 45feb1aa503..db4d5a7d569 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py @@ -96,7 +96,7 @@ def azure_vm_service_remote_exec_only(azure_auth_service: AzureAuthService) -> A @pytest.fixture -def azure_fileshare(config_persistence_service: ConfigPersistenceService) -> AzureFileShareService: +def azure_fileshare(azure_auth_service: AzureAuthService) -> AzureFileShareService: """ Creates a dummy AzureFileShareService for tests that require it. """ @@ -105,4 +105,4 @@ def azure_fileshare(config_persistence_service: ConfigPersistenceService) -> Azu "storageAccountName": "TEST_ACCOUNT_NAME", "storageFileShareName": "TEST_FS_NAME", "storageAccountKey": "TEST_ACCOUNT_KEY" - }, global_config={}, parent=config_persistence_service) + }, global_config={}, parent=azure_auth_service) From 9aeaa92decf4d86b24def42745d57e420f829ce5 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 8 Jul 2024 17:38:55 -0700 Subject: [PATCH 03/16] late initialization for azure fileshare _share_client; proper monkey patching of the class for unit tests --- .../services/remote/azure/azure_fileshare.py | 38 +++++++++---------- .../remote/azure/azure_fileshare_test.py | 30 +++++++-------- 2 files changed, 34 insertions(+), 34 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 2f8a1a3a3f5..80cd2d1a4d8 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -54,34 +54,34 @@ def __init__(self, config, global_config, parent, self.merge_methods(methods, [self.upload, self.download]) ) - check_required_params( self.config, { "storageAccountName", "storageFileShareName", } ) + self._share_client: Optional[ShareClient] = None - self._share_client = ShareClient.from_share_url( - AzureFileShareService._SHARE_URL.format( - account_name=self.config["storageAccountName"], - fs_name=self.config["storageFileShareName"], - ), - credential=self._get_access_token(), - token_intent="backup", - ) - - def _get_access_token(self) -> str: + def _get_share_client(self) -> ShareClient: """ - Get the access token for Azure file share. + Get the Azure file share client object. """ - assert self._parent is not None and isinstance(self._parent, SupportsAuth), \ - "Authorization service not provided. Include service-auth.jsonc?" - return self._parent.get_access_token() + if self._share_client is None: + assert self._parent is not None and isinstance(self._parent, SupportsAuth), \ + "Authorization service not provided. Include service-auth.jsonc?" + self._share_client = ShareClient.from_share_url( + AzureFileShareService._SHARE_URL.format( + account_name=self.config["storageAccountName"], + fs_name=self.config["storageFileShareName"], + ), + credential=self._parent.get_access_token(), + token_intent="backup", + ) + return self._share_client def download(self, params: dict, remote_path: str, local_path: str, recursive: bool = True) -> None: super().download(params, remote_path, local_path, recursive) - dir_client = self._share_client.get_directory_client(remote_path) + dir_client = self._get_share_client().get_directory_client(remote_path) if dir_client.exists(): os.makedirs(local_path, exist_ok=True) for content in dir_client.list_directories_and_files(): @@ -94,7 +94,7 @@ def download(self, params: dict, remote_path: str, local_path: str, recursive: b # Ensure parent folders exist folder, _ = os.path.split(local_path) os.makedirs(folder, exist_ok=True) - file_client = self._share_client.get_file_client(remote_path) + file_client = self._get_share_client().get_file_client(remote_path) try: data = file_client.download_file() with open(local_path, "wb") as output_file: @@ -145,7 +145,7 @@ def _upload(self, local_path: str, remote_path: str, recursive: bool, seen: Set[ # Ensure parent folders exist folder, _ = os.path.split(remote_path) self._remote_makedirs(folder) - file_client = self._share_client.get_file_client(remote_path) + file_client = self._get_share_client().get_file_client(remote_path) with open(local_path, "rb") as file_data: _LOG.debug("Upload file: %s -> %s", local_path, remote_path) file_client.upload_file(file_data) @@ -165,6 +165,6 @@ def _remote_makedirs(self, remote_path: str) -> None: if not folder: continue path += folder + "/" - dir_client = self._share_client.get_directory_client(path) + dir_client = self._get_share_client().get_directory_client(path) if not dir_client.exists(): dir_client.create_directory() diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py index 199c42f1fb0..e270952c3fc 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py @@ -24,9 +24,9 @@ def test_download_file(mock_makedirs: MagicMock, mock_open: MagicMock, azure_fil local_folder = "some/local/folder" remote_path = f"{remote_folder}/{filename}" local_path = f"{local_folder}/{filename}" - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access config: dict = {} - with patch.object(mock_share_client, "get_file_client") as mock_get_file_client, \ + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_file_client") as mock_get_file_client, \ patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client: mock_get_directory_client.return_value = Mock(exists=Mock(return_value=False)) @@ -75,9 +75,9 @@ def test_download_folder_non_recursive(mock_makedirs: MagicMock, remote_folder = "a/remote/folder" local_folder = "some/local/folder" dir_client_returns = make_dir_client_returns(remote_folder) - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access config: dict = {} - with patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, \ + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, \ patch.object(mock_share_client, "get_file_client") as mock_get_file_client: mock_get_directory_client.side_effect = lambda x: dir_client_returns[x] @@ -99,9 +99,9 @@ def test_download_folder_recursive(mock_makedirs: MagicMock, mock_open: MagicMoc remote_folder = "a/remote/folder" local_folder = "some/local/folder" dir_client_returns = make_dir_client_returns(remote_folder) - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access config: dict = {} - with patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, \ + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_directory_client") as mock_get_directory_client, \ patch.object(mock_share_client, "get_file_client") as mock_get_file_client: mock_get_directory_client.side_effect = lambda x: dir_client_returns[x] @@ -127,11 +127,11 @@ def test_upload_file(mock_isdir: MagicMock, mock_open: MagicMock, azure_fileshar local_folder = "some/local/folder" remote_path = f"{remote_folder}/{filename}" local_path = f"{local_folder}/{filename}" - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access mock_isdir.return_value = False - config: dict = {} - with patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + config: dict = {} + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_file_client") as mock_get_file_client: azure_fileshare.upload(config, local_path, remote_path) mock_get_file_client.assert_called_with(remote_path) @@ -196,10 +196,10 @@ def test_upload_directory_non_recursive(mock_scandir: MagicMock, isdir_returns = make_isdir_returns(local_folder) mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access - config: dict = {} - with patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + config: dict = {} + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_file_client") as mock_get_file_client: azure_fileshare.upload(config, local_folder, remote_folder, recursive=False) mock_get_file_client.assert_called_with(f"{remote_folder}/a_file_1.csv") @@ -218,10 +218,10 @@ def test_upload_directory_recursive(mock_scandir: MagicMock, isdir_returns = make_isdir_returns(local_folder) mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] - mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access - config: dict = {} - with patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + config: dict = {} + with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ + patch.object(mock_share_client, "get_file_client") as mock_get_file_client: azure_fileshare.upload(config, local_folder, remote_folder, recursive=True) mock_get_file_client.assert_has_calls([ From 54c308bc061b80223b7890889eecc39f11ccb285 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 8 Jul 2024 18:11:57 -0700 Subject: [PATCH 04/16] remove storageAccountKey parameter from azure_fileshare fixture --- mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py index db4d5a7d569..8ff22aa422e 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py @@ -104,5 +104,4 @@ def azure_fileshare(azure_auth_service: AzureAuthService) -> AzureFileShareServi return AzureFileShareService(config={ "storageAccountName": "TEST_ACCOUNT_NAME", "storageFileShareName": "TEST_FS_NAME", - "storageAccountKey": "TEST_ACCOUNT_KEY" }, global_config={}, parent=azure_auth_service) From ead96a103a45839fe7d3d2d5b41bd586f61c78a4 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Fri, 12 Jul 2024 15:31:22 -0700 Subject: [PATCH 05/16] black formatting updates --- .../services/remote/azure/azure_fileshare.py | 7 ++--- .../remote/azure/azure_fileshare_test.py | 27 +++++++++---------- 2 files changed, 16 insertions(+), 18 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 7b2ae0fdd90..4a33bc896a1 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -58,7 +58,7 @@ def __init__( { "storageAccountName", "storageFileShareName", - } + }, ) self._share_client: Optional[ShareClient] = None @@ -67,8 +67,9 @@ def _get_share_client(self) -> ShareClient: Get the Azure file share client object. """ if self._share_client is None: - assert self._parent is not None and isinstance(self._parent, SupportsAuth), \ - "Authorization service not provided. Include service-auth.jsonc?" + assert self._parent is not None and isinstance( + self._parent, SupportsAuth + ), "Authorization service not provided. Include service-auth.jsonc?" self._share_client = ShareClient.from_share_url( AzureFileShareService._SHARE_URL.format( account_name=self.config["storageAccountName"], diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py index 18bfe6425e8..f3acb5d77e7 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py @@ -28,9 +28,7 @@ def test_download_file( local_path = f"{local_folder}/{filename}" config: dict = {} - with patch.object( - azure_fileshare, "_share_client" - ) as mock_share_client, patch.object( + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( mock_share_client, "get_file_client" ) as mock_get_file_client, patch.object( mock_share_client, "get_directory_client" @@ -88,9 +86,7 @@ def test_download_folder_non_recursive( mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access config: dict = {} - with patch.object( - azure_fileshare, "_share_client" - ) as mock_share_client, patch.object( + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( mock_share_client, "get_directory_client" ) as mock_get_directory_client, patch.object( mock_share_client, "get_file_client" @@ -123,9 +119,7 @@ def test_download_folder_recursive( local_folder = "some/local/folder" dir_client_returns = make_dir_client_returns(remote_folder) config: dict = {} - with patch.object( - azure_fileshare, "_share_client" - ) as mock_share_client, patch.object( + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( mock_share_client, "get_directory_client" ) as mock_get_directory_client, patch.object( mock_share_client, "get_file_client" @@ -167,8 +161,9 @@ def test_upload_file( mock_isdir.return_value = False config: dict = {} - with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ - patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( + mock_share_client, "get_file_client" + ) as mock_get_file_client: azure_fileshare.upload(config, local_path, remote_path) mock_get_file_client.assert_called_with(remote_path) @@ -240,8 +235,9 @@ def test_upload_directory_non_recursive( config: dict = {} config: dict = {} - with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ - patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( + mock_share_client, "get_file_client" + ) as mock_get_file_client: azure_fileshare.upload(config, local_folder, remote_folder, recursive=False) mock_get_file_client.assert_called_with(f"{remote_folder}/a_file_1.csv") @@ -266,8 +262,9 @@ def test_upload_directory_recursive( config: dict = {} config: dict = {} - with patch.object(azure_fileshare, "_share_client") as mock_share_client, \ - patch.object(mock_share_client, "get_file_client") as mock_get_file_client: + with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( + mock_share_client, "get_file_client" + ) as mock_get_file_client: azure_fileshare.upload(config, local_folder, remote_folder, recursive=True) mock_get_file_client.assert_has_calls( From 8078d7071b34b0e1772ba09432d487eb31e6e83e Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Fri, 19 Jul 2024 15:50:30 -0700 Subject: [PATCH 06/16] docformatter fixes --- .../mlos_bench/services/remote/azure/azure_fileshare.py | 4 +--- mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 4a33bc896a1..169ec454290 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -63,9 +63,7 @@ def __init__( self._share_client: Optional[ShareClient] = None def _get_share_client(self) -> ShareClient: - """ - Get the Azure file share client object. - """ + """Get the Azure file share client object.""" if self._share_client is None: assert self._parent is not None and isinstance( self._parent, SupportsAuth diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py index 5f635843762..37d554c8972 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py @@ -103,9 +103,7 @@ def azure_vm_service_remote_exec_only(azure_auth_service: AzureAuthService) -> A @pytest.fixture def azure_fileshare(azure_auth_service: AzureAuthService) -> AzureFileShareService: - """ - Creates a dummy AzureFileShareService for tests that require it. - """ + """Creates a dummy AzureFileShareService for tests that require it.""" with patch("mlos_bench.services.remote.azure.azure_fileshare.ShareClient"): return AzureFileShareService( config={ From f2b47c6393e0b6413ea9f1bf2f58373dabf0ed1d Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 22 Jul 2024 14:13:18 -0700 Subject: [PATCH 07/16] formatting fixes --- mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 169ec454290..38e7a023215 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -12,8 +12,8 @@ from azure.storage.fileshare import ShareClient from mlos_bench.services.base_fileshare import FileShareService -from mlos_bench.services.types.authenticator_type import SupportsAuth from mlos_bench.services.base_service import Service +from mlos_bench.services.types.authenticator_type import SupportsAuth from mlos_bench.util import check_required_params _LOG = logging.getLogger(__name__) From e9e48f2db6600f4938e0a5fbe28a428332617a89 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 22 Jul 2024 15:17:11 -0700 Subject: [PATCH 08/16] remove git merge artifact --- .../tests/services/remote/azure/azure_fileshare_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py index f3acb5d77e7..0111178fb35 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py @@ -259,7 +259,6 @@ def test_upload_directory_recursive( mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access - config: dict = {} config: dict = {} with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( From 2925ebb592b8e3773735f62730942ac94cf1d508 Mon Sep 17 00:00:00 2001 From: Sergiy Matusevych Date: Mon, 22 Jul 2024 15:18:10 -0700 Subject: [PATCH 09/16] more git merge fixes --- .../tests/services/remote/azure/azure_fileshare_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py index 0111178fb35..d2aecc22755 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/azure_fileshare_test.py @@ -118,6 +118,7 @@ def test_download_folder_recursive( remote_folder = "a/remote/folder" local_folder = "some/local/folder" dir_client_returns = make_dir_client_returns(remote_folder) + config: dict = {} with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( mock_share_client, "get_directory_client" @@ -232,7 +233,6 @@ def test_upload_directory_non_recursive( mock_scandir.side_effect = lambda x: scandir_returns[process_paths(x)] mock_isdir.side_effect = lambda x: isdir_returns[process_paths(x)] mock_share_client = azure_fileshare._share_client # pylint: disable=protected-access - config: dict = {} config: dict = {} with patch.object(azure_fileshare, "_share_client") as mock_share_client, patch.object( From 00ff3bfbf8b368205f30eeb4a652604b42e3cc91 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 23 Jul 2024 22:48:05 +0000 Subject: [PATCH 10/16] move sanity check into __init__ --- .../services/remote/azure/azure_fileshare.py | 6 +++--- .../services/test_load_service_config_examples.py | 14 +++++++++++++- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 38e7a023215..78cf4ca6ddd 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -60,14 +60,14 @@ def __init__( "storageFileShareName", }, ) + assert self._parent is not None and isinstance( + self._parent, SupportsAuth + ), "Authorization service not provided. Include service-auth.jsonc?" self._share_client: Optional[ShareClient] = None def _get_share_client(self) -> ShareClient: """Get the Azure file share client object.""" if self._share_client is None: - assert self._parent is not None and isinstance( - self._parent, SupportsAuth - ), "Authorization service not provided. Include service-auth.jsonc?" self._share_client = ShareClient.from_share_url( AzureFileShareService._SHARE_URL.format( account_name=self.config["storageAccountName"], diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index 55453270808..6a92586ed27 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -48,11 +48,23 @@ def test_load_service_config_examples( config_path: str, ) -> None: """Tests loading a config example.""" + parent: Service = config_loader_service config = config_loader_service.load_config(config_path, ConfigSchema.SERVICE) + if config.get("class", "").split(".")[-1] == "AzureFileShareService": + # AzureFileShareService requires an auth service to be loaded as well. + auth_service_config = config_loader_service.load_config( + "services/remote/azure/service-auth.jsonc", + ConfigSchema.SERVICE, + ) + auth_service = config_loader_service.build_service( + config=auth_service_config, + parent=config_loader_service, + ) + parent = auth_service # Make an instance of the class based on the config. service_inst = config_loader_service.build_service( config=config, - parent=config_loader_service, + parent=parent, ) assert service_inst is not None assert isinstance(service_inst, Service) From 7895fe50e80390c356edb5407b1f58d21a4d78ab Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 24 Jul 2024 20:25:17 +0000 Subject: [PATCH 11/16] use the mock auth service --- .../tests/config/services/test_load_service_config_examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index 6a92586ed27..b3e5f88fcd6 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -53,7 +53,7 @@ def test_load_service_config_examples( if config.get("class", "").split(".")[-1] == "AzureFileShareService": # AzureFileShareService requires an auth service to be loaded as well. auth_service_config = config_loader_service.load_config( - "services/remote/azure/service-auth.jsonc", + "services/remote/mock/mock_auth_service.jsonc", ConfigSchema.SERVICE, ) auth_service = config_loader_service.build_service( From 72b6c15089ac842a02d52f46e35aa5019a93e828 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 25 Jul 2024 20:06:26 +0000 Subject: [PATCH 12/16] refactor --- .../mlos_bench/services/remote/azure/azure_fileshare.py | 2 +- .../config/services/test_load_service_config_examples.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index df3ef454010..e0597841311 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -62,7 +62,7 @@ def __init__( ) assert self._parent is not None and isinstance( self._parent, SupportsAuth - ), "Authorization service not provided. Include service-auth.jsonc?" + ), "Authorization service not provided. Include services/remote/azure/service-auth.jsonc?" self._auth_service: SupportsAuth = self._parent self._share_client: Optional[ShareClient] = None # Cache of the last access token we used to access the file share. diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index b3e5f88fcd6..e010fd140b9 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -50,7 +50,12 @@ def test_load_service_config_examples( """Tests loading a config example.""" parent: Service = config_loader_service config = config_loader_service.load_config(config_path, ConfigSchema.SERVICE) - if config.get("class", "").split(".")[-1] == "AzureFileShareService": + # Add other services that require a SupportsAuth parent service as necessary. + requires_auth_service_parent = { + "AzureFileShareService", + } + config_class_name = str(config.get("class", "MISSING CLASS")).rsplit(".", maxsplit=1)[-1] + if config_class_name in requires_auth_service_parent: # AzureFileShareService requires an auth service to be loaded as well. auth_service_config = config_loader_service.load_config( "services/remote/mock/mock_auth_service.jsonc", From a33de94ea0156d21c2e675da6490de881ba6165d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 25 Jul 2024 20:12:16 +0000 Subject: [PATCH 13/16] comments --- mlos_bench/mlos_bench/services/remote/azure/azure_auth.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py b/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py index e674a943256..c18bc6f2d32 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py @@ -108,6 +108,7 @@ def _init_sp(self) -> None: cert_bytes = b64decode(secret.value) # Reauthenticate as the service principal. + # TODO: Add managed identity support as well. self._cred = azure_id.CertificateCredential( tenant_id=tenant_id, client_id=sp_client_id, From 941a1d60e87954a783a71181f5911ac43e37ea19 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 8 Aug 2024 22:00:10 +0000 Subject: [PATCH 14/16] tweaks to make the config error load time checkable --- .../services/remote/azure/azure_fileshare.py | 23 +++++++++++-------- .../test_load_service_config_examples.py | 12 ++++++---- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index da636fb3d04..12206173559 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -6,7 +6,7 @@ import logging import os -from typing import Any, Callable, Dict, List, Optional, Set, Union +from typing import Any, Callable, Dict, List, Optional, Set, Union, get_type_hints from azure.core.credentials import TokenCredential from azure.core.exceptions import ResourceNotFoundError @@ -62,22 +62,25 @@ def __init__( }, ) # TODO: Ensure that the parent service is an authentication service that provides a TokenCredential. - assert self._parent is not None and isinstance( - self._parent, SupportsAuth - ), "Authorization service not provided. Include services/remote/azure/service-auth.jsonc?" + assert ( + self._parent is not None + and isinstance(self._parent, SupportsAuth) + and get_type_hints(self._parent.get_credential).get("return") == TokenCredential + ), ( + "Azure Authorization service not provided. " + "Include services/remote/azure/service-auth.jsonc?" + ) self._auth_service: SupportsAuth[TokenCredential] = self._parent self._share_client: Optional[ShareClient] = None - # Cache of the last access token we used to access the file share. - self._access_token: Optional[str] = None def _get_share_client(self) -> ShareClient: """Get the Azure file share client object.""" if self._share_client is None: credential = self._auth_service.get_credential() - assert isinstance( - credential, TokenCredential - ), (f"Expected a TokenCredential, but got {type(credential)} instead. " - "Include services/remote/azure/service-auth.jsonc?") + assert isinstance(credential, TokenCredential), ( + f"Expected a TokenCredential, but got {type(credential)} instead. " + "Include services/remote/azure/service-auth.jsonc?" + ) self._share_client = ShareClient.from_share_url( self._SHARE_URL.format( account_name=self.config["storageAccountName"], diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index e010fd140b9..2c2e74e3fb5 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -50,15 +50,18 @@ def test_load_service_config_examples( """Tests loading a config example.""" parent: Service = config_loader_service config = config_loader_service.load_config(config_path, ConfigSchema.SERVICE) + # Add other services that require a SupportsAuth parent service as necessary. + # AzureFileShareService requires an AzureAuth service to be loaded as well. + # mock_auth_service_config = "services/remote/mock/mock_auth_service.jsonc" + azure_auth_service_config = "services/remote/azure/service-auth.jsonc" requires_auth_service_parent = { - "AzureFileShareService", + "AzureFileShareService": azure_auth_service_config, } config_class_name = str(config.get("class", "MISSING CLASS")).rsplit(".", maxsplit=1)[-1] - if config_class_name in requires_auth_service_parent: - # AzureFileShareService requires an auth service to be loaded as well. + if auth_service_config_path := requires_auth_service_parent.get(config_class_name): auth_service_config = config_loader_service.load_config( - "services/remote/mock/mock_auth_service.jsonc", + auth_service_config_path, ConfigSchema.SERVICE, ) auth_service = config_loader_service.build_service( @@ -66,6 +69,7 @@ def test_load_service_config_examples( parent=config_loader_service, ) parent = auth_service + # Make an instance of the class based on the config. service_inst = config_loader_service.build_service( config=config, From c81e9f5f90d2ab705a2d04f9c2702426b0a5814d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 8 Aug 2024 22:01:13 +0000 Subject: [PATCH 15/16] format --- mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index 12206173559..f0f69b134a9 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -61,7 +61,8 @@ def __init__( "storageFileShareName", }, ) - # TODO: Ensure that the parent service is an authentication service that provides a TokenCredential. + # Ensure that the parent service is an authentication service that provides + # a TokenCredential. assert ( self._parent is not None and isinstance(self._parent, SupportsAuth) From c38c6113b5b61dba8279692af50021e7916031d5 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 8 Aug 2024 22:03:13 +0000 Subject: [PATCH 16/16] fixup --- mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py index f0f69b134a9..67dae5009bf 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -68,7 +68,7 @@ def __init__( and isinstance(self._parent, SupportsAuth) and get_type_hints(self._parent.get_credential).get("return") == TokenCredential ), ( - "Azure Authorization service not provided. " + "Azure Authentication service not provided. " "Include services/remote/azure/service-auth.jsonc?" ) self._auth_service: SupportsAuth[TokenCredential] = self._parent