Pyright Improvements

.devcontainer/devcontainer.json

@@ -68,7 +68,6 @@
                 "huntertran.auto-markdown-toc",
                 "ibm.output-colorizer",
                 "lextudio.restructuredtext",
-                "matangover.mypy",
                 "ms-azuretools.vscode-docker",
                 "ms-python.black-formatter",
                 "ms-python.pylint",

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ ci:
     # Let pre-commit.ci automatically update PRs with formatting fixes.
     autofix_prs: true
     # skip local hooks - they should be managed manually via conda-envs/*.yml
-    skip: [mypy, pylint, pycodestyle]
+    skip: [mypy, pylint, pycodestyle, pyright]
     autoupdate_schedule: monthly
     autoupdate_commit_msg: |
         [pre-commit.ci] pre-commit autoupdate
@@ -15,6 +15,7 @@ ci:
         See Also:
         - https://github.com/microsoft/MLOS/blob/main/conda-envs/mlos.yml
         - https://pypi.org/project/mypy/
+        - https://pypi.org/project/pyright/
         - https://pypi.org/project/pylint/
         - https://pypi.org/project/pycodestyle/
 
@@ -140,6 +141,20 @@ repos:
             (?x)^(
                 doc/source/conf.py
             )$
+    -   id: pyright
+        name: pyright
+        entry: pyright
+        language: system
+        types: [python]
+        require_serial: true
+        exclude: |
+            (?x)^(
+                doc/source/conf.py|
+                mlos_core/setup.py|
+                mlos_bench/setup.py|
+                mlos_viz/setup.py|
+                conftest.py
+            )$
     -   id: mypy
         name: mypy
         entry: mypy

.vscode/extensions.json

@@ -12,7 +12,6 @@
         "huntertran.auto-markdown-toc",
         "ibm.output-colorizer",
         "lextudio.restructuredtext",
-        "matangover.mypy",
         "ms-azuretools.vscode-docker",
         "ms-python.black-formatter",
         "ms-python.pylint",

.vscode/settings.json

@@ -170,6 +170,5 @@
     "python.testing.unittestEnabled": false,
     "debugpy.debugJustMyCode": false,
     "python.analysis.autoImportCompletions": true,
-    "python.analysis.supportRestructuredText": true,
-    "python.analysis.typeCheckingMode": "standard"
+    "python.analysis.supportRestructuredText": true
 }

conda-envs/mlos.yml

@@ -28,6 +28,7 @@ dependencies:
     - pylint==3.3.3
     - tomlkit
     - mypy==1.14.1
+    - pyright==1.1.392.post0
     - pandas-stubs
     - types-beautifulsoup4
     - types-colorama

doc/source/conf.py

@@ -29,6 +29,10 @@
 from sphinx.application import Sphinx as SphinxApp
 from sphinx.environment import BuildEnvironment
 
+# Note: doc requirements aren't installed by default.
+# To install them, run `pip install -r doc/requirements.txt`
+
+
 sys.path.insert(0, os.path.abspath("../../mlos_core/mlos_core"))
 sys.path.insert(1, os.path.abspath("../../mlos_bench/mlos_bench"))
 sys.path.insert(1, os.path.abspath("../../mlos_viz/mlos_viz"))

mlos_bench/mlos_bench/environments/mock_env.py

@@ -64,8 +64,8 @@ def __init__(  # pylint: disable=too-many-arguments
         seed = int(self.config.get("mock_env_seed", -1))
         self._run_random = random.Random(seed or None) if seed >= 0 else None
         self._status_random = random.Random(seed or None) if seed >= 0 else None
-        self._range = self.config.get("mock_env_range")
-        self._metrics = self.config.get("mock_env_metrics", ["score"])
+        self._range: tuple[int, int] | None = self.config.get("mock_env_range")
+        self._metrics: list[str] | None = self.config.get("mock_env_metrics", ["score"])
         self._is_ready = True
 
     def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue]:
@@ -80,7 +80,7 @@ def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue
         if self._range:
             score = self._range[0] + score * (self._range[1] - self._range[0])
 
