Skip to content

Commit 9c06e6d

Browse files
verveerpjverveerpj
committed
Refactor ropt event handling
1 parent b35e66a commit 9c06e6d

File tree

5 files changed

+153
-109
lines changed

5 files changed

+153
-109
lines changed

‎src/ert/run_models/event.py

+6-9
Original file line numberDiff line numberDiff line change
@@ -32,23 +32,20 @@ class RunModelStatusEvent(RunModelEvent):
3232
class EverestStatusEvent(BaseModel):
3333
batch: int | None
3434
event_type: Literal["EverestStatusEvent"] = "EverestStatusEvent"
35-
36-
# Reflects what is currently in ROPT,
37-
# changes in ROPT should appear here accordingly
3835
everest_event: Literal[
39-
"START_EVALUATION",
36+
"START_OPTIMIZER_EVALUATION",
4037
"START_SAMPLING_EVALUATION",
41-
"START_OPTIMIZER_STEP",
42-
"FINISHED_OPTIMIZER_STEP",
43-
"START_EVALUATOR_STEP",
44-
"FINISHED_EVALUATOR_STEP",
4538
]
4639

4740

4841
class EverestBatchResultEvent(BaseModel):
4942
batch: int
5043
event_type: Literal["EverestBatchResultEvent"] = "EverestBatchResultEvent"
51-
everest_event: Literal["FINISHED_EVALUATION", "FINISHED_SAMPLING_EVALUATION"]
44+
everest_event: Literal[
45+
"OPTIMIZATION_RESULT",
46+
"FINISHED_OPTIMIZER_EVALUATION",
47+
"FINISHED_SAMPLING_EVALUATION",
48+
]
5249
result_type: Literal["FunctionResult", "GradientResult"]
5350

5451

‎src/ert/run_models/everest_run_model.py

+32-33
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,7 @@ def _init_domain_transforms(
346346

347347
self.send_event(
348348
EverestStatusEvent(
349-
batch=None, # Always 0, but omitting it for consistency
349+
batch=self._batch_id,
350350
everest_event="START_SAMPLING_EVALUATION",
351351
)
352352
)
@@ -365,6 +365,9 @@ def _init_domain_transforms(
365365
)
366366
)
367367

368+
# Increase the batch ID for the next evaluation:
369+
self._batch_id += 1
370+
368371
if transforms.objectives.has_auto_scale:
369372
transforms.objectives.calculate_auto_scales(objectives)
370373
if (
@@ -399,45 +402,29 @@ def _create_optimizer(self) -> BasicOptimizer:
399402
),
400403
)
401404

402-
def _forward_ropt_event(everest_event: OptimizerEvent) -> None:
405+
def _forward_results(everest_event: OptimizerEvent) -> None:
403406
has_results = bool(everest_event.data and everest_event.data.get("results"))
404407

405-
# The batch these results pertain to
406-
# If the event has results, they usually pertain to the
407-
# batch before self._batch_id, i.e., self._batch_id - 1
408-
409408
if has_results:
410409
# A ROPT event may contain multiple results, here we send one
411410
# event per result
412411
results = everest_event.data["results"]
413-
batch_id = results[0].batch_id
414412

415413
for r in results:
416414
self.send_event(
417415
EverestBatchResultEvent(
418-
batch=batch_id,
419-
everest_event=everest_event.event_type.name,
416+
batch=r.batch_id, # From the event, the batch that produced it.
417+
everest_event="OPTIMIZATION_RESULT",
420418
result_type=(
421419
"FunctionResult"
422420
if isinstance(r, FunctionResults)
423421
else "GradientResult"
424422
),
425423
)
426424
)
427-
else:
428-
# Events indicating the start of an evaluation,
429-
# start of optimizer step holds no results
430-
# but may still be relevant to the subscriber
431-
self.send_event(
432-
EverestStatusEvent(
433-
batch=None,
434-
everest_event=everest_event.event_type.name,
435-
)
436-
)
437425

438-
# Forward ROPT events to queue
439-
for event_type in EventType:
440-
optimizer.add_observer(event_type, _forward_ropt_event)
426+
# Send events when results are received.
427+
optimizer.add_observer(EventType.FINISHED_EVALUATION, _forward_results)
441428

442429
return optimizer
443430

