diff --git a/mlos_bench/mlos_bench/services/base_fileshare.py b/mlos_bench/mlos_bench/services/base_fileshare.py index 7acf8a83ea0..83595aad1fd 100644 --- a/mlos_bench/mlos_bench/services/base_fileshare.py +++ b/mlos_bench/mlos_bench/services/base_fileshare.py @@ -9,6 +9,7 @@ import logging from abc import ABCMeta, abstractmethod +from typing import Any, Dict, Optional from mlos_bench.services.base_service import Service from mlos_bench.services.types.fileshare_type import SupportsFileShareOps @@ -21,7 +22,9 @@ class FileShareService(Service, SupportsFileShareOps, metaclass=ABCMeta): An abstract base of all file shares. """ - def __init__(self, config: dict, parent: Service): + def __init__(self, config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new file share with a given config. @@ -31,10 +34,12 @@ def __init__(self, config: dict, parent: Service): Free-format dictionary that contains the file share configuration. It will be passed as a constructor parameter of the class specified by `class_name`. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self.register([ self.download, diff --git a/mlos_bench/mlos_bench/services/base_service.py b/mlos_bench/mlos_bench/services/base_service.py index 85ba9049945..b5a8ed9fa0b 100644 --- a/mlos_bench/mlos_bench/services/base_service.py +++ b/mlos_bench/mlos_bench/services/base_service.py @@ -9,7 +9,7 @@ import json import logging -from typing import Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union from mlos_bench.services.types.config_loader_type import SupportsConfigLoading from mlos_bench.util import instantiate_from_config @@ -23,7 +23,11 @@ class Service: """ @classmethod - def new(cls, class_name: str, config: dict, parent: Optional["Service"]) -> "Service": + def new(cls, + class_name: str, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional["Service"] = None) -> "Service": """ Factory method for a new service with a given config. @@ -37,6 +41,8 @@ def new(cls, class_name: str, config: dict, parent: Optional["Service"]) -> "Ser Free-format dictionary that contains the service configuration. It will be passed as a constructor parameter of the class specified by `class_name`. + global_config : dict + Free-format dictionary of global parameters. parent : Service A parent service that can provide mixin functions. @@ -45,9 +51,12 @@ def new(cls, class_name: str, config: dict, parent: Optional["Service"]) -> "Ser svc : Service An instance of the `Service` class initialized with `config`. """ - return instantiate_from_config(cls, class_name, config, parent) + return instantiate_from_config(cls, class_name, config, global_config, parent) - def __init__(self, config: Optional[dict] = None, parent: Optional["Service"] = None): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional["Service"] = None): """ Create a new service with a given config. @@ -57,6 +66,8 @@ def __init__(self, config: Optional[dict] = None, parent: Optional["Service"] = Free-format dictionary that contains the service configuration. It will be passed as a constructor parameter of the class specified by `class_name`. + global_config : dict + Free-format dictionary of global parameters. parent : Service An optional parent service that can provide mixin functions. """ @@ -72,12 +83,14 @@ def __init__(self, config: Optional[dict] = None, parent: Optional["Service"] = self._config_loader_service = parent if _LOG.isEnabledFor(logging.DEBUG): - _LOG.debug("Service: %s Config:\n%s", - self.__class__.__name__, json.dumps(self.config, indent=2)) - _LOG.debug("Service: %s Parent mixins: %s", - self.__class__.__name__, + _LOG.debug("Service: %s Config:\n%s", self, json.dumps(self.config, indent=2)) + _LOG.debug("Service: %s Globals:\n%s", self, json.dumps(global_config or {}, indent=2)) + _LOG.debug("Service: %s Parent mixins: %s", self, [] if parent is None else list(parent._services.keys())) + def __repr__(self) -> str: + return self.__class__.__name__ + @property def config_loader_service(self) -> SupportsConfigLoading: """ diff --git a/mlos_bench/mlos_bench/services/config_persistence.py b/mlos_bench/mlos_bench/services/config_persistence.py index 218e40792fa..ea5b8ec69df 100644 --- a/mlos_bench/mlos_bench/services/config_persistence.py +++ b/mlos_bench/mlos_bench/services/config_persistence.py @@ -43,7 +43,9 @@ class ConfigPersistenceService(Service, SupportsConfigLoading): BUILTIN_CONFIG_PATH = str(files("mlos_bench.config").joinpath("")).replace("\\", "/") - def __init__(self, config: Optional[Dict[str, Any]] = None, + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, parent: Optional[Service] = None): """ Create a new instance of config persistence service. @@ -53,10 +55,12 @@ def __init__(self, config: Optional[Dict[str, Any]] = None, config : dict Free-format dictionary that contains parameters for the service. (E.g., root path for config files, etc.) + global_config : dict + Free-format dictionary of global parameters. parent : Service An optional parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self._config_path: List[str] = self.config.get("config_path", []) self._config_loader_service = self @@ -312,7 +316,7 @@ def _build_standalone_service(self, config: Dict[str, Any], An instance of the `Service` class initialized with `config`. """ (svc_class, svc_config) = self.prepare_class_load(config, global_config) - service = Service.new(svc_class, svc_config, parent) + service = Service.new(svc_class, svc_config, global_config, parent) _LOG.info("Created service: %s", service) return service @@ -483,7 +487,7 @@ def load_services(self, json_file_names: Iterable[str], """ _LOG.info("Load services: %s parent: %s", json_file_names, parent.__class__.__name__) - service = Service(global_config, parent) + service = Service({}, global_config, parent) for fname in json_file_names: config = self.load_config(fname, ConfigSchema.SERVICE) service.register(self.build_service(config, global_config, service).export()) diff --git a/mlos_bench/mlos_bench/services/local/local_exec.py b/mlos_bench/mlos_bench/services/local/local_exec.py index eb3f9ff5065..53354561a83 100644 --- a/mlos_bench/mlos_bench/services/local/local_exec.py +++ b/mlos_bench/mlos_bench/services/local/local_exec.py @@ -13,7 +13,7 @@ import subprocess import sys -from typing import Dict, Iterable, Mapping, Optional, Tuple, TYPE_CHECKING +from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, TYPE_CHECKING from mlos_bench.services.base_service import Service from mlos_bench.services.local.temp_dir_context import TempDirContextService @@ -32,7 +32,10 @@ class LocalExecService(TempDirContextService, SupportsLocalExec): due to reduced dependency management complications vs the target environment. """ - def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = None): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of a service to run scripts locally. @@ -41,10 +44,12 @@ def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = No config : dict Free-format dictionary that contains parameters for the service. (E.g., root path for config files, etc.) + global_config : dict + Free-format dictionary of global parameters. parent : Service An optional parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self.register([self.local_exec]) def local_exec(self, script_lines: Iterable[str], diff --git a/mlos_bench/mlos_bench/services/local/temp_dir_context.py b/mlos_bench/mlos_bench/services/local/temp_dir_context.py index 401bbee2ac5..cbfdb70e902 100644 --- a/mlos_bench/mlos_bench/services/local/temp_dir_context.py +++ b/mlos_bench/mlos_bench/services/local/temp_dir_context.py @@ -10,7 +10,7 @@ import logging from contextlib import nullcontext from tempfile import TemporaryDirectory -from typing import Optional, Union +from typing import Any, Dict, Optional, Union from mlos_bench.services.base_service import Service @@ -26,7 +26,10 @@ class TempDirContextService(Service, metaclass=abc.ABCMeta): This class is not supposed to be used as a standalone service. """ - def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = None): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of a service that provides temporary directory context for local exec service. @@ -36,10 +39,12 @@ def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = No config : dict Free-format dictionary that contains parameters for the service. (E.g., root path for config files, etc.) + global_config : dict + Free-format dictionary of global parameters. parent : Service An optional parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self._temp_dir = self.config.get("temp_dir") self.register([self.temp_dir_context]) 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 0c2763993cb..4b02cbc2eab 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_auth.py @@ -10,6 +10,7 @@ import json import logging import subprocess +from typing import Any, Dict, Optional from mlos_bench.services.base_service import Service from mlos_bench.services.types.authenticator_type import SupportsAuth @@ -24,7 +25,10 @@ class AzureAuthService(Service, SupportsAuth): _REQ_INTERVAL = 300 # = 5 min - def __init__(self, config: dict, parent: Service): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of Azure authentication services proxy. @@ -33,16 +37,18 @@ def __init__(self, config: dict, parent: Service): config : dict Free-format dictionary that contains the benchmark environment configuration. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) # Register methods that we want to expose to the Environment objects. self.register([self.get_access_token]) # This parameter can come from command line as strings, so conversion is needed. - self._req_interval = float(config.get("tokenRequestInterval", self._REQ_INTERVAL)) + self._req_interval = float(self.config.get("tokenRequestInterval", self._REQ_INTERVAL)) self._access_token = "RENEW *NOW*" self._token_expiration_ts = datetime.datetime.now() # Typically, some future timestamp. 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 e6a1d7701f4..f0d4617d1be 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_fileshare.py @@ -9,7 +9,7 @@ import os import logging -from typing import Set +from typing import Any, Dict, Optional, Set from azure.storage.fileshare import ShareClient @@ -27,7 +27,10 @@ class AzureFileShareService(FileShareService): _SHARE_URL = "https://{account_name}.file.core.windows.net/{fs_name}" - def __init__(self, config: dict, parent: Service): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new file share Service for Azure environments with a given config. @@ -37,13 +40,15 @@ def __init__(self, config: dict, parent: Service): Free-format dictionary that contains the file share configuration. It will be passed as a constructor parameter of the class specified by `class_name`. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) check_required_params( - config, { + self.config, { "storageAccountName", "storageFileShareName", "storageAccountKey", @@ -52,10 +57,10 @@ def __init__(self, config: dict, parent: Service): self._share_client = ShareClient.from_share_url( AzureFileShareService._SHARE_URL.format( - account_name=config["storageAccountName"], - fs_name=config["storageFileShareName"], + account_name=self.config["storageAccountName"], + fs_name=self.config["storageFileShareName"], ), - credential=config["storageAccountKey"], + credential=self.config["storageAccountKey"], ) def download(self, remote_path: str, local_path: str, recursive: bool = True) -> None: diff --git a/mlos_bench/mlos_bench/services/remote/azure/azure_services.py b/mlos_bench/mlos_bench/services/remote/azure/azure_services.py index ff15df9b7aa..5dd5349ecc3 100644 --- a/mlos_bench/mlos_bench/services/remote/azure/azure_services.py +++ b/mlos_bench/mlos_bench/services/remote/azure/azure_services.py @@ -10,7 +10,7 @@ import time import logging -from typing import Callable, Iterable, Tuple +from typing import Any, Callable, Dict, Iterable, Optional, Tuple import requests @@ -103,7 +103,10 @@ class AzureVMService(Service, SupportsVMOps, SupportsRemoteExec): "?api-version=2022-03-01" ) - def __init__(self, config: dict, parent: Service): + def __init__(self, + config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of Azure services proxy. @@ -112,13 +115,15 @@ def __init__(self, config: dict, parent: Service): config : dict Free-format dictionary that contains the benchmark environment configuration. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) check_required_params( - config, { + self.config, { "subscription", "resourceGroup", "deploymentName", @@ -141,18 +146,18 @@ def __init__(self, config: dict, parent: Service): ]) # These parameters can come from command line as strings, so conversion is needed. - self._poll_interval = float(config.get("pollInterval", self._POLL_INTERVAL)) - self._poll_timeout = float(config.get("pollTimeout", self._POLL_TIMEOUT)) - self._request_timeout = float(config.get("requestTimeout", self._REQUEST_TIMEOUT)) + self._poll_interval = float(self.config.get("pollInterval", self._POLL_INTERVAL)) + self._poll_timeout = float(self.config.get("pollTimeout", self._POLL_TIMEOUT)) + self._request_timeout = float(self.config.get("requestTimeout", self._REQUEST_TIMEOUT)) # TODO: Provide external schema validation? template = self.config_loader_service.load_config( - config['deploymentTemplatePath'], schema_type=None) + self.config['deploymentTemplatePath'], schema_type=None) assert template is not None and isinstance(template, dict) self._deploy_template = template self._deploy_params = merge_parameters( - dest=config['deploymentTemplateParameters'].copy(), source=self.config) + dest=self.config['deploymentTemplateParameters'].copy(), source=global_config) def _get_headers(self) -> dict: """ diff --git a/mlos_bench/mlos_bench/tests/services/local/mock/mock_local_exec_service.py b/mlos_bench/mlos_bench/tests/services/local/mock/mock_local_exec_service.py index 852256f5fed..503c5e1956d 100644 --- a/mlos_bench/mlos_bench/tests/services/local/mock/mock_local_exec_service.py +++ b/mlos_bench/mlos_bench/tests/services/local/mock/mock_local_exec_service.py @@ -7,7 +7,7 @@ """ import logging -from typing import Iterable, Mapping, Optional, Tuple, TYPE_CHECKING +from typing import Any, Dict, Iterable, Mapping, Optional, Tuple, TYPE_CHECKING from mlos_bench.services.base_service import Service from mlos_bench.services.local.temp_dir_context import TempDirContextService @@ -24,8 +24,10 @@ class MockLocalExecService(TempDirContextService, SupportsLocalExec): Mock methods for LocalExecService testing. """ - def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = None): - super().__init__(config, parent) + def __init__(self, config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): + super().__init__(config, global_config, parent) self.register([self.local_exec]) def local_exec(self, script_lines: Iterable[str], 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 89bc633251e..1b161174fbc 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py +++ b/mlos_bench/mlos_bench/tests/services/remote/azure/conftest.py @@ -30,7 +30,7 @@ def azure_auth_service(config_persistence_service: ConfigPersistenceService, """ Creates a dummy AzureAuthService for tests that require it. """ - auth = AzureAuthService(config={}, parent=config_persistence_service) + auth = AzureAuthService(config={}, global_config={}, parent=config_persistence_service) monkeypatch.setattr(auth, "get_access_token", lambda: "TEST_TOKEN") return auth @@ -51,7 +51,7 @@ def azure_vm_service(azure_auth_service: AzureAuthService) -> AzureVMService: "vmName": "test-vm", # Should come from the upper-level config "pollInterval": 1, "pollTimeout": 2 - }, parent=azure_auth_service) + }, global_config={}, parent=azure_auth_service) @pytest.fixture @@ -64,4 +64,4 @@ def azure_fileshare(config_persistence_service: ConfigPersistenceService) -> Azu "storageAccountName": "TEST_ACCOUNT_NAME", "storageFileShareName": "TEST_FS_NAME", "storageAccountKey": "TEST_ACCOUNT_KEY" - }, parent=config_persistence_service) + }, global_config={}, parent=config_persistence_service) diff --git a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_fileshare_service.py b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_fileshare_service.py index 1030051dc18..d50fcf19c3c 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_fileshare_service.py +++ b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_fileshare_service.py @@ -7,7 +7,9 @@ """ import logging +from typing import Any, Dict, Optional +from mlos_bench.services.base_service import Service from mlos_bench.services.base_fileshare import FileShareService from mlos_bench.services.types.fileshare_type import SupportsFileShareOps @@ -19,8 +21,10 @@ class MockFileShareService(FileShareService, SupportsFileShareOps): A collection Service functions for mocking file share ops. """ - def __init__(self, config: dict, parent: FileShareService): - super().__init__(config, parent) + def __init__(self, config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): + super().__init__(config, global_config, parent) self.register([ self.download, diff --git a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_remote_exec_service.py b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_remote_exec_service.py index fe0384d88c8..c7005471171 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_remote_exec_service.py +++ b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_remote_exec_service.py @@ -6,6 +6,8 @@ A collection Service functions for mocking remote script execution. """ +from typing import Any, Dict, Optional + from mlos_bench.services.base_service import Service from mlos_bench.services.types.remote_exec_type import SupportsRemoteExec from mlos_bench.tests.services.remote.mock import mock_operation @@ -16,7 +18,9 @@ class MockRemoteExecService(Service, SupportsRemoteExec): Mock remote script execution service. """ - def __init__(self, config: dict, parent: Service): + def __init__(self, config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of mock remote exec service. @@ -25,10 +29,12 @@ def __init__(self, config: dict, parent: Service): config : dict Free-format dictionary that contains the benchmark environment configuration. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self.register({ "remote_exec": mock_operation, "get_remote_exec_results": mock_operation, diff --git a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_vm_service.py b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_vm_service.py index c885f0fdf5e..2490c900d33 100644 --- a/mlos_bench/mlos_bench/tests/services/remote/mock/mock_vm_service.py +++ b/mlos_bench/mlos_bench/tests/services/remote/mock/mock_vm_service.py @@ -6,6 +6,8 @@ A collection Service functions for mocking managing VMs. """ +from typing import Any, Dict, Optional + from mlos_bench.services.base_service import Service from mlos_bench.services.types.vm_provisioner_type import SupportsVMOps from mlos_bench.tests.services.remote.mock import mock_operation @@ -16,7 +18,9 @@ class MockVMService(Service, SupportsVMOps): Mock VM service for testing. """ - def __init__(self, config: dict, parent: Service): + def __init__(self, config: Optional[Dict[str, Any]] = None, + global_config: Optional[Dict[str, Any]] = None, + parent: Optional[Service] = None): """ Create a new instance of mock VM services proxy. @@ -25,10 +29,12 @@ def __init__(self, config: dict, parent: Service): config : dict Free-format dictionary that contains the benchmark environment configuration. + global_config : dict + Free-format dictionary of global parameters. parent : Service Parent service that can provide mixin functions. """ - super().__init__(config, parent) + super().__init__(config, global_config, parent) self.register({ name: mock_operation for name in ( "wait_vm_deployment",