-        return {metric: score for metric in self._metrics}
+        return {metric: float(score) for metric in self._metrics or []}
 
     def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]:
         """

mlos_bench/mlos_bench/optimizers/mock_optimizer.py

@@ -32,7 +32,7 @@ def __init__(
         self._random: dict[str, Callable[[Tunable], TunableValue]] = {
             "categorical": lambda tunable: rnd.choice(tunable.categories),
             "float": lambda tunable: rnd.uniform(*tunable.range),
-            "int": lambda tunable: rnd.randint(*tunable.range),
+            "int": lambda tunable: rnd.randint(*(int(x) for x in tunable.range)),
         }
 
     def bulk_register(

mlos_bench/mlos_bench/services/base_service.py

@@ -14,6 +14,7 @@
 from typing import Any, Literal
 
 from mlos_bench.config.schemas import ConfigSchema
+from mlos_bench.services.types.bound_method import BoundMethod
 from mlos_bench.services.types.config_loader_type import SupportsConfigLoading
 from mlos_bench.util import instantiate_from_config
 
@@ -278,7 +279,7 @@ def register(self, services: dict[str, Callable] | list[Callable]) -> None:
             for _, svc_method in self._service_methods.items()
             # Note: some methods are actually stand alone functions, so we need
             # to filter them out.
-            if hasattr(svc_method, "__self__") and isinstance(svc_method.__self__, Service)
+            if isinstance(svc_method, BoundMethod) and isinstance(svc_method.__self__, Service)
         }
 
     def export(self) -> dict[str, Callable]:

mlos_bench/mlos_bench/services/remote/ssh/ssh_fileshare.py

@@ -26,6 +26,8 @@ class CopyMode(Enum):
 class SshFileShareService(FileShareService, SshService):
     """A collection of functions for interacting with SSH servers as file shares."""
 
+    # pylint: disable=too-many-ancestors
+
     async def _start_file_copy(
         self,
         params: dict,

mlos_bench/mlos_bench/services/remote/ssh/ssh_host_service.py

@@ -24,6 +24,7 @@
 class SshHostService(SshService, SupportsOSOps, SupportsRemoteExec):
     """Helper methods to manage machines via SSH."""
 
+    # pylint: disable=too-many-ancestors
     # pylint: disable=too-many-instance-attributes
 
     def __init__(

mlos_bench/mlos_bench/services/types/authenticator_type.py

@@ -14,6 +14,9 @@
 class SupportsAuth(Protocol[T_co]):
     """Protocol interface for authentication for the cloud services."""
 
+    # Needed by pyright
+    # pylint: disable=unnecessary-ellipsis,redundant-returns-doc
+
     def get_access_token(self) -> str:
         """
         Get the access token for cloud services.
@@ -23,6 +26,7 @@ def get_access_token(self) -> str:
         access_token : str
             Access token.
         """
+        ...
 
     def get_auth_headers(self) -> dict:
         """
@@ -33,6 +37,7 @@ def get_auth_headers(self) -> dict:
         access_header : dict
             HTTP header containing the access token.
         """
+        ...
 
     def get_credential(self) -> T_co:
         """
@@ -43,3 +48,4 @@ def get_credential(self) -> T_co:
         credential : T_co
             Cloud-specific credential object.
         """
+        ...

mlos_bench/mlos_bench/services/types/bound_method.py

@@ -0,0 +1,24 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Protocol representing a bound method."""
+
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class BoundMethod(Protocol):
+    """A callable method bound to an object."""
+
+    # pylint: disable=too-few-public-methods
+    # pylint: disable=unnecessary-ellipsis
+
+    @property
+    def __self__(self) -> Any:
+        """The self object of the bound method."""
+        ...
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the bound method."""
+        ...

mlos_bench/mlos_bench/services/types/config_loader_type.py

@@ -23,6 +23,9 @@
 class SupportsConfigLoading(Protocol):
     """Protocol interface for helper functions to lookup and load configs."""
 
+    # Needed by pyright
+    # pylint: disable=unnecessary-ellipsis,redundant-returns-doc
+
     def get_config_paths(self) -> list[str]:
         """
         Gets the list of config paths this service will search for config files.
@@ -31,6 +34,7 @@ def get_config_paths(self) -> list[str]:
         -------
         list[str]
         """
+        ...
 
     def resolve_path(self, file_path: str, extra_paths: Iterable[str] | None = None) -> str:
         """
@@ -49,6 +53,7 @@ def resolve_path(self, file_path: str, extra_paths: Iterable[str] | None = None)
         path : str
             An actual path to the config or script.
         """
+        ...
 
     def load_config(
         self,
@@ -71,6 +76,7 @@ def load_config(
         config : Union[dict, list[dict]]
             Free-format dictionary that contains the configuration.
         """