@@ -513,9 +500,6 @@ def _run_forward_model(
513500
# Gather the results
514501
objectives, constraints = self._gather_simulation_results(ensemble)
515502

516-
# Increase the batch ID for the next evaluation:
517-
self._batch_id += 1
518-
519503
# Return the results, together with the indices of the evaluated controls:
520504
return objectives, constraints
521505

@@ -705,8 +689,6 @@ def _forward_model_evaluator(
705689
all_results=all_results,
706690
)
707691

708-
batch_id = self._batch_id # Save the batch ID, it will be modified.
709-
710692
control_values_to_simulate = np.array(
711693
[
712694
c.control_vector
@@ -720,19 +702,31 @@ def _forward_model_evaluator(
720702
sim_infos = [
721703
c for c in evaluation_infos if c.status == _EvaluationStatus.TO_SIMULATE
722704
]
705+
706+
self.send_event(
707+
EverestStatusEvent(
708+
batch=self._batch_id,
709+
everest_event="START_OPTIMIZER_EVALUATION",
710+
)
711+
)
712+
723713
sim_objectives, sim_constraints = self._run_forward_model(
724714
control_values=control_values_to_simulate,
725715
model_realizations=[c.model_realization for c in sim_infos],
726716
perturbations=[c.perturbation for c in sim_infos],
727717
)
718+
719+
self.send_event(
720+
EverestBatchResultEvent(
721+
batch=self._batch_id,
722+
everest_event="FINISHED_OPTIMIZER_EVALUATION",
723+
result_type="FunctionResult",
724+
)
725+
)
728726
else:
729727
sim_objectives = np.array([], dtype=np.float64)
730728
sim_constraints = np.array([], dtype=np.float64)
731729

732-
# Note: removing this causes some tests to fail, but this behavior
733-
# should maybe be revised?
734-
self._batch_id += 1
735-
736730
# Assign simulated results to corresponding evaluation infos
737731
for ei in evaluation_infos:
738732
if ei.simulation_id is not None:
@@ -774,13 +768,18 @@ def _forward_model_evaluator(
774768
dtype=np.int32,
775769
)
776770

777-
return EvaluatorResult(
771+
evaluator_result = EvaluatorResult(
778772
objectives=objectives,
779773
constraints=constraints,
780-
batch_id=batch_id,
774+
batch_id=self._batch_id,
781775
evaluation_ids=sim_ids,
782776
)
783777

778+
# increase the batch ID for the next evaluation:
779+
self._batch_id += 1
780+
781+
return evaluator_result
782+
784783
def _create_simulation_controls(
785784
self,
786785
control_values: NDArray[np.float64],

‎tests/everest/snapshots/test_everest_runmodel_events/test_everest_events/config_advanced.yml/snapshot.json

+57-31
Original file line numberDiff line numberDiff line change
@@ -1,116 +1,142 @@
11
{
22
"everest_events": [
33
{
4-
"batch": null,
4+
"batch": 0,
55
"event_type": "EverestStatusEvent",
6-
"everest_event": "START_OPTIMIZER_STEP"
6+
"everest_event": "START_OPTIMIZER_EVALUATION"
77
},
88
{
9-
"batch": null,
10-
"event_type": "EverestStatusEvent",
11-
"everest_event": "START_EVALUATION"
9+
"batch": 0,
10+
"event_type": "EverestBatchResultEvent",
11+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
12+
"result_type": "FunctionResult"
1213
},
1314
{
1415
"batch": 0,
1516
"event_type": "EverestBatchResultEvent",
16-
"everest_event": "FINISHED_EVALUATION",
17+
"everest_event": "OPTIMIZATION_RESULT",
1718
"result_type": "FunctionResult"
1819
},
1920
{
2021
"batch": 0,
2122
"event_type": "EverestBatchResultEvent",
22-
"everest_event": "FINISHED_EVALUATION",
23+
"everest_event": "OPTIMIZATION_RESULT",
2324
"result_type": "GradientResult"
2425
},
2526
{
26-
"batch": null,
27+
"batch": 1,
2728
"event_type": "EverestStatusEvent",
28-
"everest_event": "START_EVALUATION"
29+
"everest_event": "START_OPTIMIZER_EVALUATION"
30+
},
31+
{
32+
"batch": 1,
33+
"event_type": "EverestBatchResultEvent",
34+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
35+
"result_type": "FunctionResult"
2936
},
3037
{
3138
"batch": 1,
3239
"event_type": "EverestBatchResultEvent",
33-
"everest_event": "FINISHED_EVALUATION",
40+
"everest_event": "OPTIMIZATION_RESULT",
3441
"result_type": "FunctionResult"
3542
},
3643
{
3744
"batch": 1,
3845
"event_type": "EverestBatchResultEvent",
39-
"everest_event": "FINISHED_EVALUATION",
46+
"everest_event": "OPTIMIZATION_RESULT",
4047
"result_type": "GradientResult"
4148
},
4249
{
43-
"batch": null,
50+
"batch": 2,
4451
"event_type": "EverestStatusEvent",
45-
"everest_event": "START_EVALUATION"
52+
"everest_event": "START_OPTIMIZER_EVALUATION"
4653
},
4754
{
4855
"batch": 2,
4956
"event_type": "EverestBatchResultEvent",
50-
"everest_event": "FINISHED_EVALUATION",
57+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
5158
"result_type": "FunctionResult"
5259
},
5360
{
5461
"batch": 2,
5562
"event_type": "EverestBatchResultEvent",
56-
"everest_event": "FINISHED_EVALUATION",
63+
"everest_event": "OPTIMIZATION_RESULT",
64+
"result_type": "FunctionResult"
65+
},
66+
{
67+
"batch": 2,
68+
"event_type": "EverestBatchResultEvent",
69+
"everest_event": "OPTIMIZATION_RESULT",
5770
"result_type": "GradientResult"
5871
},
5972
{
60-
"batch": null,
73+
"batch": 3,
6174
"event_type": "EverestStatusEvent",
62-
"everest_event": "START_EVALUATION"
75+
"everest_event": "START_OPTIMIZER_EVALUATION"
6376
},
6477
{
6578
"batch": 3,
6679
"event_type": "EverestBatchResultEvent",
67-
"everest_event": "FINISHED_EVALUATION",
80+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
6881
"result_type": "FunctionResult"
6982
},
7083
{
7184
"batch": 3,
7285
"event_type": "EverestBatchResultEvent",
73-
"everest_event": "FINISHED_EVALUATION",
86+
"everest_event": "OPTIMIZATION_RESULT",
87+
"result_type": "FunctionResult"
88+
},
89+
{
90+
"batch": 3,
91+
"event_type": "EverestBatchResultEvent",
92+
"everest_event": "OPTIMIZATION_RESULT",
7493
"result_type": "GradientResult"
7594
},
7695
{
77-
"batch": null,
96+
"batch": 4,
7897
"event_type": "EverestStatusEvent",
79-
"everest_event": "START_EVALUATION"
98+
"everest_event": "START_OPTIMIZER_EVALUATION"
99+
},
100+
{
101+
"batch": 4,
102+
"event_type": "EverestBatchResultEvent",
103+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
104+
"result_type": "FunctionResult"
80105
},
81106
{
82107
"batch": 4,
83108
"event_type": "EverestBatchResultEvent",
84-
"everest_event": "FINISHED_EVALUATION",
109+
"everest_event": "OPTIMIZATION_RESULT",
85110
"result_type": "FunctionResult"
86111
},
87112
{
88113
"batch": 4,
89114
"event_type": "EverestBatchResultEvent",
90-
"everest_event": "FINISHED_EVALUATION",
115+
"everest_event": "OPTIMIZATION_RESULT",
91116
"result_type": "GradientResult"
92117
},
93118
{
94-
"batch": null,
119+
"batch": 5,
95120
"event_type": "EverestStatusEvent",
96-
"everest_event": "START_EVALUATION"
121+
"everest_event": "START_OPTIMIZER_EVALUATION"
97122
},
98123
{
99124
"batch": 5,
100125
"event_type": "EverestBatchResultEvent",
101-
"everest_event": "FINISHED_EVALUATION",
126+
"everest_event": "FINISHED_OPTIMIZER_EVALUATION",
102127
"result_type": "FunctionResult"
103128
},
104129
{
105130
"batch": 5,
106131
"event_type": "EverestBatchResultEvent",
107-
"everest_event": "FINISHED_EVALUATION",
108-
"result_type": "GradientResult"
132+
"everest_event": "OPTIMIZATION_RESULT",
133+
"result_type": "FunctionResult"
109134
},
110135
{
111-
"batch": null,
112-
"event_type": "EverestStatusEvent",
113-
"everest_event": "FINISHED_OPTIMIZER_STEP"
136+
"batch": 5,
137+
"event_type": "EverestBatchResultEvent",
138+
"everest_event": "OPTIMIZATION_RESULT",
139+
"result_type": "GradientResult"
114140
}
115141
],
116142
"num_full_snapshots": 6

0 commit comments

Comments
 (0)