84 lines
3 KiB
Python
84 lines
3 KiB
Python
from __future__ import annotations
|
|
from typing import Dict, Tuple, Callable
|
|
from optune.trial import Trial
|
|
import logging
|
|
|
|
import time
|
|
import subprocess
|
|
import pandas as pd
|
|
|
|
from hyperparameter_tuner.loss import Loss
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_objective(
|
|
slurm_script: str,
|
|
results_path: str,
|
|
loss: Loss,
|
|
trial_param_types: Dict[str, Tuple[str, Tuple, Dict]],
|
|
) -> Callable[[Trial], float]:
|
|
"""
|
|
Create an objective function for Optuna to optimize, which submits SLURM jobs and waits for results.
|
|
|
|
Args:
|
|
slurm_script: str - Path to the SLURM script to execute.
|
|
results_path: str - Path to the CSV file where results will be logged.
|
|
loss: Loss - An instance of a Loss class that implements a `calculate` method.
|
|
trial_param_types: Dict[str, Tuple[str, Tuple, Dict]]: Dictionary mapping parameter names to their types
|
|
and arguments. Each entry is structured as:
|
|
- key str: The name of the parameter.
|
|
- value: Tuple[str, Tuple, Dict]
|
|
- str - The parameter type ('int', 'float', or 'categorical').
|
|
- Tuple - Positional arguments for the parameter's sampling method.
|
|
- Dict - Keyword arguments for the parameter's sampling method.
|
|
"""
|
|
|
|
def objective(trial: Trial) -> float:
|
|
"""
|
|
Objective function to be passed to Optuna for trial evaluation.
|
|
|
|
Args:
|
|
trial (Trial): A trial object from Optuna, used to suggest parameter values.
|
|
|
|
Returns:
|
|
float: The calculated loss based on the results from the CSV file.
|
|
|
|
Raises:
|
|
ValueError: If an invalid parameter type is provided.
|
|
subprocess.CalledProcessError: If the SLURM job submission fails.
|
|
"""
|
|
param_methods = {'int': trial.suggest_int, 'float': trial.suggest_float, 'categorial': trial.suggest_categorical}
|
|
|
|
trial_id = trial.number
|
|
trial_params = {}
|
|
for param_name, (arg_type, args, kwargs) in trial_param_types.items():
|
|
if arg_type in param_methods:
|
|
trial_params[param_name] = param_methods[arg_type](param_name, *args, **kwargs)
|
|
else:
|
|
logger.error(f'Invalid parameter type: {arg_type}')
|
|
raise
|
|
|
|
command = f'sbatch {slurm_script} {results_path} {trial_id} {" ".join(str(v) for v in trial_params.values())}'
|
|
|
|
try:
|
|
subprocess.run(command, shell=True, check=True)
|
|
logger.info('SLURM job submitted with trial ID: {trial_id}')
|
|
logger.info('Paramters: {trial_params}')
|
|
except subprocess.CalledProcessError:
|
|
logger.error('Error submitting SLURM job')
|
|
raise
|
|
|
|
while True:
|
|
df = pd.read_csv(results_path)
|
|
|
|
matching_row = df.loc[df['trial'] == trial_id]
|
|
if matching_row.empty:
|
|
time.sleep(5)
|
|
continue
|
|
|
|
row_data = matching_row.iloc[0]
|
|
return loss.calculate(row_data)
|
|
|
|
return objective
|