Coverage for mlos_bench/mlos_bench/tests/environments/composite_env_test.py: 100%

34 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-07 01:52 +0000

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Unit tests for composite environment.""" 

6 

7import pytest 

8 

9from mlos_bench.environments.composite_env import CompositeEnv 

10from mlos_bench.services.config_persistence import ConfigPersistenceService 

11from mlos_bench.tunables.tunable_groups import TunableGroups 

12 

13# pylint: disable=redefined-outer-name 

14 

15 

16@pytest.fixture 

17def composite_env(tunable_groups: TunableGroups) -> CompositeEnv: 

18 """Test fixture for CompositeEnv.""" 

19 return CompositeEnv( 

20 name="Composite Test Environment", 

21 config={ 

22 "tunable_params": ["provision", "boot"], 

23 "const_args": { 

24 "vm_server_name": "Mock Server VM", 

25 "vm_client_name": "Mock Client VM", 

26 "someConst": "root", 

27 "global_param": "default", 

28 }, 

29 "children": [ 

30 { 

31 "name": "Mock Client Environment 1", 

32 "class": "mlos_bench.environments.mock_env.MockEnv", 

33 "config": { 

34 "tunable_params": ["provision"], 

35 "const_args": { 

36 "vmName": "$vm_client_name", 

37 "EnvId": 1, 

38 }, 

39 "required_args": ["vmName", "someConst", "global_param"], 

40 "mock_env_range": [60, 120], 

41 "mock_env_metrics": ["score"], 

42 }, 

43 }, 

44 { 

45 "name": "Mock Server Environment 2", 

46 "class": "mlos_bench.environments.mock_env.MockEnv", 

47 "config": { 

48 "tunable_params": ["boot"], 

49 "const_args": { 

50 "vmName": "$vm_server_name", 

51 "EnvId": 2, 

52 "global_param": "local", 

53 }, 

54 "required_args": ["vmName"], 

55 "mock_env_range": [60, 120], 

56 "mock_env_metrics": ["score"], 

57 }, 

58 }, 

59 { 

60 "name": "Mock Control Environment 3", 

61 "class": "mlos_bench.environments.mock_env.MockEnv", 

62 "config": { 

63 "tunable_params": ["boot"], 

64 "const_args": { 

65 "vmName": "Mock Control VM", 

66 "EnvId": 3, 

67 }, 

68 "required_args": ["vmName", "vm_server_name", "vm_client_name"], 

69 "mock_env_range": [60, 120], 

70 "mock_env_metrics": ["score"], 

71 }, 

72 }, 

73 ], 

74 }, 

75 tunables=tunable_groups, 

76 service=ConfigPersistenceService({}), 

77 global_config={"global_param": "global_value"}, 

78 ) 

79 

80 

81def test_composite_env_params(composite_env: CompositeEnv) -> None: 

82 """ 

83 Check that the const_args from the parent environment get propagated to the 

84 children. 

85 

86 NOTE: The current logic is that variables flow down via required_args and const_args, parent 