+        ...
 
     def build_environment(  # pylint: disable=too-many-arguments
         self,
@@ -108,6 +114,7 @@ def build_environment(  # pylint: disable=too-many-arguments
         env : Environment
             An instance of the `Environment` class initialized with `config`.
         """
+        ...
 
     def load_environment(
         self,
@@ -140,6 +147,7 @@ def load_environment(
         env : Environment
             A new benchmarking environment.
         """
+        ...
 
     def load_environment_list(
         self,
@@ -173,6 +181,7 @@ def load_environment_list(
         env : list[Environment]
             A list of new benchmarking environments.
         """
+        ...
 
     def load_services(
         self,
@@ -198,3 +207,4 @@ def load_services(
         service : Service
             A collection of service methods.
         """
+        ...

mlos_bench/mlos_bench/services/types/host_ops_type.py

@@ -14,6 +14,8 @@
 class SupportsHostOps(Protocol):
     """Protocol interface for Host/VM boot operations."""
 
+    # pylint: disable=unnecessary-ellipsis
+
     def start_host(self, params: dict) -> tuple["Status", dict]:
         """
         Start a Host/VM.
@@ -29,6 +31,7 @@ def start_host(self, params: dict) -> tuple["Status", dict]:
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...
 
     def stop_host(self, params: dict, force: bool = False) -> tuple["Status", dict]:
         """
@@ -47,6 +50,7 @@ def stop_host(self, params: dict, force: bool = False) -> tuple["Status", dict]:
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...
 
     def restart_host(self, params: dict, force: bool = False) -> tuple["Status", dict]:
         """
@@ -65,6 +69,7 @@ def restart_host(self, params: dict, force: bool = False) -> tuple["Status", dic
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...
 
     def wait_host_operation(self, params: dict) -> tuple["Status", dict]:
         """
@@ -85,3 +90,4 @@ def wait_host_operation(self, params: dict) -> tuple["Status", dict]:
             Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT}
             Result is info on the operation runtime if SUCCEEDED, otherwise {}.
         """
+        ...

mlos_bench/mlos_bench/services/types/host_provisioner_type.py

@@ -14,6 +14,8 @@
 class SupportsHostProvisioning(Protocol):
     """Protocol interface for Host/VM provisioning operations."""
 
+    # pylint: disable=unnecessary-ellipsis
+
     def provision_host(self, params: dict) -> tuple["Status", dict]:
         """
         Check if Host/VM is ready. Deploy a new Host/VM, if necessary.
@@ -31,6 +33,7 @@ def provision_host(self, params: dict) -> tuple["Status", dict]:
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...
 
     def wait_host_deployment(self, params: dict, *, is_setup: bool) -> tuple["Status", dict]:
         """
@@ -52,6 +55,7 @@ def wait_host_deployment(self, params: dict, *, is_setup: bool) -> tuple["Status
             Status is one of {PENDING, SUCCEEDED, FAILED, TIMED_OUT}
             Result is info on the operation runtime if SUCCEEDED, otherwise {}.
         """
+        ...
 
     def deprovision_host(self, params: dict) -> tuple["Status", dict]:
         """
@@ -68,6 +72,7 @@ def deprovision_host(self, params: dict) -> tuple["Status", dict]:
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...
 
     def deallocate_host(self, params: dict) -> tuple["Status", dict]:
         """
@@ -88,3 +93,4 @@ def deallocate_host(self, params: dict) -> tuple["Status", dict]:
             A pair of Status and result. The result is always {}.
             Status is one of {PENDING, SUCCEEDED, FAILED}
         """
+        ...

mlos_bench/mlos_bench/services/types/local_exec_type.py

@@ -24,6 +24,9 @@ class SupportsLocalExec(Protocol):
     vs the target environment. Used in LocalEnv and provided by LocalExecService.
     """
 
+    # Needed by pyright
+    # pylint: disable=unnecessary-ellipsis,redundant-returns-doc
+
     def local_exec(
         self,
         script_lines: Iterable[str],
@@ -49,6 +52,7 @@ def local_exec(
         (return_code, stdout, stderr) : (int, str, str)
             A 3-tuple of return code, stdout, and stderr of the script process.
         """
+        ...
 
     def temp_dir_context(
         self,
@@ -59,11 +63,12 @@ def temp_dir_context(
 
         Parameters
         ----------
-        path : str
+        path : str | None
             A path to the temporary directory. Create a new one if None.
 
         Returns
         -------
         temp_dir_context : tempfile.TemporaryDirectory
             Temporary directory context to use in the `with` clause.
         """
+        ...