Coverage for mlos_bench/mlos_bench/tests/conftest.py: 98%
51 statements
« prev ^ index » next coverage.py v7.6.10, created at 2025-01-21 01:50 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Common fixtures for mock TunableGroups and Environment objects."""
7import os
8import sys
9from collections.abc import Generator
10from typing import Any
12import pytest
13from fasteners import InterProcessLock, InterProcessReaderWriterLock
14from pytest_docker.plugin import Services as DockerServices
15from pytest_docker.plugin import get_docker_services
17from mlos_bench.environments.mock_env import MockEnv
18from mlos_bench.tests import SEED, tunable_groups_fixtures
19from mlos_bench.tunables.tunable_groups import TunableGroups
21# pylint: disable=redefined-outer-name
22# -- Ignore pylint complaints about pytest references to
23# `tunable_groups` fixture as both a function and a parameter.
25# Expose some of those as local names so they can be picked up as fixtures by pytest.
26tunable_groups_config = tunable_groups_fixtures.tunable_groups_config
27tunable_groups = tunable_groups_fixtures.tunable_groups
28mixed_numerics_tunable_groups = tunable_groups_fixtures.mixed_numerics_tunable_groups
29covariant_group = tunable_groups_fixtures.covariant_group
32@pytest.fixture
33def mock_env(tunable_groups: TunableGroups) -> MockEnv:
34 """Test fixture for MockEnv."""
35 return MockEnv(
36 name="Test Env",
37 config={
38 "tunable_params": ["provision", "boot", "kernel"],
39 "mock_env_seed": SEED,
40 "mock_env_range": [60, 120],
41 "mock_env_metrics": ["score"],
42 },
43 tunables=tunable_groups,
44 )
47@pytest.fixture
48def mock_env_no_noise(tunable_groups: TunableGroups) -> MockEnv:
49 """Test fixture for MockEnv."""
50 return MockEnv(
51 name="Test Env No Noise",
52 config={
53 "tunable_params": ["provision", "boot", "kernel"],
54 "mock_env_seed": -1,
55 "mock_env_range": [60, 120],
56 "mock_env_metrics": ["score", "other_score"],
57 },
58 tunables=tunable_groups,
59 )
62# Fixtures to configure the pytest-docker plugin.
63@pytest.fixture(scope="session")
64def docker_setup() -> list[str] | str:
65 """Setup for docker services."""
66 if sys.platform == "darwin" or os.environ.get("HOST_OSTYPE", "").lower().startswith("darwin"):
67 # Workaround an oddity on macOS where the "docker-compose up"
68 # command always recreates the containers.
69 # That leads to races when multiple workers are trying to
70 # start and use the same services.
71 return ["up --build -d --no-recreate"]
72 else:
73 return ["up --build -d"]
76@pytest.fixture(scope="session")
77def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
78 """
79 Returns the path to the docker-compose file.
81 Parameters
82 ----------
83 pytestconfig : pytest.Config
85 Returns
86 -------
87 str
88 Path to the docker-compose file.
89 """
90 _ = pytestconfig # unused
91 return [
92 os.path.join(os.path.dirname(__file__), "services", "remote", "ssh", "docker-compose.yml"),
93 # Add additional configs as necessary here.
94 ]
97@pytest.fixture(scope="session")
98def docker_compose_project_name(short_testrun_uid: str) -> str:
99 """
100 Returns the name of the docker-compose project.
102 Returns
103 -------
104 str
105 Name of the docker-compose project.
106 """
107 # Use the xdist testrun UID to ensure that the docker-compose project name
108 # is unique across sessions, but shared amongst workers.
109 return f"mlos_bench-test-{short_testrun_uid}"
112@pytest.fixture(scope="session")
113def docker_services_lock(
114 shared_temp_dir: str,
115 short_testrun_uid: str,
116) -> InterProcessReaderWriterLock:
117 """
118 Gets a pytest session lock for xdist workers to mark when they're using the docker
119 services.
121 Yields
122 ------
123 A lock to ensure that setup/teardown operations don't happen while a
124 worker is using the docker services.
125 """
126 return InterProcessReaderWriterLock(
127 f"{shared_temp_dir}/pytest_docker_services-{short_testrun_uid}.lock"
128 )
131@pytest.fixture(scope="session")
132def docker_setup_teardown_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessLock:
133 """
134 Gets a pytest session lock between xdist workers for the docker setup/teardown
135 operations.
137 Yields
138 ------
139 A lock to ensure that only one worker is doing setup/teardown at a time.
140 """
141 return InterProcessLock(
142 f"{shared_temp_dir}/pytest_docker_services-setup-teardown-{short_testrun_uid}.lock"
143 )
146@pytest.fixture(scope="session")
147def locked_docker_services(
148 docker_compose_command: Any,
149 docker_compose_file: Any,
150 docker_compose_project_name: Any,
151 docker_setup: Any,
152 docker_cleanup: Any,
153 docker_setup_teardown_lock: InterProcessLock,
154 docker_services_lock: InterProcessReaderWriterLock,
155) -> Generator[DockerServices, Any, None]:
156 """A locked version of the docker_services fixture to implement xdist single
157 instance locking.
158 """
159 # pylint: disable=too-many-arguments,too-many-positional-arguments
160 # Mark the services as in use with the reader lock.
161 docker_services_lock.acquire_read_lock()
162 # Acquire the setup lock to prevent multiple setup operations at once.
163 docker_setup_teardown_lock.acquire()
164 # This "with get_docker_services(...)"" pattern is in the default fixture.
165 # We call it instead of docker_services() to avoid pytest complaints about
166 # calling fixtures directly.
167 with get_docker_services(
168 docker_compose_command,
169 docker_compose_file,
170 docker_compose_project_name,
171 docker_setup,
172 docker_cleanup,
173 ) as docker_services:
174 # Release the setup/tear down lock in order to let the setup operation
175 # continue for other workers (should be a no-op at this point).
176 docker_setup_teardown_lock.release()
177 # Yield the services so that tests within this worker can use them.
178 yield docker_services
179 # Now tests that use those services get to run on this worker...
180 # Once the tests are done, release the read lock that marks the services as in use.
181 docker_services_lock.release_read_lock()
182 # Now as we prepare to execute the cleanup code on context exit we need
183 # to acquire the setup/teardown lock again.
184 # First we attempt to get the write lock so that we wait for other
185 # readers to finish and guard against a lock inversion possibility.
186 docker_services_lock.acquire_write_lock()
187 # Next, acquire the setup/teardown lock
188 # First one here is the one to do actual work, everyone else is basically a no-op.
189 # Upon context exit, we should execute the docker_cleanup code.
190 # And try to get the setup/tear down lock again.
191 docker_setup_teardown_lock.acquire()
192 # Finally, after the docker_cleanup code has finished, remove both locks.
193 docker_setup_teardown_lock.release()
194 docker_services_lock.release_write_lock()