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