wehub-resource-sync
77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from ray import tune
|
|
from ray.tune.search import BasicVariantGenerator
|
|
|
|
# Hang full script until this marker is deleted
|
|
HANG_RUN_MARKER = os.environ.get("HANG_RUN_MARKER", "")
|
|
|
|
# Delete this marker when a trial is started
|
|
DELETE_TRIAL_MARKER = os.environ.get("DELETE_TRIAL_MARKER", "")
|
|
|
|
# Hang in trial until this marker is deleted
|
|
HANG_TRIAL_MARKER = os.environ.get("HANG_TRIAL_MARKER", "")
|
|
|
|
# Delete this marker after tuning finished
|
|
DELETE_RUN_MARKER = os.environ.get("DELETE_RUN_MARKER", "")
|
|
|
|
# Hang at end of run until this marker is deleted
|
|
HANG_END_MARKER = os.environ.get("HANG_END_MARKER", "")
|
|
|
|
# Report this val as the "fixed" metric in the trial.
|
|
# This value is captured in the trainer and will conflict when a trainable
|
|
# is overwritten!
|
|
FIXED_VAL = int(os.environ["FIXED_VAL"])
|
|
|
|
# Grid search over these vals and report as "param" metric in the trial.
|
|
# Even with conflicting trainables, these will be reported correctly as they
|
|
# are tracked by the driver, not the trainable.
|
|
VALS = [int(os.environ["VAL_1"]), int(os.environ["VAL_2"])]
|
|
|
|
# Wait for HANG_RUN_MARKER
|
|
while HANG_RUN_MARKER and Path(HANG_RUN_MARKER).exists():
|
|
time.sleep(0.1)
|
|
|
|
|
|
def train_func(config):
|
|
# Delete DELETE_TRIAL_MARKER
|
|
delete_marker = config["delete_marker"]
|
|
if delete_marker and Path(delete_marker).exists():
|
|
Path(delete_marker).unlink()
|
|
|
|
# Wait for HANG_TRIAL_MARKER
|
|
hang_marker = config["hang_marker"]
|
|
while hang_marker and Path(hang_marker).exists():
|
|
time.sleep(0.1)
|
|
|
|
# Finish trial
|
|
tune.report({"param": config["param"], "fixed": config["fixed"]})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tuner = tune.Tuner(
|
|
tune.with_resources(train_func, {"CPU": 2}),
|
|
param_space={
|
|
"fixed": FIXED_VAL,
|
|
"param": tune.grid_search(VALS),
|
|
"delete_marker": DELETE_TRIAL_MARKER,
|
|
"hang_marker": HANG_TRIAL_MARKER,
|
|
},
|
|
tune_config=tune.TuneConfig(search_alg=BasicVariantGenerator(max_concurrent=1)),
|
|
)
|
|
results = tuner.fit()
|
|
|
|
# Delete DELETE_RUN_MARKER
|
|
if DELETE_RUN_MARKER and Path(DELETE_RUN_MARKER).exists():
|
|
Path(DELETE_RUN_MARKER).unlink()
|
|
|
|
# Wait for HANG_END_MARKER
|
|
while HANG_END_MARKER and Path(HANG_END_MARKER).exists():
|
|
time.sleep(0.1)
|
|
|
|
# Put assertions last, so we don't finish early because of failures
|
|
assert sorted([result.metrics["param"] for result in results]) == VALS
|
|
assert [result.metrics["fixed"] for result in results] == [FIXED_VAL, FIXED_VAL]
|