[GENERAL SUPPORT]: Reusing/importing trials.

[cesare-montresor]

Question

I'm in the condition of repeating the same NAS search using different settings ( optimizing against a different set of metrics, or fixing some experiment paramerter)

So, i'm trying to run a common the sobol phase ( with all metrics and unlocked params ) and then, i setup 4 different clients, for the experiements variants (namely A,B,C,D, which can run in parallel ).

In order to achieve this, currently I'm using multiple clients:
1x COMMON: strategy sobol, 30 trials
4x VARIANT (A,B,C,D): strategy sobol 30 trials + bo 30 trials
variants runs after COMMON and they are initialized using
client.experiment.warm_start_from_old_experiment() and they are succefully imported.
At this point, i would expect to pull from trial 31, however no, the client still starts from 1, and run 60 more trials, 30 sobol, 30 bo, ending up with 90 trials in total.
It doesn't look like there is a clean way to advance the generation strategy to reflect the current number of trials.

Alternatively i tried to create the new strategy using only the 'bo' steps, or if i l put 0 trials for the sobol step, crash at:

client.get_next_trial()

with:

Exception has occurred: DataRequiredError
`MapKeyToFloat` transform requires non-empty data.
  File "/home/cesare/Projects/StrangeNet/modularnet/src/modularnet/metamodel/experiment.py", line 500, in run_phase
    trials, stop = phase.client.get_next_trials(max_trials=parallel_trials)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ax.exceptions.core.DataRequiredError: `MapKeyToFloat` transform requires non-empty data.

What's the correct way of doing so ?

configure client

def configure_client(self, template:MetaModelTemplate, strategy:str='bo'):
        # seed_everything(self.seed) -> seed what ? ax client ?
        params = template.get_params()

        min_acc = template.get_value('nas_min_acc', 0.8) 
        objectives = {'val_acc': ObjectiveProperties(minimize=False, threshold=min_acc)}


        if template.get_value('exp_metric_modularity', False):
            min_mod_score = template.get_value('nas_min_mod_score', 0.2)
            objectives['modularity_score'] = ObjectiveProperties(minimize=False, threshold=min_mod_score) 
        
        sobol_step = GenerationStep(Generators.SOBOL, self.trials_prior )
        bo_step = GenerationStep(Generators.BOTORCH_MODULAR, self.trials_search)

        generation_steps = [ sobol_step ] 
        
        if strategy == 'bo':
            generation_steps.append(bo_step)
        

        ax_client = AxClient(
            generation_strategy=GenerationStrategy(steps=generation_steps),
        )

        ax_client.create_experiment( 
            name=self.name, 
            parameters=params,
            objectives=objectives,
            tracking_metric_names=['modularity_score','val_acc'],
        )


        
        return ax_client

        
        num_trials = client.experiment.num_trials
        if old_client is not None:
            imported_trials = client.experiment.warm_start_from_old_experiment(old_client.experiment, search_space_check_membership_raise_error=False)
            print(f"Imported {len(imported_trials)}/{len(old_client.experiment.trials)} trials from previous experiment")
        start_trials = client.experiment.num_trials
        while not stop:
            params, trial_index = client.get_next_trial()
            print(start_trials,params, trial_index)

Please provide any relevant code snippet if applicable.

Code of Conduct

• I agree to follow this Ax's Code of Conduct

[mgarrard]

Hi @cesare-montresor, thanks so much for reaching out! What you are referring to here is an active area of work for our team, and we hope to be able to cut a minor release with full API functionality for this in 2026 :)

In the meantime, I think you are on the right track, using warm_start_from_old_experiment - with the initial trials are you marking them as running/completed and attaching data anywhere? If so, could you include that code snippet?

What I believe is happening is that the GenerationStrategy your experiment is using wants trials that have, or are expecting data, in order to transition from SOBOL to MBM, and then the BO phase needs data for modeling training etc. So I think we're close, just need to do some trial state + data attaching management :)

[cesare-montresor]

Hi @mgarrard !
Thank you for looking into it, I'm currently in need of this feature, therefore i'm running a bit more testing.
While my testing is not comprehensive, and my experience limited, I've been putting together a more reproducible test case for my issue, you can find it here:

