Skip to content

[FEATURE REQUEST]: setting fit_tracking_metrics in configure_generation_strategy #4688

@chanansh

Description

@chanansh

Motivation

Currently, to disable modeling of tracked metrics, one needs to manually edit the strategy. See #4600 for details.

Describe the solution you'd like to see implemented in Ax.

The feature request is to enable easy way to disable modeling of tracking metrics in client.configure_generation_strategy(fit_tracking_metrics=False). This way, one can use the high level interface while disabling modeling the tracked metrics, which usually is not needed and are tracked just for book-keeping.

Describe any alternatives you've considered to the above solution.

Current workaround implementation suggestion

def disable_tracked_metrics_modeling(client: Client) -> None:
    """
    Modify the client's generation strategy to disable tracking metrics from being modeled by GP.
    
    This function modifies the existing generation strategy created by configure_generation_strategy()
    to add fit_tracking_metrics=False to the MBM (Model-Based Method) node's model_kwargs.
    This prevents tracking metrics from being modeled by the Gaussian Process while still storing them.
    
    See: https://github.com/facebook/Ax/issues/4600
    
    Args:
        client: Ax Client instance with an already configured generation strategy
    """
    from ax.adapter.registry import Generators
    
    gs = client._generation_strategy
    logger.info("Disabling tracking metrics from being modeled by GP (fit_tracking_metrics=False)")
    
    # Find the MBM (Model-Based Method) node - it's typically the last node
    # The MBM node uses BoTorch_MODULAR generator
    modified = False
    for node in reversed(gs._nodes):
        for spec in node.generator_specs:
            # Check if this is the BoTorch_MODULAR generator (MBM)
            if hasattr(spec, 'generator_enum') and spec.generator_enum == Generators.BOTORCH_MODULAR:
                # Add fit_tracking_metrics=False to model_kwargs
                if not hasattr(spec, 'model_kwargs') or spec.model_kwargs is None:
                    spec.model_kwargs = {}
                spec.model_kwargs['fit_tracking_metrics'] = False
                logger.info(f"  Modified {node.name} node: fit_tracking_metrics=False")
                modified = True
                break
        if modified:
            break
    
    if modified:
        # Set the modified generation strategy back
        client.set_generation_strategy(generation_strategy=gs)
        logger.info("  ✓ Generation strategy updated")
    else:
        logger.warning("  ⚠ Could not find MBM node with BoTorch_MODULAR generator")

Is this related to an existing issue in Ax or another repository? If so please include links to those Issues here.

No response

Code of Conduct

  • I agree to follow Ax's Code of Conduct

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions