Merge branch 'main' into pre-commit-parallelism-tweaks

.vscode/settings.json

@@ -167,5 +167,9 @@
         "--log-level=DEBUG",
         "."
     ],
-    "python.testing.unittestEnabled": false
+    "python.testing.unittestEnabled": false,
+    "debugpy.debugJustMyCode": false,
+    "python.analysis.autoImportCompletions": true,
+    "python.analysis.supportRestructuredText": true,
+    "python.analysis.typeCheckingMode": "standard"
 }

mlos_bench/mlos_bench/environments/base_environment.py

@@ -8,6 +8,7 @@
 import json
 import logging
 from collections.abc import Iterable, Sequence
+from contextlib import AbstractContextManager as ContextManager
 from datetime import datetime
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Literal
@@ -28,7 +29,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-class Environment(metaclass=abc.ABCMeta):
+class Environment(ContextManager, metaclass=abc.ABCMeta):
     # pylint: disable=too-many-instance-attributes
     """An abstract base of all benchmark environments."""
 

mlos_bench/mlos_bench/environments/local/local_fileshare_env.py

@@ -91,7 +91,7 @@ def _template_from_to(self, config_key: str) -> list[tuple[Template, Template]]:
     def _expand(
         from_to: Iterable[tuple[Template, Template]],
         params: Mapping[str, TunableValue],
-    ) -> Generator[tuple[str, str], None, None]:
+    ) -> Generator[tuple[str, str]]:
         """
         Substitute $var parameters in from/to path templates.
 

mlos_bench/mlos_bench/optimizers/base_optimizer.py

@@ -9,6 +9,7 @@
 import logging
 from abc import ABCMeta, abstractmethod
 from collections.abc import Sequence
+from contextlib import AbstractContextManager as ContextManager
 from types import TracebackType
 from typing import Literal
 
@@ -25,7 +26,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-class Optimizer(metaclass=ABCMeta):  # pylint: disable=too-many-instance-attributes
+class Optimizer(ContextManager, metaclass=ABCMeta):  # pylint: disable=too-many-instance-attributes
     """An abstract interface between the benchmarking framework and mlos_core
     optimizers.
     """

mlos_bench/mlos_bench/schedulers/base_scheduler.py

@@ -7,6 +7,7 @@
 import json
 import logging
 from abc import ABCMeta, abstractmethod
+from contextlib import AbstractContextManager as ContextManager
 from datetime import datetime
 from types import TracebackType
 from typing import Any, Literal
@@ -23,7 +24,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-class Scheduler(metaclass=ABCMeta):
+class Scheduler(ContextManager, metaclass=ABCMeta):
     # pylint: disable=too-many-instance-attributes
     """Base class for the optimization loop scheduling policies."""
 

mlos_bench/mlos_bench/services/base_service.py

@@ -9,6 +9,7 @@
 import json
 import logging
 from collections.abc import Callable
+from contextlib import AbstractContextManager as ContextManager
 from types import TracebackType
 from typing import Any, Literal
 
@@ -19,7 +20,7 @@
 _LOG = logging.getLogger(__name__)
 
 
-class Service:
+class Service(ContextManager):
     """An abstract base of all Environment Services and used to build up mix-ins."""
 
     @classmethod

mlos_bench/mlos_bench/services/config_persistence.py

@@ -67,7 +67,7 @@ def __init__(
             Free-format dictionary of global parameters.
         parent : Service
             An optional parent service that can provide mixin functions.
-        methods : Union[dict[str, Callable], list[Callable], None]
+        methods : dict[str, Callable] | list[Callable] | None
             New methods to register with the service.
         """
         super().__init__(
@@ -166,7 +166,7 @@ def load_config(
 
         Returns
         -------
-        config : Union[dict, list[dict]]
+        config : dict | list[dict]
             Free-format dictionary that contains the configuration.
         """
         assert isinstance(json, str)

mlos_bench/mlos_bench/storage/base_storage.py

@@ -25,6 +25,7 @@
 import logging
 from abc import ABCMeta, abstractmethod
 from collections.abc import Iterator
+from contextlib import AbstractContextManager as ContextManager
 from datetime import datetime
 from types import TracebackType
 from typing import Any, Literal
@@ -132,7 +133,7 @@ def experiment(  # pylint: disable=too-many-arguments
             the results of the experiment and related data.
         """
 
-    class Experiment(metaclass=ABCMeta):
+    class Experiment(ContextManager, metaclass=ABCMeta):
         # pylint: disable=too-many-instance-attributes
         """
         Base interface for storing the results of the experiment.

mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py

@@ -51,7 +51,7 @@ def ssh_test_server(
     ssh_test_server_hostname: str,
     docker_compose_project_name: str,
     locked_docker_services: DockerServices,
-) -> Generator[SshTestServerInfo, None, None]:
+) -> Generator[SshTestServerInfo]:
     """
     Fixture for getting the ssh test server services setup via docker-compose using
     pytest-docker.

mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_fileshare.py

@@ -23,7 +23,7 @@
 
 
 @contextmanager
-def closeable_temp_file(**kwargs: Any) -> Generator[_TemporaryFileWrapper, None, None]:
+def closeable_temp_file(**kwargs: Any) -> Generator[_TemporaryFileWrapper]:
     """
     Provides a context manager for a temporary file that can be closed and still
     unlinked.

mlos_bench/mlos_bench/tests/storage/sql/fixtures.py

@@ -38,7 +38,7 @@ def storage() -> SqlStorage:
 def exp_storage(
     storage: SqlStorage,
     tunable_groups: TunableGroups,
-) -> Generator[SqlStorage.Experiment, None, None]:
+) -> Generator[SqlStorage.Experiment]:
     """
     Test fixture for Experiment using in-memory SQLite3 storage.
 
@@ -60,7 +60,7 @@ def exp_storage(
 @pytest.fixture
 def exp_no_tunables_storage(
     storage: SqlStorage,
-) -> Generator[SqlStorage.Experiment, None, None]:
+) -> Generator[SqlStorage.Experiment]:
     """
     Test fixture for Experiment using in-memory SQLite3 storage.
 
@@ -84,7 +84,7 @@ def exp_no_tunables_storage(
 def mixed_numerics_exp_storage(
     storage: SqlStorage,
     mixed_numerics_tunable_groups: TunableGroups,
-) -> Generator[SqlStorage.Experiment, None, None]:
+) -> Generator[SqlStorage.Experiment]:
     """
     Test fixture for an Experiment with mixed numerics tunables using in-memory SQLite3
     storage.

mlos_bench/mlos_bench/tunables/tunable_groups.py

@@ -179,13 +179,13 @@ def __setitem__(
         self._index[name][name] = value
         return self._index[name][name]
 
-    def __iter__(self) -> Generator[tuple[Tunable, CovariantTunableGroup], None, None]:
+    def __iter__(self) -> Generator[tuple[Tunable, CovariantTunableGroup]]:
         """
         An iterator over all tunables in the group.
 
         Returns
         -------
-        [(tunable, group), ...] : Generator[tuple[Tunable, CovariantTunableGroup], None, None]
+        [(tunable, group), ...] : Generator[tuple[Tunable, CovariantTunableGroup]]
             An iterator over all tunables in all groups. Each element is a 2-tuple
             of an instance of the Tunable parameter and covariant group it belongs to.
         """

mlos_core/mlos_core/tests/spaces/monkey_patch_quantization_test.py

@@ -93,7 +93,7 @@ def test_configspace_quant_repatch() -> None:
     new_meta[QUANTIZATION_BINS_META_KEY] = 21
     hp.meta = new_meta
     monkey_patch_hp_quantization(hp)
-    samples_set = set(hp.sample_value(100, seed=RandomState(SEED)))
+    samples_set: set[int] = set(hp.sample_value(100, seed=RandomState(SEED)))
     quantized_values_new = set(range(5, 96, 10))
     assert samples_set.issubset(set(range(0, 101, 5)))
     assert len(samples_set - quantized_values_new) < len(samples_set)

pyproject.toml

@@ -62,7 +62,6 @@ disable = [
     "consider-using-assignment-expr",
     "docstring-first-line-empty",
     "missing-raises-doc",
-    "unnecessary-default-type-args",      # affects Generator type hints, but we still support python 3.8
 ]
 
 [tool.pylint.string]