87 """ 

88 assert composite_env.children[0].parameters == { 

89 "vmName": "Mock Client VM", # const_args from the parent thru variable substitution 

90 "EnvId": 1, # const_args from the child 

91 "vmSize": "Standard_B4ms", # tunable_params from the parent 

92 "someConst": "root", # pulled in from parent via required_args 

93 "global_param": "global_value", # pulled in from the global_config 

94 } 

95 assert composite_env.children[1].parameters == { 

96 "vmName": "Mock Server VM", # const_args from the parent 

97 "EnvId": 2, # const_args from the child 

98 "idle": "halt", # tunable_params from the parent 

99 # "someConst": "root" # not required, so not passed from the parent 

100 "global_param": "global_value", # pulled in from the global_config 

101 } 

102 assert composite_env.children[2].parameters == { 

103 "vmName": "Mock Control VM", # const_args from the parent 

104 "EnvId": 3, # const_args from the child 

105 "idle": "halt", # tunable_params from the parent 

106 # "someConst": "root" # not required, so not passed from the parent 

107 "vm_client_name": "Mock Client VM", 

108 "vm_server_name": "Mock Server VM", 

109 # "global_param": "global_value" # not required, so not picked from the global_config 

110 } 

111 

112 

113def test_composite_env_setup(composite_env: CompositeEnv, tunable_groups: TunableGroups) -> None: 

114 """Check that the child environments update their tunable parameters.""" 

115 tunable_groups.assign( 

116 { 

117 "vmSize": "Standard_B2s", 

118 "idle": "mwait", 

119 "kernel_sched_migration_cost_ns": 100000, 

120 } 

121 ) 

122 

123 with composite_env as env_context: 

124 assert env_context.setup(tunable_groups) 

125 

126 assert composite_env.children[0].parameters == { 

127 "vmName": "Mock Client VM", # const_args from the parent 

128 "EnvId": 1, # const_args from the child 

129 "vmSize": "Standard_B2s", # tunable_params from the parent 

130 "someConst": "root", # pulled in from parent via required_args 

131 "global_param": "global_value", # pulled in from the global_config 

132 } 

133 assert composite_env.children[1].parameters == { 

134 "vmName": "Mock Server VM", # const_args from the parent 

135 "EnvId": 2, # const_args from the child 

136 "idle": "mwait", # tunable_params from the parent 

137 # "someConst": "root" # not required, so not passed from the parent 

138 "global_param": "global_value", # pulled in from the global_config 

139 } 

140 assert composite_env.children[2].parameters == { 

141 "vmName": "Mock Control VM", # const_args from the parent 

142 "EnvId": 3, # const_args from the child 

143 "idle": "mwait", # tunable_params from the parent 

144 "vm_client_name": "Mock Client VM", 

145 "vm_server_name": "Mock Server VM", 

146 # "global_param": "global_value" # not required, so not picked from the global_config 

147 } 

148 

149 

150@pytest.fixture 

151def nested_composite_env(tunable_groups: TunableGroups) -> CompositeEnv: 

152 """Test fixture for CompositeEnv.""" 

153 return CompositeEnv( 

154 name="Composite Test Environment", 

155 config={ 

156 "tunable_params": ["provision", "boot"], 

157 "const_args": { 

158 "vm_server_name": "Mock Server VM", 

159 "vm_client_name": "Mock Client VM", 

160 "someConst": "root", 

161 }, 

162 "children": [ 

163 { 

164 "name": "Nested Composite Client Environment 1", 

165 "class": "mlos_bench.environments.composite_env.CompositeEnv", 

166 "config": { 

167 "tunable_params": ["provision"], 

168 "const_args": { 

169 "vmName": "$vm_client_name", 

170 "EnvId": 1, 

171 }, 

172 "required_args": ["vmName", "EnvId", "someConst", "vm_server_name"], 

173 "children": [ 

174 { 

175 "name": "Mock Client Environment 1", 

176 "class": "mlos_bench.environments.mock_env.MockEnv", 

177 "config": { 

178 "tunable_params": ["provision"], 

179 # TODO: Might be nice to include a "^" or "*" option 

180 # here to indicate that all required_args from 

181 # the parent should be included here too in 

182 # order to reduce duplication. 

183 "required_args": [ 

184 "vmName", 

185 "EnvId", 

186 "someConst", 

187 "vm_server_name", 

188 "global_param", 

189 ], 

190 "mock_env_range": [60, 120], 

191 "mock_env_metrics": ["score"], 

192 }, 

193 }, 

194 # ... 

195 ], 

196 }, 

197 }, 

198 { 

199 "name": "Nested Composite Server Environment 2", 

200 "class": "mlos_bench.environments.composite_env.CompositeEnv", 

201 "config": { 

202 "tunable_params": ["boot"], 

203 "const_args": { 

204 "vmName": "$vm_server_name", 

205 "EnvId": 2, 

206 }, 

207 "required_args": ["vmName", "EnvId", "vm_client_name"], 

208 "children": [ 

209 { 

210 "name": "Mock Server Environment 2", 

211 "class": "mlos_bench.environments.mock_env.MockEnv", 

212 "config": { 

213 "tunable_params": ["boot"], 

214 "required_args": ["vmName", "EnvId", "vm_client_name"], 

215 "mock_env_range": [60, 120], 

216 "mock_env_metrics": ["score"], 

217 }, 

218 }, 

219 # ... 

220 ], 

221 }, 

222 }, 

223 ], 

224 }, 

225 tunables=tunable_groups, 

226 service=ConfigPersistenceService({}), 

227 global_config={"global_param": "global_value"}, 

228 ) 

229 

230 

231def test_nested_composite_env_params(nested_composite_env: CompositeEnv) -> None: 

232 """ 

233 Check that the const_args from the parent environment get propagated to the 

234 children. 

235 

236 NOTE: The current logic is that variables flow down via required_args and const_args, parent 

237 """ 

238 assert isinstance(nested_composite_env.children[0], CompositeEnv) 

239 assert nested_composite_env.children[0].children[0].parameters == { 

240 "vmName": "Mock Client VM", # const_args from the parent thru variable substitution 

241 "EnvId": 1, # const_args from the child 

242 "vmSize": "Standard_B4ms", # tunable_params from the parent 

243 "someConst": "root", # pulled in from parent via required_args 

244 "vm_server_name": "Mock Server VM", 

245 "global_param": "global_value", # pulled in from the global_config 

246 } 

247 assert isinstance(nested_composite_env.children[1], CompositeEnv) 

248 assert nested_composite_env.children[1].children[0].parameters == { 

249 "vmName": "Mock Server VM", # const_args from the parent 

250 "EnvId": 2, # const_args from the child 

251 "idle": "halt", # tunable_params from the parent 

252 # "someConst": "root" # not required, so not passed from the parent 

253 "vm_client_name": "Mock Client VM", 

254 # "global_param": "global_value" # not required, so not picked from the global_config 

255 } 

256 

257 

258def test_nested_composite_env_setup( 

259 nested_composite_env: CompositeEnv, 

260 tunable_groups: TunableGroups, 

261) -> None: 

262 """Check that the child environments update their tunable parameters.""" 

263 tunable_groups.assign( 

264 { 

265 "vmSize": "Standard_B2s", 

266 "idle": "mwait", 

267 "kernel_sched_migration_cost_ns": 100000, 

268 } 

269 ) 

270 

271 with nested_composite_env as env_context: 

272 assert env_context.setup(tunable_groups) 

273 

274 assert isinstance(nested_composite_env.children[0], CompositeEnv) 

275 assert nested_composite_env.children[0].children[0].parameters == { 

276 "vmName": "Mock Client VM", # const_args from the parent 

277 "EnvId": 1, # const_args from the child 

278 "vmSize": "Standard_B2s", # tunable_params from the parent 

279 "someConst": "root", # pulled in from parent via required_args 

280 "vm_server_name": "Mock Server VM", 

281 "global_param": "global_value", # pulled in from the global_config 

282 } 

283 

284 assert isinstance(nested_composite_env.children[1], CompositeEnv) 

285 assert nested_composite_env.children[1].children[0].parameters == { 

286 "vmName": "Mock Server VM", # const_args from the parent 

287 "EnvId": 2, # const_args from the child 

288 "idle": "mwait", # tunable_params from the parent 

289 # "someConst": "root" # not required, so not passed from the parent 

290 "vm_client_name": "Mock Client VM", 

291 }