Modify Trial when using a Sweeper

[michelkok]

🚀 Feature Request

A way to modify a Trial object while using a sweeper.

Motivation

Using a sweeper the Trial object is not accessible from Python code anywhere.

A potential use case is the setting of User Attributes before or after the Sweeper launched the Trial. For example, when doing hyperparameter optimization I use the Optuna sweeper plugin. The Study gives me back the best Trial with the optimal configuration. To find my model (that is stored in a separate database), I do not want to retrain my model with the optimal setup. I'd rather just extract the reference ID of this model. User Attributes seems like a logical place to store such information. However, using the sweeper there is no way to set these as the Trial object is not accessible.

I have tried asking on StackOverflow but no answer was given. Then I looked through the codebase and it just seems impossible at the moment (unless I am missing something).

Pitch

Describe the solution you'd like
Ideally, I'd like the Trial object to be available within the main function. Something along these lines:

@hydra.main(version_base=None, config_path="conf", config_name="sweep")
def sphere(cfg: TrainConfig, trial: Trial) -> float:
    model, loss = train_model(cfg)
    model_db_id = save_model(model)
    trial.set_user_attr("model_db_id": model_id)

    return loss, trial

Describe alternatives you've considered
An alternative could be to provide some callback function that can be used to modify the Trial object like custom_search_space. I like this less though, as there is no control over when to modify the Trial (e.g., the function below is called before actually computing anything - in my example, I'd like to use it after training my model).

I have even tried to implement this for my use case but this is really just a hack and not a solution as the function is clearly not intended to do the stuff I'm doing here.

hydra:
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    study_name: train_model
    storage: sqlite:///trials.db
    n_trials: 20
    n_jobs: 1
    params:
      x: range(-5.5, 5.5, step=0.5)
      y: choice(-5 ,0 ,5)
      z: choice([0, 1], [2, 3], [2, 5])
    custom_search_space: package.run.configure

def configure(_, trial: Trial) -> None:
    trial.set_user_attr("model_db_id", 123456)

Are you willing to open a pull request? (See CONTRIBUTING)
Yes, but I am not familiar with most of the codebase. But for relatively simple changes, sure!

[Jasha10]

For reference, there is some related discussion in #1710.

[Jasha10]

there is no control over when to modify the Trial (e.g., the function below is called before actually computing anything - in my example, I'd like to use it after training my model

@mkardas, what is the use-case for modifying the trial object after training the model?

[mkardas]

@mkardas, what is the use-case for modifying the trial object after training the model?

@Jasha10 did you mean @michelkok?

[Jasha10]

Yes, I did mean @michelkok! Sorry about the ping Marcin!

[michelkok]

For example, to attach User Attributes.

[omry]

Since Hydra is abstracting the underlying sweeper, it does not seem likely that this can be supported.
There are a few issues:

1. The abstraction is getting in the way. We can't easily expose objects from the underlying sweeper.
2. Launchers means that the actual job process may be running somewhere else. There isn't a full fledged communication channel between the launching/sweeping process and the child processes which means you may not have access to the information you intend to attach in the sweeping process.

[michelkok]

Sorry, I'm only coming back to this now.
I understand that my proposed solution does not really fit with the layered design of Hydra.

Is there any other way to achieve what I want?

Again, the main thing is that I'm struggling to find back the model + parameters of the best run after a sweep.
The current workflow is like this:

1. Run a sweep (with the Optuna Sweeper plugin).
2. Load the sth.db study and find the best trial.
3. Find parameters, model weights and metrics of the trial.

Going from step 2 to 3 seems really hard since the trial object itself contains only a bit of information and it cannot be modified.

Another approach I have tried for example is using the experimental Callback. But in that I have to maintain a state with similar properties to the sth.db which seems redundant (e.g., the best run and it's configured output directory).