import json
import os
import random
from ax.adapter import Generators
from ax.analysis.analysis import GenerationStrategy
from ax.generation_strategy.generation_node import GenerationStep
from ax.service.ax_client import AxClient, ObjectiveProperties



def main():
    #crash_1: choice param with single value are turned into fixed params, causing crash when importing ( by design ? )
    crash_1=True 

    #crash_2: disable crash_1 to replicate crash_2
    #crash_2: choice param with 2 values are then compared to choice values with 1 params ( the become fixed )
    #crash_2: this causes crash BO at the time of fitting ( with a not very readable error ) 
    crash_2=True  
    
    pc = params_common(crash_2)  #crash_2: see below
    pa = params_a()
    pb = params_b()


    o_a = objective_a()
    o_b = objective_b()

    m = metrics()

    so = sobol_client()
    bo_a = bo_client()
    bo_b = bo_client()

    print("Create experiments")
    so.create_experiment(parameters=pc, objectives=o_a, tracking_metric_names=m)
    bo_a.create_experiment(parameters=pa, objectives=o_a, tracking_metric_names=m)
    bo_b.create_experiment(parameters=pb, objectives=o_b, tracking_metric_names=m)
    
    print("------------------------")
    print("SOBOL Optimization")
    run_trials(so)


    print("------------------------")
    print("Import Trials to BO A")
    trials = bo_a.experiment.warm_start_from_old_experiment( 
        so.experiment, 
        search_space_check_membership_raise_error=crash_1 #crash_1: trigger
    )
    print("Imported Trials: ", len(trials) )


    print("------------------------")
    print("Import Trials to BO B")
    trials = bo_a.experiment.warm_start_from_old_experiment( 
        so.experiment, 
        search_space_check_membership_raise_error=crash_1 #crash_1: trigger
    )
    print("Imported Trials: ", len(trials) )
    

    
    print("------------------------")
    print("BO Optimization A")
    run_trials(bo_a)   # successfully runs
    print("------------------------")
    print("BO Optimization B")
    # crash_2: trupy happens of "fit"
    # bo_b.fit_model() # uncommenting this line makes the error appear earlier
    run_trials(bo_b)   #crash_2: fails with mystical error ( truly, it fails on BO fit )
    

# sobol
def params_common(crash:bool=False):
    if crash:
        return params_common_crash()
    else:
        return params_common_no_crash()


def params_common_no_crash():
    return [
        {"name": "learning_rate", "type": "range", "bounds": [1e-5, 1e-2], "log_scale": True},
        {"name": "batch_size", "type": "range", "bounds": [16, 128], "value_type": "int"},
        {"name": "exp_metric_modularity", "type": "choice", "values": [False, True], "value_type": "bool"},  #crash_1: choice 2 values -> choice
        {"name": "exp_modularizer", "type": "choice", "values": [False, True], "value_type": "bool"},
    ]

def params_common_crash():
    return [
        {"name": "learning_rate", "type": "range", "bounds": [1e-5, 1e-2], "log_scale": True},
        {"name": "batch_size", "type": "range", "bounds": [16, 128], "value_type": "int"},
        {"name": "exp_metric_modularity", "type": "choice", "values": [True], "value_type": "bool"},  #crash_2: single value True
        {"name": "exp_modularizer", "type": "choice", "values": [False, True], "value_type": "bool"},
    ]

# bo

def params_a():
    return [
        {"name": "learning_rate", "type": "range", "bounds": [1e-5, 1e-2], "log_scale": True},
        {"name": "batch_size", "type": "range", "bounds": [16, 128], "value_type": "int"},
        {"name": "exp_metric_modularity", "type": "choice", "values": [True], "value_type": "bool"}, 
        {"name": "exp_modularizer", "type": "choice", "values": [False], "value_type": "bool"}, #crash_1: choice 1 values -> fixed
    ]

def params_b():
    return [
        {"name": "learning_rate", "type": "range", "bounds": [1e-5, 1e-2], "log_scale": True},
        {"name": "batch_size", "type": "range", "bounds": [16, 128], "value_type": "int"},
        {"name": "exp_metric_modularity", "type": "choice", "values": [False], "value_type": "bool"}, #crash_2: single value False
        {"name": "exp_modularizer", "type": "choice", "values": [True], "value_type": "bool"}, #crash_1: choice 1 values -> fixed
    ]




