diff --git a/submission_runner.py b/submission_runner.py index 717ea2dc4..47730d3fc 100644 --- a/submission_runner.py +++ b/submission_runner.py @@ -340,9 +340,12 @@ def train_once( train_state['accumulated_submission_time'] += ( train_step_end_time - train_state['last_step_end_time']) + # Use 3x the runtime budget for the self-tuning ruleset. + max_allowed_runtime_sec = ( + workload.max_allowed_runtime_sec if FLAGS.tuning_ruleset == 'external' + else 3 * workload.max_allowed_runtime_sec) train_state['is_time_remaining'] = ( - train_state['accumulated_submission_time'] < - workload.max_allowed_runtime_sec) + train_state['accumulated_submission_time'] < max_allowed_runtime_sec) # Check if submission is eligible for an untimed eval. if ((train_step_end_time - train_state['last_eval_time']) >= workload.eval_period_time_sec or train_state['training_complete']): @@ -567,6 +570,9 @@ def score_submission_on_workload(workload: spec.Workload, logging.info(f'Total number of evals: {num_evals}') logging.info('=' * 20) else: + if tuning_search_space is not None: + raise ValueError( + 'Cannot provide a tuning search space when using self tuning.') if not rng_seed: rng_seed = struct.unpack('q', os.urandom(8))[0] rng = prng.PRNGKey(rng_seed)