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