def objective_a():
    return {
        'val_acc': ObjectiveProperties(minimize=False, threshold=0.80),
        'modularity_score': ObjectiveProperties(minimize=False, threshold=0.20),
    }

def objective_b():
    return {
        'val_acc': ObjectiveProperties(minimize=False, threshold=0.80),
    }

def metrics():
    return ['val_acc', 'modularity_score']


def load_client(path:str):
    if not os.path.exists(path): return None
    with open(path, 'r') as f:
        data = json.load(f)
        return AxClient.from_json_snapshot(data)
    


def run_trials(ax_client:AxClient, acc_only:bool=False):
    stop = False
    while not stop:
        trail_data, stop = ax_client.get_next_trials(max_trials=1)
        for trail_index, params in trail_data.items():
            print(trail_index)  #, params)
            result = run_trial(trail_index, params, acc_only)
            ax_client.complete_trial(trial_index=trail_index, raw_data=result)


def sobol_client():
    ax_client = AxClient(
        generation_strategy=GenerationStrategy(steps=[
            GenerationStep(Generators.SOBOL, num_trials=5),
        ]),
    )
    return ax_client

def bo_client():
    ax_client = AxClient(
        generation_strategy=GenerationStrategy(steps=[
            GenerationStep(Generators.SOBOL, num_trials=5, use_all_trials_in_exp=True),
            GenerationStep(Generators.BOTORCH_MODULAR, num_trials=5),
        ]),
    )
    return ax_client




def run_trial(trial_index, params, acc_only: bool=False):
    if acc_only:
        return {'val_acc': random.random()}
    
    return {'modularity_score': random.random(),'val_acc': random.random()}



if __name__ == "__main__": main()

this is the best error I got when attempting to manually try to fit the BO model

client.fit_model()
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/service/ax_client.py", line 1213, in fit_model
    self.generation_strategy._curr._fit(experiment=self.experiment)
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/generation_strategy/generation_node.py", line 359, in _fit
    generator_spec.fit(
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/generation_strategy/generator_spec.py", line 140, in fit
    self._fitted_adapter = self.generator_enum(
                           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/registry.py", line 348, in __call__
    adapter = adapter_class(
              ^^^^^^^^^^^^^^
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/torch.py", line 148, in __init__
    super().__init__(
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/base.py", line 226, in __init__
    experiment_data, search_space = self._transform_data(
                                    ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/base.py", line 318, in _transform_data
    t_instance = t(
                 ^^
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/transforms/map_key_to_float.py", line 64, in __init__
    Transform.__init__(
  File "/home/cesare/Projects/StrangeNet/venv/lib64/python3.11/site-packages/ax/adapter/transforms/base.py", line 81, in __init__
    raise DataRequiredError(
ax.exceptions.core.DataRequiredError: `MapKeyToFloat` transform requires non-empty data.

this is otherwise the error when fit is called implicitly get_next_trials()

Exception has occurred: DataRequiredError
`MapKeyToFloat` transform requires non-empty data.
  File "/home/cesare/Projects/StrangeNet/crash/ax_crash_min.py", line 145, in run_trials
    trail_data, stop = ax_client.get_next_trials(max_trials=1)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cesare/Projects/StrangeNet/crash/ax_crash_min.py", line 71, in main
    run_trials(bo_b)   #crash_2: fails with mystical error ( truly, it fails on BO fit )
    ^^^^^^^^^^^^^^^^
  File "/home/cesare/Projects/StrangeNet/crash/ax_crash_min.py", line 180, in <module>
    if __name__ == "__main__": main()
                               ^^^^^^
ax.exceptions.core.DataRequiredError: `MapKeyToFloat` transform requires non-empty data.

By the way, if you point me at the latest branch for this features I'll gladly provide more testing, in case is needed.

[cesare-montresor]

On a side thought, this issue happen because I ended up using the Ax template to hold all experiment configurations, eventually using fixed values, to avoid messing with the search.
However in this case, i use it to control Objectives ( NAS vs multiobjective NAS ).
My reasoning is that while collecting data with sobol the experiment which, shoudn't need any objective, in theory, just tracked metrics. This quirk happens because my "common trials" are genereted by a client with only a sobol step, which is the only one not really requiring an objective.
It is as if the objective would be more tied to the trial generation step, then the experiment itself ... but it still makes more sense as experiment, from a practical prospective.