various cosmetic fixes to please pylint and flake8; shorter version of mock service implementations

mlos_bench/mlos_bench/services/config_persistence.py

@@ -9,7 +9,6 @@
 """
 
 import os
-import re
 import sys
 
 import json    # For logging only

mlos_bench/mlos_bench/services/local/local_exec.py

@@ -6,18 +6,17 @@
 Helper functions to run scripts and commands locally on the scheduler side.
 """
 
-import contextlib
 import errno
 import logging
 import os
 import shlex
 import subprocess
 import sys
-import tempfile
 
-from typing import Dict, Iterable, Mapping, Optional, Tuple, Union, TYPE_CHECKING
+from typing import 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
 from mlos_bench.services.types.local_exec_type import SupportsLocalExec
 
 if TYPE_CHECKING:
@@ -26,7 +25,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-class LocalExecService(Service, SupportsLocalExec):
+class LocalExecService(TempDirContextService, SupportsLocalExec):
     """
     Collection of methods to run scripts and commands in an external process
     on the node acting as the scheduler. Can be useful for data processing
@@ -46,29 +45,7 @@ def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = No
             An optional parent service that can provide mixin functions.
         """
         super().__init__(config, parent)
-        self._temp_dir = self.config.get("temp_dir")
-        self.register([
-            self.temp_dir_context,
-            self.local_exec,
-        ])
-
-    def temp_dir_context(self, path: Optional[str] = None) -> Union[tempfile.TemporaryDirectory, contextlib.nullcontext]:
-        """
-        Create a temp directory or use the provided path.
-
-        Parameters
-        ----------
-        path : str
-            A path to the temporary directory. Create a new one if None.
-
-        Returns
-        -------
-        temp_dir_context : TemporaryDirectory
-            Temporary directory context to use in the `with` clause.
-        """
-        if path is None and self._temp_dir is None:
-            return tempfile.TemporaryDirectory()
-        return contextlib.nullcontext(path or self._temp_dir)
+        self.register([self.local_exec])
 
     def local_exec(self, script_lines: Iterable[str],
                    env: Optional[Mapping[str, "TunableValue"]] = None,

mlos_bench/mlos_bench/services/local/temp_dir_context.py

@@ -0,0 +1,62 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""
+Helper functions to work with temp files locally on the scheduler side.
+"""
+
+import abc
+import logging
+from contextlib import nullcontext
+from tempfile import TemporaryDirectory
+from typing import Optional, Union
+
+from mlos_bench.services.base_service import Service
+
+_LOG = logging.getLogger(__name__)
+
+
+class TempDirContextService(Service, metaclass=abc.ABCMeta):
+    """
+    A *base* service class that provides a method to create a temporary
+    directory context for local scripts.
+
+    It is inherited by LocalExecService and MockLocalExecService.
+    This class is not supposed to be used as a standalone service.
+    """
+
+    def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = None):
+        """
+        Create a new instance of a service that provides temporary directory context
+        for local exec service.
+
+        Parameters
+        ----------
+        config : dict
+            Free-format dictionary that contains parameters for the service.
+            (E.g., root path for config files, etc.)
+        parent : Service
+            An optional parent service that can provide mixin functions.
+        """
+        super().__init__(config, parent)
+        self._temp_dir = self.config.get("temp_dir")
+        self.register([self.temp_dir_context])
+
+    def temp_dir_context(self, path: Optional[str] = None) -> Union[TemporaryDirectory, nullcontext]:
+        """
+        Create a temp directory or use the provided path.
+
+        Parameters
+        ----------
+        path : str
+            A path to the temporary directory. Create a new one if None.
+
+        Returns
+        -------
+        temp_dir_context : TemporaryDirectory
+            Temporary directory context to use in the `with` clause.
+        """
+        if path is None and self._temp_dir is None:
+            return TemporaryDirectory()
+        return nullcontext(path or self._temp_dir)

mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py

@@ -76,9 +76,7 @@ def load_environment_config_examples(config_loader_service: ConfigPersistenceSer
     for mock_service_config_path in mock_service_configs:
         mock_service_config = config_loader_service.load_config(mock_service_config_path)
         config_loader_service.register(config_loader_service.build_service(
-                                       config=mock_service_config,
-                                       parent=config_loader_service,
-                                      ).export())
+            config=mock_service_config, parent=config_loader_service).export())
 
     envs = config_loader_service.load_environment_list(
         config_path, tunable_groups, global_config, service=config_loader_service)

mlos_bench/mlos_bench/tests/services/local/mock/mock_local_exec_service.py

@@ -6,13 +6,11 @@
 A collection Service functions for mocking local exec.
 """
 
-import contextlib
 import logging
-import tempfile
-
-from typing import Iterable, Mapping, Optional, Tuple, Union, TYPE_CHECKING
+from typing import Iterable, Mapping, Optional, Tuple, TYPE_CHECKING
 
 from mlos_bench.services.base_service import Service
+from mlos_bench.services.local.temp_dir_context import TempDirContextService
 from mlos_bench.services.types.local_exec_type import SupportsLocalExec
 
 if TYPE_CHECKING:
@@ -21,72 +19,17 @@
 _LOG = logging.getLogger(__name__)
 
 
-class MockLocalExecService(Service, SupportsLocalExec):
+class MockLocalExecService(TempDirContextService, SupportsLocalExec):
     """
     Mock methods for LocalExecService testing.
     """
 
     def __init__(self, config: Optional[dict] = None, parent: Optional[Service] = None):
-        """
-        Create a new instance of a service to run scripts locally.
-
-        Parameters
-        ----------
-        config : dict
-            Free-format dictionary that contains parameters for the service.
-            (E.g., root path for config files, etc.)
-        parent : Service
-            An optional parent service that can provide mixin functions.
-        """
         super().__init__(config, parent)
-        self._temp_dir = self.config.get("temp_dir")
-        self.register([
-            self.temp_dir_context,
-            self.local_exec,
-        ])
-
-    def temp_dir_context(self, path: Optional[str] = None) -> Union[tempfile.TemporaryDirectory, contextlib.nullcontext]:
-        """
-        Create a temp directory or use the provided path.
-
-        Parameters
-        ----------
-        path : str
-            A path to the temporary directory. Create a new one if None.
-
-        Returns
-        -------
-        temp_dir_context : TemporaryDirectory
-            Temporary directory context to use in the `with` clause.
-        """
-        if path is None and self._temp_dir is None:
-            return tempfile.TemporaryDirectory()
-        return contextlib.nullcontext(path or self._temp_dir)
+        self.register([self.local_exec])
 
     def local_exec(self, script_lines: Iterable[str],
                    env: Optional[Mapping[str, "TunableValue"]] = None,
                    cwd: Optional[str] = None,
                    return_on_error: bool = False) -> Tuple[int, str, str]:
-        """
-        Execute the script lines from `script_lines` in a local process.
-
-        Parameters
-        ----------
-        script_lines : Iterable[str]
-            Lines of the script to run locally.
-            Treat every line as a separate command to run.
-        env : Mapping[str, Union[int, float, str]]
-            Environment variables (optional).
-        cwd : str
-            Work directory to run the script at.
-            If omitted, use `temp_dir` or create a temporary dir.
-        return_on_error : bool
-            If True, stop running script lines on first non-zero return code.
-            The default is False.
-
-        Returns
-        -------
-        (return_code, stdout, stderr) : (int, str, str)
-            A 3-tuple of return code, stdout, and stderr of the script process.
-        """
         return (0, "", "")

mlos_bench/mlos_bench/tests/services/remote/mock/__init__.py

@@ -5,3 +5,19 @@
 """
 Mock remote services for testing purposes.
 """
+
+from typing import Any, Tuple
+
+from mlos_bench.environments.status import Status
+
+
+def mock_operation(*_args: Any, **_kwargs: Any) -> Tuple[Status, dict]:
+    """
+    Mock VM operation that always succeeds.
+
+    Returns
+    -------
+    result : (Status, dict)
+        A pair of Status and result, always (SUCCEEDED, {}).
+    """
+    return Status.SUCCEEDED, {}

mlos_bench/mlos_bench/tests/services/remote/mock/mock_remote_exec_service.py

@@ -6,16 +6,9 @@
 A collection Service functions for mocking remote script execution.
 """
 
-import logging
-
-from typing import Iterable, Tuple
-
-from mlos_bench.environments.status import Status
 from mlos_bench.services.base_service import Service
 from mlos_bench.services.types.remote_exec_type import SupportsRemoteExec
-from mlos_bench.util import check_required_params
-
-_LOG = logging.getLogger(__name__)
+from mlos_bench.tests.services.remote.mock import mock_operation
 
 
 class MockRemoteExecService(Service, SupportsRemoteExec):
@@ -36,55 +29,7 @@ def __init__(self, config: dict, parent: Service):
             Parent service that can provide mixin functions.
         """
         super().__init__(config, parent)
-
-        check_required_params(
-            config, {
-                "vmName",
-            }
-        )
-
-        # Register methods that we want to expose to the Environment objects.
-        self.register([
-            self.remote_exec,
-            self.get_remote_exec_results,
-        ])
-
-    def remote_exec(self, script: Iterable[str], params: dict) -> Tuple[Status, dict]:
-        """
-        Run a command on remote host OS.
-
-        Parameters
-        ----------
-        script : Iterable[str]
-            A list of lines to execute as a script on a remote VM.
-        params : dict
-            Flat dictionary of (key, value) pairs of parameters.
-            They usually come from `const_args` and `tunable_params`
-            properties of the Environment.
-
-        Returns
-        -------
-        result : (Status, dict)
-            A pair of Status and result.
-            Status is one of {PENDING, SUCCEEDED, FAILED}
-        """
-        return (Status.SUCCEEDED, {})
-
-    def get_remote_exec_results(self, params: dict) -> Tuple[Status, dict]:
-        """
-        Get the results of the asynchronously running command.
-
-        Parameters
-        ----------
-        params : dict
-            Flat dictionary of (key, value) pairs of tunable parameters.
-            Must have the "asyncResultsUrl" key to get the results.
-            If the key is not present, return Status.PENDING.
-
-        Returns
-        -------
-        result : (Status, dict)
-            A pair of Status and result.
-            Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT}
-        """
-        return (Status.SUCCEEDED, {})
+        self.register({
+            "remote_exec": mock_operation,
+            "get_remote_exec_results": mock_operation,
+        })