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

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions