Skip to content

Commit

Permalink
Clean up no_search mode to not use search_strategy (#643)
Browse files Browse the repository at this point in the history
## Describe your changes
We currently use a hacky method to run `no_search` using
`search_strategy`. This makes the logic confusing and unnecessarily
complicated.

Clean up the logic by just iterating over the pass flows directly. 

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [ ] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [ ] Format your code by running `pre-commit run --all-files`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link
  • Loading branch information
jambayk authored Oct 14, 2023
1 parent 989964d commit 4a00cad
Showing 1 changed file with 8 additions and 44 deletions.
52 changes: 8 additions & 44 deletions olive/engine/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -467,66 +467,30 @@ def run_no_search(
output_dir: str = None,
output_name: str = None,
):
"""Run all the registered Olive passes in no-search model where search strategy is None."""
assert (
self.search_strategy._config.execution_order == "joint"
), "run_no_search only supports default joint execution order"

"""Run all the registered Olive pass flows in no-search mode."""
for pass_item in self.passes.values():
if len(pass_item["pass"].search_space()) > 0:
pass_name = pass_item["name"]
raise ValueError(f"Pass {pass_name} has search space but search strategy is None")

evaluator_config = self.evaluator_for_pass(list(self.passes.keys())[-1])
if evaluator_config is None:
# provide dummy objective
objective_dict = {"dummy": {"higher_is_better": True, "goal": 0}}
else:
objective_dict = self.resolve_objectives(
input_model_config, input_model_id, data_root, evaluator_config.metrics, accelerator_spec
)

# initialize the search strategy
self.search_strategy.initialize(self.pass_flows_search_spaces, input_model_id, objective_dict)

iter_num = 0
output_models = {}
while True:
iter_num += 1

# get the next step
next_step = self.search_strategy.next_step()

if iter_num == 1:
assert next_step is not None, "Search strategy returned None for the first step"
# if no more steps, break
if next_step is None:
break

assert iter_num <= len(self.pass_flows), "No more pass flows to run"

# get the model id of the first input model
model_id = next_step["model_id"]
if model_id == input_model_id:
model_config = input_model_config
else:
model_config = self._load_model(model_id)

logger.debug(f"Step no search with search point {next_step['search_point']} ...")
for pass_flow in self.pass_flows:
# search point is empty since there is no search
passes_to_run = [(pass_id, {}) for pass_id in pass_flow]

# run all the passes in the step
# run all the passes in the pass flow
logger.debug(f"Running {pass_flow} with no search ...")
should_prune, signal, model_ids = self._run_passes(
next_step["passes"], model_config, model_id, data_root, accelerator_spec
passes_to_run, input_model_config, input_model_id, data_root, accelerator_spec
)

pass_flow = self.pass_flows[iter_num - 1]
if should_prune:
failed_pass = pass_flow[len(model_ids)]
logger.warning(f"Flow {pass_flow} is pruned due to failed or invalid config for pass '{failed_pass}'")
continue

# names of the output models of the passes
pass_output_names = [self.passes[pass_name]["output_name"] for pass_name, _ in next_step["passes"]]
pass_output_names = [self.passes[pass_id]["output_name"] for pass_id in pass_flow]
pass_output_names = [f"{name}_{accelerator_spec}" if name else None for name in pass_output_names]

# output dir with pass flow
Expand Down

0 comments on commit 4a00cad

Please sign in to comment.