From c6c45b696f01d05e8f8ba4b8746bec0fe42e2a65 Mon Sep 17 00:00:00 2001 From: Daniel Cohen Date: Mon, 2 Dec 2024 08:30:00 -0800 Subject: [PATCH] Move telemetry out of OSS ax (#3089) Summary: Pull Request resolved: https://github.com/facebook/Ax/pull/3089 It currently isn't providing value in OSS and preventing us from reflecting non OSS features Reviewed By: mpolson64 Differential Revision: D66239268 --- ax/telemetry/__init__.py | 5 - ax/telemetry/ax_client.py | 132 ---- ax/telemetry/common.py | 113 ---- ax/telemetry/experiment.py | 328 ---------- ax/telemetry/generation_strategy.py | 85 --- ax/telemetry/optimization.py | 572 ------------------ ax/telemetry/scheduler.py | 166 ----- ax/telemetry/tests/test_ax_client.py | 173 ------ ax/telemetry/tests/test_experiment.py | 128 ---- .../tests/test_generation_strategy.py | 49 -- ax/telemetry/tests/test_optimization.py | 292 --------- ax/telemetry/tests/test_scheduler.py | 280 --------- sphinx/source/telemetry.rst | 56 -- 13 files changed, 2379 deletions(-) delete mode 100644 ax/telemetry/__init__.py delete mode 100644 ax/telemetry/ax_client.py delete mode 100644 ax/telemetry/common.py delete mode 100644 ax/telemetry/experiment.py delete mode 100644 ax/telemetry/generation_strategy.py delete mode 100644 ax/telemetry/optimization.py delete mode 100644 ax/telemetry/scheduler.py delete mode 100644 ax/telemetry/tests/test_ax_client.py delete mode 100644 ax/telemetry/tests/test_experiment.py delete mode 100644 ax/telemetry/tests/test_generation_strategy.py delete mode 100644 ax/telemetry/tests/test_optimization.py delete mode 100644 ax/telemetry/tests/test_scheduler.py delete mode 100644 sphinx/source/telemetry.rst diff --git a/ax/telemetry/__init__.py b/ax/telemetry/__init__.py deleted file mode 100644 index 4b87eb9e4d0..00000000000 --- a/ax/telemetry/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. diff --git a/ax/telemetry/ax_client.py b/ax/telemetry/ax_client.py deleted file mode 100644 index 13df668ca25..00000000000 --- a/ax/telemetry/ax_client.py +++ /dev/null @@ -1,132 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from __future__ import annotations - -from dataclasses import asdict, dataclass -from typing import Any - -from ax.service.ax_client import AxClient -from ax.telemetry.common import _get_max_transformed_dimensionality -from ax.telemetry.experiment import ExperimentCompletedRecord, ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord - - -@dataclass(frozen=True) -class AxClientCreatedRecord: - """ - Record of the AxClient creation event. This can be used for telemetry in settings - where many AxClients are being created either manually or programatically. In - order to facilitate easy serialization only include simple types: numbers, strings, - bools, and None. - """ - - experiment_created_record: ExperimentCreatedRecord - generation_strategy_created_record: GenerationStrategyCreatedRecord - - arms_per_trial: int - early_stopping_strategy_cls: str | None - global_stopping_strategy_cls: str | None - - # Dimensionality of transformed SearchSpace can often be much higher due to one-hot - # encoding of unordered ChoiceParameters - transformed_dimensionality: int | None - - @classmethod - def from_ax_client(cls, ax_client: AxClient) -> AxClientCreatedRecord: - # Some AxClients may implement `batch_size`, those that do not use - # one trial arms. - if getattr(ax_client, "batch_size", None) is not None: - # pyre-fixme[16] `AxClient` has no attribute `batch_size` - arms_per_trial = ax_client.batch_size - else: - arms_per_trial = 1 - - return cls( - experiment_created_record=ExperimentCreatedRecord.from_experiment( - experiment=ax_client.experiment - ), - generation_strategy_created_record=( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=ax_client.generation_strategy - ) - ), - arms_per_trial=arms_per_trial, - early_stopping_strategy_cls=( - None - if ax_client.early_stopping_strategy is None - else ax_client.early_stopping_strategy.__class__.__name__ - ), - global_stopping_strategy_cls=( - None - if ax_client.global_stopping_strategy is None - else ax_client.global_stopping_strategy.__class__.__name__ - ), - transformed_dimensionality=_get_max_transformed_dimensionality( - search_space=ax_client.experiment.search_space, - generation_strategy=ax_client.generation_strategy, - ), - ) - - def flatten(self) -> dict[str, Any]: - """ - Flatten into an appropriate format for logging to a tabular database. - """ - - self_dict = asdict(self) - experiment_created_record_dict = self_dict.pop("experiment_created_record") - generation_strategy_created_record_dict = self_dict.pop( - "generation_strategy_created_record" - ) - - return { - **self_dict, - **experiment_created_record_dict, - **generation_strategy_created_record_dict, - } - - -@dataclass(frozen=True) -class AxClientCompletedRecord: - """ - Record of the AxClient completion event. This will have information only available - after the optimization has completed. - """ - - experiment_completed_record: ExperimentCompletedRecord - - best_point_quality: float - model_fit_quality: float - model_std_quality: float - model_fit_generalization: float - model_std_generalization: float - - @classmethod - def from_ax_client(cls, ax_client: AxClient) -> AxClientCompletedRecord: - return cls( - experiment_completed_record=ExperimentCompletedRecord.from_experiment( - experiment=ax_client.experiment - ), - best_point_quality=float("nan"), # TODO[T147907632] - model_fit_quality=float("nan"), # TODO[T147907632] - model_std_quality=float("nan"), - model_fit_generalization=float("nan"), - model_std_generalization=float("nan"), - ) - - def flatten(self) -> dict[str, Any]: - """ - Flatten into an appropriate format for logging to a tabular database. - """ - - self_dict = asdict(self) - experiment_completed_record_dict = self_dict.pop("experiment_completed_record") - - return { - **self_dict, - **experiment_completed_record_dict, - } diff --git a/ax/telemetry/common.py b/ax/telemetry/common.py deleted file mode 100644 index 267e5b88b70..00000000000 --- a/ax/telemetry/common.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import warnings -from datetime import datetime -from typing import Any - -from ax.core.experiment import Experiment -from ax.core.parameter import FixedParameter -from ax.exceptions.core import AxWarning -from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy -from ax.modelbridge.modelbridge_utils import ( - extract_search_space_digest, - transform_search_space, -) -from ax.modelbridge.registry import ModelRegistryBase, Models, SearchSpace -from ax.modelbridge.transforms.base import Transform -from ax.modelbridge.transforms.cast import Cast - -# Models whose generated trails will count towards initialization_trials -INITIALIZATION_MODELS: list[Models] = [Models.SOBOL, Models.UNIFORM] - -# Models whose generated trails will count towards other_trials -OTHER_MODELS: list[Models] = [] -# Product surface to use if none is provided -DEFAULT_PRODUCT_SURFACE = "unknown" - - -def _get_max_transformed_dimensionality( - search_space: SearchSpace, generation_strategy: GenerationStrategy -) -> int | None: - """ - Get dimensionality of transformed SearchSpace for all steps in the - GenerationStrategy and return the maximum. - """ - if generation_strategy.is_node_based: - warnings.warn( - "`_get_max_transformed_dimensionality` does not fully support node-based " - "generation strategies. This will result in an incomplete record.", - category=AxWarning, - stacklevel=4, - ) - # TODO [T192965545]: Support node-based generation strategies in telemetry - return None - - transforms_by_step = [ - _extract_transforms_and_configs(step=step) - for step in generation_strategy._steps - ] - - transformed_search_spaces = [ - transform_search_space( - search_space=search_space, - transforms=[Cast] + transforms, - transform_configs=transform_configs, - ) - for transforms, transform_configs in transforms_by_step - ] - - # The length of the bounds of a SearchSpaceDigest is equal to the number of - # dimensions present. - dimensionalities = [ - len( - extract_search_space_digest( - search_space=tf_search_space, - param_names=[ - name - for name, param in tf_search_space.parameters.items() - if not isinstance(param, FixedParameter) - ], - ).bounds - ) - for tf_search_space in transformed_search_spaces - ] - - return max(dimensionalities) - - -def _extract_transforms_and_configs( - step: GenerationStep, -) -> tuple[list[type[Transform]], dict[str, Any]]: - """ - Extract Transforms and their configs from the GenerationStep. Prefer kwargs - provided over the model's defaults. - """ - - kwargs = step.model_spec.model_kwargs or {} - transforms = kwargs.get("transforms") - transform_configs = kwargs.get("transform_configs") - - if transforms is not None and transform_configs is not None: - return transforms, transform_configs - - model = step.model - if isinstance(model, ModelRegistryBase): - _, bridge_kwargs = model.view_defaults() - transforms = transforms or bridge_kwargs.get("transforms") - transform_configs = transform_configs or bridge_kwargs.get("transform_configs") - - return (transforms or [], transform_configs or {}) - - -def get_unique_identifier(experiment: Experiment) -> str: - """ - Return a unique identifier for an experiment so creation and completion - events can be joined. - """ - str_time = datetime.strftime(experiment.time_created, "%Y-%m-%d %H:%M:%S") - return f"{experiment.name}_{str_time}" diff --git a/ax/telemetry/experiment.py b/ax/telemetry/experiment.py deleted file mode 100644 index cddc2a0a002..00000000000 --- a/ax/telemetry/experiment.py +++ /dev/null @@ -1,328 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from __future__ import annotations - -from dataclasses import dataclass -from math import inf - -from ax.core.base_trial import TrialStatus - -from ax.core.experiment import Experiment -from ax.core.map_metric import MapMetric -from ax.core.parameter import ( - ChoiceParameter, - FixedParameter, - ParameterType, - RangeParameter, -) -from ax.core.search_space import HierarchicalSearchSpace, SearchSpace -from ax.core.utils import get_model_times -from ax.telemetry.common import INITIALIZATION_MODELS, OTHER_MODELS - -INITIALIZATION_MODEL_STRS: list[str] = [enum.value for enum in INITIALIZATION_MODELS] -OTHER_MODEL_STRS: list[str] = [enum.value for enum in OTHER_MODELS] + [None] - - -@dataclass(frozen=True) -class ExperimentCreatedRecord: - """ - Record of the Experiment creation event. This can be used for telemetry in settings - where many Experiments are being created either manually or programatically. In - order to facilitate easy serialization only include simple types: numbers, strings, - bools, and None. - """ - - experiment_name: str | None - experiment_type: str | None - - # SearchSpace info - num_continuous_range_parameters: int - - # Note: ordered ChoiceParameters and int RangeParameters should both utilize the - # following fields - num_int_range_parameters_small: int # 2 - 3 elements - num_int_range_parameters_medium: int # 4 - 7 elements - num_int_range_parameters_large: int # 8 or more elements - - # Any RangeParameter can specify log space sampling - num_log_scale_range_parameters: int - - num_unordered_choice_parameters_small: int # 2 - 3 elements - num_unordered_choice_parameters_medium: int # 4 - 7 elements - num_unordered_choice_parameters_large: int # 8 or more elements - - num_fixed_parameters: int - - dimensionality: int - hierarchical_tree_height: int # Height of tree for HSS, 1 for base SearchSpace - num_parameter_constraints: int - - # OptimizationConfig info - num_objectives: int - num_tracking_metrics: int - num_outcome_constraints: int # Includes ObjectiveThresholds in MOO - - # General Metrics info - num_map_metrics: int - metric_cls_to_quantity: dict[str, int] - - # Runner info - runner_cls: str - - @classmethod - def from_experiment(cls, experiment: Experiment) -> ExperimentCreatedRecord: - ( - num_continuous_range_parameters, - num_int_range_parameters_small, - num_int_range_parameters_medium, - num_int_range_parameters_large, - num_log_scale_range_parameters, - num_unordered_choice_parameters_small, - num_unordered_choice_parameters_medium, - num_unordered_choice_parameters_large, - num_fixed_parameters, - ) = cls._get_param_counts_from_search_space( - search_space=experiment.search_space - ) - - return cls( - experiment_name=experiment._name, - experiment_type=experiment._experiment_type, - num_continuous_range_parameters=num_continuous_range_parameters, - num_int_range_parameters_small=num_int_range_parameters_small, - num_int_range_parameters_medium=num_int_range_parameters_medium, - num_int_range_parameters_large=num_int_range_parameters_large, - num_log_scale_range_parameters=num_log_scale_range_parameters, - num_unordered_choice_parameters_small=num_unordered_choice_parameters_small, - num_unordered_choice_parameters_medium=( - num_unordered_choice_parameters_medium - ), - num_unordered_choice_parameters_large=num_unordered_choice_parameters_large, - num_fixed_parameters=num_fixed_parameters, - dimensionality=sum( - 1 for param in experiment.parameters.values() if param.cardinality() > 1 - ), - hierarchical_tree_height=( - experiment.search_space.height - if isinstance(experiment.search_space, HierarchicalSearchSpace) - else 1 - ), - num_parameter_constraints=len( - experiment.search_space.parameter_constraints - ), - num_objectives=( - len(experiment.optimization_config.objective.metrics) - if experiment.optimization_config is not None - else 0 - ), - num_tracking_metrics=len(experiment.tracking_metrics), - num_outcome_constraints=( - len(experiment.optimization_config.outcome_constraints) - if experiment.optimization_config is not None - else 0 - ), - num_map_metrics=sum( - 1 - for metric in experiment.metrics.values() - if isinstance(metric, MapMetric) - ), - metric_cls_to_quantity={ - cls_name: sum( - 1 - for metric in experiment.metrics.values() - if metric.__class__.__name__ == cls_name - ) - for cls_name in { - metric.__class__.__name__ for metric in experiment.metrics.values() - } - }, - runner_cls=experiment.runner.__class__.__name__, - ) - - @staticmethod - def _get_param_counts_from_search_space( - search_space: SearchSpace, - ) -> tuple[int, int, int, int, int, int, int, int, int]: - """ - Return counts of different types of parameters. - - returns: - num_continuous_range_parameters - - num_int_range_parameters_small - num_int_range_parameters_medium - num_int_range_parameters_large - - num_log_scale_range_parameters - - num_unordered_choice_parameters_small - num_unordered_choice_parameters_medium - num_unordered_choice_parameters_large - - num_fixed_parameters - """ - - num_continuous_range_parameters = sum( - 1 - for param in search_space.parameters.values() - if isinstance(param, RangeParameter) - and param.parameter_type == ParameterType.FLOAT - ) - num_int_range_parameters_small = sum( - 1 - for param in search_space.parameters.values() - if ( - isinstance(param, RangeParameter) - or (isinstance(param, ChoiceParameter) and param.is_ordered) - ) - and (1.0 < param.cardinality() <= 3.0) - ) - num_int_range_parameters_medium = sum( - 1 - for param in search_space.parameters.values() - if ( - isinstance(param, RangeParameter) - or (isinstance(param, ChoiceParameter) and param.is_ordered) - ) - and (3.0 < param.cardinality() <= 7.0) - ) - num_int_range_parameters_large = sum( - 1 - for param in search_space.parameters.values() - if ( - isinstance(param, RangeParameter) - or (isinstance(param, ChoiceParameter) and param.is_ordered) - ) - and (7.0 < param.cardinality() < inf) - ) - num_log_scale_range_parameters = sum( - 1 - for param in search_space.parameters.values() - if isinstance(param, RangeParameter) and param.log_scale - ) - num_unordered_choice_parameters_small = sum( - 1 - for param in search_space.parameters.values() - if (isinstance(param, ChoiceParameter) and not param.is_ordered) - and (1 < param.cardinality() <= 3) - ) - num_unordered_choice_parameters_medium = sum( - 1 - for param in search_space.parameters.values() - if (isinstance(param, ChoiceParameter) and not param.is_ordered) - and (3 < param.cardinality() <= 7) - ) - num_unordered_choice_parameters_large = sum( - 1 - for param in search_space.parameters.values() - if (isinstance(param, ChoiceParameter) and not param.is_ordered) - and param.cardinality() > 7 - ) - num_fixed_parameters = sum( - 1 - for param in search_space.parameters.values() - if isinstance(param, FixedParameter) - ) - - return ( - num_continuous_range_parameters, - num_int_range_parameters_small, - num_int_range_parameters_medium, - num_int_range_parameters_large, - num_log_scale_range_parameters, - num_unordered_choice_parameters_small, - num_unordered_choice_parameters_medium, - num_unordered_choice_parameters_large, - num_fixed_parameters, - ) - - -@dataclass(frozen=True) -class ExperimentCompletedRecord: - """ - Record of the Experiment completion event. This can be used for telemetry in - settings where many Experiments are being created either manually or - programatically. In order to facilitate easy serialization only include simple - types: numbers, strings, bools, and None. - """ - - num_initialization_trials: int - num_bayesopt_trials: int - num_other_trials: int - - num_completed_trials: int - num_failed_trials: int - num_abandoned_trials: int - num_early_stopped_trials: int - - total_fit_time: int - total_gen_time: int - - # OptimizationConfig info which might be updated for human in the - # loop experiments - num_objectives: int - num_tracking_metrics: int - num_outcome_constraints: int # Includes ObjectiveThresholds in MOO - - @classmethod - def from_experiment(cls, experiment: Experiment) -> ExperimentCompletedRecord: - trial_count_by_status = { - status: len(trials) - for status, trials in experiment.trials_by_status.items() - } - ignored_statuses = {TrialStatus.STAGED, TrialStatus.CANDIDATE} - - model_keys = [ - [gr._model_key for gr in trial.generator_runs] - for trial in experiment.trials.values() - if trial.status not in ignored_statuses - ] - - fit_time, gen_time = get_model_times(experiment=experiment) - - return cls( - num_initialization_trials=sum( - 1 - for model_key_list in model_keys - if all( - model_key in INITIALIZATION_MODEL_STRS - for model_key in model_key_list - ) - ), - num_bayesopt_trials=sum( - 1 - for model_key_list in model_keys - if any( - model_key not in INITIALIZATION_MODEL_STRS - and model_key not in OTHER_MODEL_STRS - for model_key in model_key_list - ) - ), - num_other_trials=sum( - 1 - for model_key_list in model_keys - if all(model_key in OTHER_MODEL_STRS for model_key in model_key_list) - ), - num_completed_trials=trial_count_by_status[TrialStatus.COMPLETED], - num_failed_trials=trial_count_by_status[TrialStatus.FAILED], - num_abandoned_trials=trial_count_by_status[TrialStatus.ABANDONED], - num_early_stopped_trials=trial_count_by_status[TrialStatus.EARLY_STOPPED], - total_fit_time=int(fit_time), - total_gen_time=int(gen_time), - num_objectives=( - len(experiment.optimization_config.objective.metrics) - if experiment.optimization_config is not None - else 0 - ), - num_tracking_metrics=len(experiment.tracking_metrics), - num_outcome_constraints=( - len(experiment.optimization_config.outcome_constraints) - if experiment.optimization_config is not None - else 0 - ), - ) diff --git a/ax/telemetry/generation_strategy.py b/ax/telemetry/generation_strategy.py deleted file mode 100644 index 2bfb760aa4c..00000000000 --- a/ax/telemetry/generation_strategy.py +++ /dev/null @@ -1,85 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from __future__ import annotations - -import warnings -from dataclasses import dataclass -from math import inf - -from ax.exceptions.core import AxWarning -from ax.modelbridge.generation_strategy import GenerationStrategy -from ax.telemetry.common import INITIALIZATION_MODELS, OTHER_MODELS - - -@dataclass(frozen=True) -class GenerationStrategyCreatedRecord: - """ - Record of the GenerationStrategy creation event. This can be used for telemetry in - settings where many GenerationStrategy are being created either manually or - programatically. In order to facilitate easy serialization only include simple - types: numbers, strings, bools, and None. - """ - - generation_strategy_name: str - - # -1 indicates unlimited trials requested, 0 indicates no trials requested - num_requested_initialization_trials: None | ( - int # Typically the number of Sobol trials - ) - num_requested_bayesopt_trials: int | None - num_requested_other_trials: int | None - - # Minimum `max_parallelism` across GenerationSteps, i.e. the bottleneck - max_parallelism: int | None - - @classmethod - def from_generation_strategy( - cls, generation_strategy: GenerationStrategy - ) -> GenerationStrategyCreatedRecord: - if generation_strategy.is_node_based: - warnings.warn( - "`GenerationStrategyCreatedRecord` does not fully support node-based " - "generation strategies. This will result in an incomplete record.", - category=AxWarning, - stacklevel=4, - ) - # TODO [T192965545]: Support node-based generation strategies in telemetry - return cls( - generation_strategy_name=generation_strategy.name, - num_requested_initialization_trials=None, - num_requested_bayesopt_trials=None, - num_requested_other_trials=None, - max_parallelism=None, - ) - - # Minimum `max_parallelism` across GenerationSteps, i.e. the bottleneck - true_max_parallelism = min( - step.max_parallelism or inf for step in generation_strategy._steps - ) - - return cls( - generation_strategy_name=generation_strategy.name, - num_requested_initialization_trials=sum( - step.num_trials - for step in generation_strategy._steps - if step.model in INITIALIZATION_MODELS - ), - num_requested_bayesopt_trials=sum( - step.num_trials - for step in generation_strategy._steps - if step.model not in INITIALIZATION_MODELS + OTHER_MODELS - ), - num_requested_other_trials=sum( - step.num_trials - for step in generation_strategy._steps - if step.model in OTHER_MODELS - ), - max_parallelism=( - true_max_parallelism if isinstance(true_max_parallelism, int) else -1 - ), - ) diff --git a/ax/telemetry/optimization.py b/ax/telemetry/optimization.py deleted file mode 100644 index 95546caf4d9..00000000000 --- a/ax/telemetry/optimization.py +++ /dev/null @@ -1,572 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from __future__ import annotations - -from dataclasses import dataclass - -from ax.core.experiment import Experiment -from ax.modelbridge.generation_strategy import GenerationStrategy - -from ax.service.ax_client import AxClient -from ax.service.scheduler import Scheduler -from ax.telemetry.ax_client import AxClientCompletedRecord, AxClientCreatedRecord -from ax.telemetry.common import ( - _get_max_transformed_dimensionality, - DEFAULT_PRODUCT_SURFACE, -) -from ax.telemetry.experiment import ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord -from ax.telemetry.scheduler import SchedulerCompletedRecord, SchedulerCreatedRecord - - -@dataclass(frozen=True) -class OptimizationCreatedRecord: - """ - Record of the "Optimization" creation event. This can come from either an - AxClient or a Scheduler. This Record is especially useful for logging Ax-backed - optimization results to a tabular database (i.e. one row per Record). - """ - - unique_identifier: str - owner: str - - # ExperimentCreatedRecord fields - experiment_name: str | None - experiment_type: str | None - num_continuous_range_parameters: int - num_int_range_parameters_small: int - num_int_range_parameters_medium: int - num_int_range_parameters_large: int - num_log_scale_range_parameters: int - num_unordered_choice_parameters_small: int - num_unordered_choice_parameters_medium: int - num_unordered_choice_parameters_large: int - num_fixed_parameters: int - dimensionality: int - hierarchical_tree_height: int - num_parameter_constraints: int - num_objectives: int - num_tracking_metrics: int - num_outcome_constraints: int - num_map_metrics: int - metric_cls_to_quantity: dict[str, int] - runner_cls: str - - # GenerationStrategyCreatedRecord fields - generation_strategy_name: str | None - num_requested_initialization_trials: int | None - num_requested_bayesopt_trials: int | None - num_requested_other_trials: int | None - max_parallelism: int | None - - # {AxClient, Scheduler}CreatedRecord fields - early_stopping_strategy_cls: str | None - global_stopping_strategy_cls: str | None - transformed_dimensionality: int | None - scheduler_total_trials: int | None - scheduler_max_pending_trials: int - arms_per_trial: int - - # Top level info - product_surface: str - launch_surface: str - - deployed_job_id: int - trial_evaluation_identifier: str | None - - # Miscellaneous product info - is_manual_generation_strategy: bool - warm_started_from: str | None - num_custom_trials: int - support_tier: str - - @classmethod - def from_scheduler( - cls, - scheduler: Scheduler, - unique_identifier: str, - owner: str, - product_surface: str, - launch_surface: str, - deployed_job_id: int, - trial_evaluation_identifier: str | None, - is_manual_generation_strategy: bool, - warm_started_from: str | None, - num_custom_trials: int, - support_tier: str, - ) -> OptimizationCreatedRecord: - scheduler_created_record = SchedulerCreatedRecord.from_scheduler( - scheduler=scheduler - ) - - experiment_created_record = scheduler_created_record.experiment_created_record - generation_strategy_created_record = ( - scheduler_created_record.generation_strategy_created_record - ) - - return cls( - experiment_name=experiment_created_record.experiment_name, - experiment_type=experiment_created_record.experiment_type, - num_continuous_range_parameters=( - experiment_created_record.num_continuous_range_parameters - ), - num_int_range_parameters_small=( - experiment_created_record.num_int_range_parameters_small - ), - num_int_range_parameters_medium=( - experiment_created_record.num_int_range_parameters_medium - ), - num_int_range_parameters_large=( - experiment_created_record.num_int_range_parameters_large - ), - num_log_scale_range_parameters=( - experiment_created_record.num_log_scale_range_parameters - ), - num_unordered_choice_parameters_small=( - experiment_created_record.num_unordered_choice_parameters_small - ), - num_unordered_choice_parameters_medium=( - experiment_created_record.num_unordered_choice_parameters_medium - ), - num_unordered_choice_parameters_large=( - experiment_created_record.num_unordered_choice_parameters_large - ), - num_fixed_parameters=experiment_created_record.num_fixed_parameters, - dimensionality=experiment_created_record.dimensionality, - hierarchical_tree_height=experiment_created_record.hierarchical_tree_height, - num_parameter_constraints=( - experiment_created_record.num_parameter_constraints - ), - num_objectives=experiment_created_record.num_objectives, - num_tracking_metrics=experiment_created_record.num_tracking_metrics, - num_outcome_constraints=experiment_created_record.num_outcome_constraints, - num_map_metrics=experiment_created_record.num_map_metrics, - metric_cls_to_quantity=experiment_created_record.metric_cls_to_quantity, - runner_cls=experiment_created_record.runner_cls, - generation_strategy_name=( - generation_strategy_created_record.generation_strategy_name - ), - num_requested_initialization_trials=( - generation_strategy_created_record.num_requested_initialization_trials - ), - num_requested_bayesopt_trials=( - generation_strategy_created_record.num_requested_bayesopt_trials - ), - num_requested_other_trials=( - generation_strategy_created_record.num_requested_other_trials - ), - max_parallelism=generation_strategy_created_record.max_parallelism, - early_stopping_strategy_cls=( - scheduler_created_record.early_stopping_strategy_cls - ), - global_stopping_strategy_cls=( - scheduler_created_record.global_stopping_strategy_cls - ), - transformed_dimensionality=( - scheduler_created_record.transformed_dimensionality - ), - scheduler_total_trials=scheduler_created_record.scheduler_total_trials, - scheduler_max_pending_trials=( - scheduler_created_record.scheduler_max_pending_trials - ), - arms_per_trial=scheduler_created_record.arms_per_trial, - unique_identifier=unique_identifier, - owner=owner, - product_surface=product_surface, - launch_surface=launch_surface, - deployed_job_id=deployed_job_id, - trial_evaluation_identifier=trial_evaluation_identifier, - is_manual_generation_strategy=is_manual_generation_strategy, - warm_started_from=warm_started_from, - num_custom_trials=num_custom_trials, - support_tier=support_tier, - ) - - @classmethod - def from_ax_client( - cls, - ax_client: AxClient, - unique_identifier: str, - owner: str, - product_surface: str, - launch_surface: str, - deployed_job_id: int, - trial_evaluation_identifier: str | None, - is_manual_generation_strategy: bool, - warm_started_from: str | None, - num_custom_trials: int, - ) -> OptimizationCreatedRecord: - ax_client_created_record = AxClientCreatedRecord.from_ax_client( - ax_client=ax_client - ) - - experiment_created_record = ax_client_created_record.experiment_created_record - generation_strategy_created_record = ( - ax_client_created_record.generation_strategy_created_record - ) - - return cls( - experiment_name=experiment_created_record.experiment_name, - experiment_type=experiment_created_record.experiment_type, - num_continuous_range_parameters=( - experiment_created_record.num_continuous_range_parameters - ), - num_int_range_parameters_small=( - experiment_created_record.num_int_range_parameters_small - ), - num_int_range_parameters_medium=( - experiment_created_record.num_int_range_parameters_medium - ), - num_int_range_parameters_large=( - experiment_created_record.num_int_range_parameters_large - ), - num_log_scale_range_parameters=( - experiment_created_record.num_log_scale_range_parameters - ), - num_unordered_choice_parameters_small=( - experiment_created_record.num_unordered_choice_parameters_small - ), - num_unordered_choice_parameters_medium=( - experiment_created_record.num_unordered_choice_parameters_medium - ), - num_unordered_choice_parameters_large=( - experiment_created_record.num_unordered_choice_parameters_large - ), - num_fixed_parameters=experiment_created_record.num_fixed_parameters, - dimensionality=experiment_created_record.dimensionality, - hierarchical_tree_height=( - experiment_created_record.hierarchical_tree_height - ), - num_parameter_constraints=( - experiment_created_record.num_parameter_constraints - ), - num_objectives=experiment_created_record.num_objectives, - num_tracking_metrics=experiment_created_record.num_tracking_metrics, - num_outcome_constraints=experiment_created_record.num_outcome_constraints, - num_map_metrics=experiment_created_record.num_map_metrics, - metric_cls_to_quantity=experiment_created_record.metric_cls_to_quantity, - runner_cls=experiment_created_record.runner_cls, - generation_strategy_name=( - generation_strategy_created_record.generation_strategy_name - ), - num_requested_initialization_trials=( - generation_strategy_created_record.num_requested_initialization_trials - ), - num_requested_bayesopt_trials=( - generation_strategy_created_record.num_requested_bayesopt_trials - ), - num_requested_other_trials=( - generation_strategy_created_record.num_requested_other_trials - ), - max_parallelism=generation_strategy_created_record.max_parallelism, - early_stopping_strategy_cls=( - ax_client_created_record.early_stopping_strategy_cls - ), - global_stopping_strategy_cls=( - ax_client_created_record.global_stopping_strategy_cls - ), - transformed_dimensionality=( - ax_client_created_record.transformed_dimensionality - ), - arms_per_trial=ax_client_created_record.arms_per_trial, - unique_identifier=unique_identifier, - owner=owner, - product_surface=product_surface, - launch_surface=launch_surface, - deployed_job_id=deployed_job_id, - trial_evaluation_identifier=trial_evaluation_identifier, - is_manual_generation_strategy=is_manual_generation_strategy, - warm_started_from=warm_started_from, - num_custom_trials=num_custom_trials, - # The following are not applicable for AxClient - scheduler_total_trials=None, - scheduler_max_pending_trials=-1, - support_tier="", # This support may be added in the future - ) - - @classmethod - def from_experiment( - cls, - experiment: Experiment, - generation_strategy: GenerationStrategy | None, - unique_identifier: str, - owner: str, - product_surface: str, - launch_surface: str, - deployed_job_id: int, - is_manual_generation_strategy: bool, - num_custom_trials: int, - warm_started_from: str | None = None, - arms_per_trial: int | None = None, - trial_evaluation_identifier: str | None = None, - ) -> OptimizationCreatedRecord: - experiment_created_record = ExperimentCreatedRecord.from_experiment( - experiment=experiment, - ) - generation_strategy_created_record = ( - None - if generation_strategy is None - else ( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=generation_strategy, - ) - ) - ) - arms_per_trial = -1 if arms_per_trial is None else arms_per_trial - product_surface = ( - DEFAULT_PRODUCT_SURFACE if product_surface is None else product_surface - ) - - num_requested_initialization_trials = ( - None - if generation_strategy_created_record is None - else generation_strategy_created_record.num_requested_initialization_trials - ) - return cls( - experiment_name=experiment_created_record.experiment_name, - experiment_type=experiment_created_record.experiment_type, - num_continuous_range_parameters=( - experiment_created_record.num_continuous_range_parameters - ), - num_int_range_parameters_small=( - experiment_created_record.num_int_range_parameters_small - ), - num_int_range_parameters_medium=( - experiment_created_record.num_int_range_parameters_medium - ), - num_int_range_parameters_large=( - experiment_created_record.num_int_range_parameters_large - ), - num_log_scale_range_parameters=( - experiment_created_record.num_log_scale_range_parameters - ), - num_unordered_choice_parameters_small=( - experiment_created_record.num_unordered_choice_parameters_small - ), - num_unordered_choice_parameters_medium=( - experiment_created_record.num_unordered_choice_parameters_medium - ), - num_unordered_choice_parameters_large=( - experiment_created_record.num_unordered_choice_parameters_large - ), - num_fixed_parameters=experiment_created_record.num_fixed_parameters, - dimensionality=experiment_created_record.dimensionality, - hierarchical_tree_height=( - experiment_created_record.hierarchical_tree_height - ), - num_parameter_constraints=( - experiment_created_record.num_parameter_constraints - ), - num_objectives=experiment_created_record.num_objectives, - num_tracking_metrics=experiment_created_record.num_tracking_metrics, - num_outcome_constraints=experiment_created_record.num_outcome_constraints, - num_map_metrics=experiment_created_record.num_map_metrics, - metric_cls_to_quantity=experiment_created_record.metric_cls_to_quantity, - runner_cls=experiment_created_record.runner_cls, - generation_strategy_name=( - None - if generation_strategy_created_record is None - else generation_strategy_created_record.generation_strategy_name - ), - num_requested_initialization_trials=num_requested_initialization_trials, - num_requested_bayesopt_trials=( - None - if generation_strategy_created_record is None - else generation_strategy_created_record.num_requested_bayesopt_trials - ), - num_requested_other_trials=( - None - if generation_strategy_created_record is None - else generation_strategy_created_record.num_requested_other_trials - ), - max_parallelism=( - None - if generation_strategy_created_record is None - else generation_strategy_created_record.max_parallelism - ), - early_stopping_strategy_cls=None, - global_stopping_strategy_cls=None, - transformed_dimensionality=( - None - if generation_strategy is None - else _get_max_transformed_dimensionality( - search_space=experiment.search_space, - generation_strategy=generation_strategy, - ) - ), - arms_per_trial=arms_per_trial, - unique_identifier=unique_identifier, - owner=owner, - product_surface=product_surface, - launch_surface=launch_surface, - deployed_job_id=deployed_job_id, - trial_evaluation_identifier=trial_evaluation_identifier, - is_manual_generation_strategy=is_manual_generation_strategy, - warm_started_from=warm_started_from, - num_custom_trials=num_custom_trials, - # The following are not applicable for AxClient - scheduler_total_trials=None, - scheduler_max_pending_trials=-1, - support_tier="", # This support may be added in the future - ) - - -@dataclass(frozen=True) -class OptimizationCompletedRecord: - """ - Record of the "Optimization" completion event. This can come from either an - AxClient or a Scheduler. This Record is especially useful for logging Ax-backed - optimization results to a tabular database (i.e. one row per Record) - """ - - unique_identifier: str - - # ExperimentCompletedRecord fields - num_initialization_trials: int - num_bayesopt_trials: int - num_other_trials: int - - num_completed_trials: int - num_failed_trials: int - num_abandoned_trials: int - num_early_stopped_trials: int - - total_fit_time: int - total_gen_time: int - - # SchedulerCompletedRecord fields - best_point_quality: float - model_fit_quality: float - model_std_quality: float - model_fit_generalization: float - model_std_generalization: float - - improvement_over_baseline: float - - num_metric_fetch_e_encountered: int - num_trials_bad_due_to_err: int - - # TODO[mpolson64] Deprecate this field as it is redundant with unique_identifier - deployed_job_id: int | None - - # Miscellaneous deployment specific info - estimated_early_stopping_savings: float - estimated_global_stopping_savings: float - - # OptimizationConfig info which might be updated for human in the - # loop experiments - num_objectives: int - num_tracking_metrics: int - num_outcome_constraints: int # Includes ObjectiveThresholds in MOO - - @classmethod - def from_scheduler( - cls, - scheduler: Scheduler, - unique_identifier: str, - deployed_job_id: int | None, - estimated_early_stopping_savings: float, - estimated_global_stopping_savings: float, - ) -> OptimizationCompletedRecord: - scheduler_completed_record = SchedulerCompletedRecord.from_scheduler( - scheduler=scheduler - ) - experiment_completed_record = ( - scheduler_completed_record.experiment_completed_record - ) - - return cls( - num_initialization_trials=( - experiment_completed_record.num_initialization_trials - ), - num_bayesopt_trials=experiment_completed_record.num_bayesopt_trials, - num_other_trials=experiment_completed_record.num_other_trials, - num_completed_trials=experiment_completed_record.num_completed_trials, - num_failed_trials=experiment_completed_record.num_failed_trials, - num_abandoned_trials=experiment_completed_record.num_abandoned_trials, - num_early_stopped_trials=( - experiment_completed_record.num_early_stopped_trials - ), - total_fit_time=experiment_completed_record.total_fit_time, - total_gen_time=experiment_completed_record.total_gen_time, - best_point_quality=scheduler_completed_record.best_point_quality, - **_extract_model_fit_dict(scheduler_completed_record), - improvement_over_baseline=( - scheduler_completed_record.improvement_over_baseline - ), - num_metric_fetch_e_encountered=( - scheduler_completed_record.num_metric_fetch_e_encountered - ), - num_trials_bad_due_to_err=( - scheduler_completed_record.num_trials_bad_due_to_err - ), - unique_identifier=unique_identifier, - deployed_job_id=deployed_job_id, - estimated_early_stopping_savings=estimated_early_stopping_savings, - estimated_global_stopping_savings=estimated_global_stopping_savings, - num_objectives=experiment_completed_record.num_objectives, - num_tracking_metrics=experiment_completed_record.num_tracking_metrics, - num_outcome_constraints=experiment_completed_record.num_outcome_constraints, - ) - - @classmethod - def from_ax_client( - cls, - ax_client: AxClient, - unique_identifier: str, - deployed_job_id: int | None, - estimated_early_stopping_savings: float, - estimated_global_stopping_savings: float, - ) -> OptimizationCompletedRecord: - ax_client_completed_record = AxClientCompletedRecord.from_ax_client( - ax_client=ax_client - ) - experiment_completed_record = ( - ax_client_completed_record.experiment_completed_record - ) - - return cls( - num_initialization_trials=( - experiment_completed_record.num_initialization_trials - ), - num_bayesopt_trials=experiment_completed_record.num_bayesopt_trials, - num_other_trials=experiment_completed_record.num_other_trials, - num_completed_trials=experiment_completed_record.num_completed_trials, - num_failed_trials=experiment_completed_record.num_failed_trials, - num_abandoned_trials=experiment_completed_record.num_abandoned_trials, - num_early_stopped_trials=( - experiment_completed_record.num_early_stopped_trials - ), - total_fit_time=experiment_completed_record.total_fit_time, - total_gen_time=experiment_completed_record.total_gen_time, - best_point_quality=ax_client_completed_record.best_point_quality, - **_extract_model_fit_dict(ax_client_completed_record), - unique_identifier=unique_identifier, - deployed_job_id=deployed_job_id, - estimated_early_stopping_savings=estimated_early_stopping_savings, - estimated_global_stopping_savings=estimated_global_stopping_savings, - num_objectives=experiment_completed_record.num_objectives, - num_tracking_metrics=experiment_completed_record.num_tracking_metrics, - num_outcome_constraints=experiment_completed_record.num_outcome_constraints, - # The following are not applicable for AxClient - improvement_over_baseline=float("nan"), - num_metric_fetch_e_encountered=-1, - num_trials_bad_due_to_err=-1, - ) - - -def _extract_model_fit_dict( - completed_record: SchedulerCompletedRecord | AxClientCompletedRecord, -) -> dict[str, float]: - model_fit_names = [ - "model_fit_quality", - "model_std_quality", - "model_fit_generalization", - "model_std_generalization", - ] - return {n: getattr(completed_record, n) for n in model_fit_names} diff --git a/ax/telemetry/scheduler.py b/ax/telemetry/scheduler.py deleted file mode 100644 index 9c8059c8b37..00000000000 --- a/ax/telemetry/scheduler.py +++ /dev/null @@ -1,166 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from __future__ import annotations - -from dataclasses import asdict, dataclass -from typing import Any -from warnings import warn - -from ax.modelbridge.cross_validation import ( - get_fit_and_std_quality_and_generalization_dict, -) - -from ax.service.scheduler import get_fitted_model_bridge, Scheduler -from ax.telemetry.common import _get_max_transformed_dimensionality - -from ax.telemetry.experiment import ExperimentCompletedRecord, ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord - - -@dataclass(frozen=True) -class SchedulerCreatedRecord: - """ - Record of the Scheduler creation event. This can be used for telemetry in settings - where many Schedulers are being created either manually or programatically. In - order to facilitate easy serialization only include simple types: numbers, strings, - bools, and None. - """ - - experiment_created_record: ExperimentCreatedRecord - generation_strategy_created_record: GenerationStrategyCreatedRecord - - # SchedulerOptions info - scheduler_total_trials: int | None - scheduler_max_pending_trials: int - arms_per_trial: int - early_stopping_strategy_cls: str | None - global_stopping_strategy_cls: str | None - - # Dimensionality of transformed SearchSpace can often be much higher due to one-hot - # encoding of unordered ChoiceParameters - transformed_dimensionality: int | None - - @classmethod - def from_scheduler(cls, scheduler: Scheduler) -> SchedulerCreatedRecord: - return cls( - experiment_created_record=ExperimentCreatedRecord.from_experiment( - experiment=scheduler.experiment - ), - generation_strategy_created_record=( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=scheduler.standard_generation_strategy - ) - ), - scheduler_total_trials=scheduler.options.total_trials, - scheduler_max_pending_trials=scheduler.options.max_pending_trials, - # If batch_size is None then we are using single-Arm trials - arms_per_trial=scheduler.options.batch_size or 1, - early_stopping_strategy_cls=( - None - if scheduler.options.early_stopping_strategy is None - else scheduler.options.early_stopping_strategy.__class__.__name__ - ), - global_stopping_strategy_cls=( - None - if scheduler.options.global_stopping_strategy is None - else scheduler.options.global_stopping_strategy.__class__.__name__ - ), - transformed_dimensionality=_get_max_transformed_dimensionality( - search_space=scheduler.experiment.search_space, - generation_strategy=scheduler.standard_generation_strategy, - ), - ) - - def flatten(self) -> dict[str, Any]: - """ - Flatten into an appropriate format for logging to a tabular database. - """ - - self_dict = asdict(self) - experiment_created_record_dict = self_dict.pop("experiment_created_record") - generation_strategy_created_record_dict = self_dict.pop( - "generation_strategy_created_record" - ) - - return { - **self_dict, - **experiment_created_record_dict, - **generation_strategy_created_record_dict, - } - - -@dataclass(frozen=True) -class SchedulerCompletedRecord: - """ - Record of the Scheduler completion event. This will have information only available - after the optimization has completed. - """ - - experiment_completed_record: ExperimentCompletedRecord - - best_point_quality: float - model_fit_quality: float | None - model_std_quality: float | None - model_fit_generalization: float | None - model_std_generalization: float | None - - improvement_over_baseline: float - - num_metric_fetch_e_encountered: int - num_trials_bad_due_to_err: int - - @classmethod - def from_scheduler(cls, scheduler: Scheduler) -> SchedulerCompletedRecord: - try: - model_bridge = get_fitted_model_bridge(scheduler) - quality_and_generalizations_dict = ( - get_fit_and_std_quality_and_generalization_dict( - fitted_model_bridge=model_bridge, - ) - ) - except Exception as e: - warn("Encountered exception in computing model fit quality: " + str(e)) - quality_and_generalizations_dict = { - "model_fit_quality": None, - "model_std_quality": None, - "model_fit_generalization": None, - "model_std_generalization": None, - } - - try: - improvement_over_baseline = scheduler.get_improvement_over_baseline() - except Exception as e: - warn( - "Encountered exception in computing improvement over baseline: " - + str(e) - ) - improvement_over_baseline = float("nan") - - return cls( - experiment_completed_record=ExperimentCompletedRecord.from_experiment( - experiment=scheduler.experiment - ), - best_point_quality=float("nan"), # TODO[T147907632] - improvement_over_baseline=improvement_over_baseline, - num_metric_fetch_e_encountered=scheduler._num_metric_fetch_e_encountered, - num_trials_bad_due_to_err=scheduler._num_trials_bad_due_to_err, - **quality_and_generalizations_dict, - ) - - def flatten(self) -> dict[str, Any]: - """ - Flatten into an appropriate format for logging to a tabular database. - """ - - self_dict = asdict(self) - experiment_completed_record_dict = self_dict.pop("experiment_completed_record") - - return { - **self_dict, - **experiment_completed_record_dict, - } diff --git a/ax/telemetry/tests/test_ax_client.py b/ax/telemetry/tests/test_ax_client.py deleted file mode 100644 index ad811d37061..00000000000 --- a/ax/telemetry/tests/test_ax_client.py +++ /dev/null @@ -1,173 +0,0 @@ -#!/usr/bin/env fbpython -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -import logging -from collections.abc import Sequence - -import numpy as np - -from ax.core.types import TParamValue -from ax.service.ax_client import AxClient, ObjectiveProperties -from ax.telemetry.ax_client import AxClientCompletedRecord, AxClientCreatedRecord -from ax.telemetry.experiment import ExperimentCompletedRecord, ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord -from ax.utils.common.testutils import TestCase - - -class TestAxClient(TestCase): - def test_ax_client_created_record_from_ax_client(self) -> None: - ax_client = AxClient() - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - is_test=True, - ) - - record = AxClientCreatedRecord.from_ax_client(ax_client=ax_client) - - expected = AxClientCreatedRecord( - experiment_created_record=ExperimentCreatedRecord.from_experiment( - experiment=ax_client.experiment - ), - generation_strategy_created_record=( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=ax_client.generation_strategy - ) - ), - arms_per_trial=1, - early_stopping_strategy_cls=None, - global_stopping_strategy_cls=None, - transformed_dimensionality=2, - ) - self.assertEqual(record, expected) - - # Test with HSS & MOO. - ax_client = AxClient() - parameters: list[ - dict[str, TParamValue | Sequence[TParamValue] | dict[str, list[str]]] - ] = [ - { - "name": "SearchSpace.optimizer", - "type": "choice", - "values": ["Adam", "SGD", "Adagrad"], - "dependents": None, - "is_ordered": False, - }, - {"name": "SearchSpace.lr", "type": "range", "bounds": [0.001, 0.1]}, - {"name": "SearchSpace.fixed", "type": "fixed", "value": 12.0}, - { - "name": "SearchSpace", - "type": "fixed", - "value": "SearchSpace", - "dependents": { - "SearchSpace": [ - "SearchSpace.optimizer", - "SearchSpace.lr", - "SearchSpace.fixed", - ] - }, - }, - ] - ax_client.create_experiment( - name="hss_experiment", - parameters=parameters, - objectives={ - "branin": ObjectiveProperties(minimize=True), - "b2": ObjectiveProperties(minimize=False), - }, - is_test=True, - ) - record = AxClientCreatedRecord.from_ax_client(ax_client=ax_client) - - expected = AxClientCreatedRecord( - experiment_created_record=ExperimentCreatedRecord.from_experiment( - experiment=ax_client.experiment - ), - generation_strategy_created_record=( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=ax_client.generation_strategy - ) - ), - arms_per_trial=1, - early_stopping_strategy_cls=None, - global_stopping_strategy_cls=None, - transformed_dimensionality=4, - ) - self.assertEqual(record, expected) - self.assertEqual(record.experiment_created_record.hierarchical_tree_height, 2) - - def test_ax_client_completed_record_from_ax_client(self) -> None: - ax_client = AxClient() - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - is_test=True, - ) - - record = AxClientCompletedRecord.from_ax_client(ax_client=ax_client) - - expected = AxClientCompletedRecord( - experiment_completed_record=ExperimentCompletedRecord.from_experiment( - experiment=ax_client.experiment - ), - best_point_quality=float("nan"), - model_fit_quality=float("nan"), - model_std_quality=float("nan"), - model_fit_generalization=float("nan"), - model_std_generalization=float("nan"), - ) - self._compare_axclient_completed_records(record, expected) - - def test_batch_trial_warning(self) -> None: - ax_client = AxClient() - warning_msg = "GenerationStrategy when using BatchTrials is in beta." - with self.assertLogs(AxClient.__module__, logging.WARNING) as logger: - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - is_test=True, - choose_generation_strategy_kwargs={ - "use_batch_trials": True, - }, - ) - self.assertTrue( - any(warning_msg in output for output in logger.output), - logger.output, - ) - - def _compare_axclient_completed_records( - self, record: AxClientCompletedRecord, expected: AxClientCompletedRecord - ) -> None: - self.assertEqual( - record.experiment_completed_record, expected.experiment_completed_record - ) - numeric_fields = [ - "best_point_quality", - "model_fit_quality", - "model_std_quality", - "model_fit_generalization", - "model_std_generalization", - ] - for field in numeric_fields: - rec_field = getattr(record, field) - exp_field = getattr(expected, field) - if np.isnan(rec_field): - self.assertTrue(np.isnan(exp_field)) - else: - self.assertAlmostEqual(rec_field, exp_field) diff --git a/ax/telemetry/tests/test_experiment.py b/ax/telemetry/tests/test_experiment.py deleted file mode 100644 index f67ddfa28d0..00000000000 --- a/ax/telemetry/tests/test_experiment.py +++ /dev/null @@ -1,128 +0,0 @@ -#!/usr/bin/env fbpython -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from ax.core.utils import get_model_times -from ax.modelbridge.registry import Models -from ax.telemetry.experiment import ExperimentCompletedRecord, ExperimentCreatedRecord -from ax.utils.common.testutils import TestCase -from ax.utils.testing.core_stubs import ( - get_branin_experiment, - get_experiment_with_custom_runner_and_metric, -) -from ax.utils.testing.mock import mock_botorch_optimize - - -class TestExperiment(TestCase): - def test_experiment_created_record_from_experiment(self) -> None: - experiment = get_experiment_with_custom_runner_and_metric( - has_outcome_constraint=True - ) - - record = ExperimentCreatedRecord.from_experiment(experiment=experiment) - expected = ExperimentCreatedRecord( - experiment_name="test", - experiment_type=None, - num_continuous_range_parameters=1, - num_int_range_parameters_small=0, - num_int_range_parameters_medium=0, - num_int_range_parameters_large=1, - num_log_scale_range_parameters=0, - num_unordered_choice_parameters_small=1, - num_unordered_choice_parameters_medium=0, - num_unordered_choice_parameters_large=0, - num_fixed_parameters=1, - dimensionality=3, - hierarchical_tree_height=1, - num_parameter_constraints=3, - num_objectives=1, - num_tracking_metrics=1, - num_outcome_constraints=1, - num_map_metrics=0, - metric_cls_to_quantity={"Metric": 2, "CustomTestMetric": 1}, - runner_cls="CustomTestRunner", - ) - self.assertEqual(record, expected) - - def test_experiment_completed_record_from_experiment(self) -> None: - experiment = get_experiment_with_custom_runner_and_metric( - has_outcome_constraint=True, num_trials=1 - ) - record = ExperimentCompletedRecord.from_experiment(experiment=experiment) - - # Calculate these here, may change from run to run - fit_time, gen_time = get_model_times(experiment=experiment) - expected = ExperimentCompletedRecord( - num_initialization_trials=1, - num_bayesopt_trials=0, - num_other_trials=0, - num_completed_trials=1, - num_failed_trials=0, - num_abandoned_trials=0, - num_early_stopped_trials=0, - total_fit_time=int(fit_time), - total_gen_time=int(gen_time), - num_objectives=1, - num_outcome_constraints=1, - num_tracking_metrics=1, - ) - self.assertEqual(record, expected) - - @mock_botorch_optimize - def test_bayesopt_trials_are_trials_containing_bayesopt(self) -> None: - experiment = get_branin_experiment() - sobol = Models.SOBOL(search_space=experiment.search_space) - trial = experiment.new_batch_trial().add_generator_run( - generator_run=sobol.gen(5) - ) - trial.mark_completed(unsafe=True) - - # create a trial that among other things does bayesopt - data = experiment.fetch_data() - botorch = Models.BOTORCH_MODULAR( - experiment=experiment, - data=data, - ) - trial = ( - experiment.new_batch_trial() - .add_generator_run(generator_run=sobol.gen(2)) - .add_generator_run(generator_run=botorch.gen(5)) - ) - trial.add_arm(experiment.arms_by_name["0_0"]) - trial.mark_completed(unsafe=True) - - # create another BO trial but leave it as a candidate - botorch = Models.BOTORCH_MODULAR( - experiment=experiment, - data=data, - ) - trial = experiment.new_batch_trial().add_generator_run( - generator_run=botorch.gen(5) - ) - - record = ExperimentCompletedRecord.from_experiment(experiment=experiment) - self.assertEqual(record.num_initialization_trials, 1) - self.assertEqual(record.num_bayesopt_trials, 1) - self.assertEqual(record.num_other_trials, 0) - - def test_other_trials_are_trials_with_no_models(self) -> None: - experiment = get_branin_experiment() - sobol = Models.SOBOL(search_space=experiment.search_space) - trial = experiment.new_batch_trial().add_generator_run( - generator_run=sobol.gen(5) - ) - trial.mark_completed(unsafe=True) - - # create a trial that has no GRs that used models - trial = experiment.new_batch_trial() - trial.add_arm(experiment.arms_by_name["0_0"]) - trial.mark_completed(unsafe=True) - - record = ExperimentCompletedRecord.from_experiment(experiment=experiment) - self.assertEqual(record.num_initialization_trials, 1) - self.assertEqual(record.num_bayesopt_trials, 0) - self.assertEqual(record.num_other_trials, 1) diff --git a/ax/telemetry/tests/test_generation_strategy.py b/ax/telemetry/tests/test_generation_strategy.py deleted file mode 100644 index 8f6fa83b561..00000000000 --- a/ax/telemetry/tests/test_generation_strategy.py +++ /dev/null @@ -1,49 +0,0 @@ -#!/usr/bin/env fbpython -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from ax.exceptions.core import AxWarning -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord -from ax.utils.common.testutils import TestCase -from ax.utils.testing.modeling_stubs import ( - get_generation_strategy, - sobol_gpei_generation_node_gs, -) - - -class TestGenerationStrategy(TestCase): - def test_generation_strategy_created_record_from_generation_strategy(self) -> None: - gs = get_generation_strategy() - record = GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=gs - ) - expected = GenerationStrategyCreatedRecord( - generation_strategy_name="Sobol+BO_MIXED", - num_requested_initialization_trials=6, - num_requested_bayesopt_trials=-1, - num_requested_other_trials=0, - max_parallelism=3, - ) - self.assertEqual(record, expected) - - def test_generation_strategy_created_record_node_based(self) -> None: - gs = sobol_gpei_generation_node_gs() - with self.assertWarnsRegex( - AxWarning, - "`GenerationStrategyCreatedRecord` does not fully support node-based*", - ): - record = GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=gs - ) - expected = GenerationStrategyCreatedRecord( - generation_strategy_name="Sobol+MBM_Nodes", - num_requested_initialization_trials=None, - num_requested_bayesopt_trials=None, - num_requested_other_trials=None, - max_parallelism=None, - ) - self.assertEqual(record, expected) diff --git a/ax/telemetry/tests/test_optimization.py b/ax/telemetry/tests/test_optimization.py deleted file mode 100644 index fa43ea6ee94..00000000000 --- a/ax/telemetry/tests/test_optimization.py +++ /dev/null @@ -1,292 +0,0 @@ -#!/usr/bin/env fbpython -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from dataclasses import asdict -from datetime import datetime - -from ax.service.ax_client import AxClient, ObjectiveProperties -from ax.service.scheduler import Scheduler, SchedulerOptions -from ax.telemetry.ax_client import AxClientCompletedRecord, AxClientCreatedRecord -from ax.telemetry.common import get_unique_identifier -from ax.telemetry.experiment import ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord -from ax.telemetry.optimization import ( - OptimizationCompletedRecord, - OptimizationCreatedRecord, -) -from ax.telemetry.scheduler import SchedulerCompletedRecord, SchedulerCreatedRecord -from ax.utils.common.testutils import TestCase -from ax.utils.testing.core_stubs import get_branin_experiment -from ax.utils.testing.modeling_stubs import get_generation_strategy - - -class TestOptimization(TestCase): - def test_optimization_created_record_from_scheduler(self) -> None: - scheduler = Scheduler( - experiment=get_branin_experiment(), - generation_strategy=get_generation_strategy(), - options=SchedulerOptions( - total_trials=0, - tolerated_trial_failure_rate=0.2, - init_seconds_between_polls=10, - ), - ) - - record = OptimizationCreatedRecord.from_scheduler( - scheduler=scheduler, - unique_identifier="foo", - owner="bar", - product_surface="Axolotl", - launch_surface="web", - deployed_job_id=1118, - trial_evaluation_identifier="train", - is_manual_generation_strategy=True, - warm_started_from=None, - num_custom_trials=0, - support_tier="good!", - ) - - expected_dict = { - **SchedulerCreatedRecord.from_scheduler(scheduler=scheduler).flatten(), - "unique_identifier": "foo", - "owner": "bar", - "product_surface": "Axolotl", - "launch_surface": "web", - "deployed_job_id": 1118, - "trial_evaluation_identifier": "train", - "is_manual_generation_strategy": True, - "warm_started_from": None, - "num_custom_trials": 0, - "support_tier": "good!", - } - - self.assertEqual(asdict(record), expected_dict) - - def test_optimization_created_record_from_ax_client(self) -> None: - ax_client = AxClient() - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - is_test=True, - ) - - record = OptimizationCreatedRecord.from_ax_client( - ax_client=ax_client, - unique_identifier="foo", - owner="bar", - product_surface="Axolotl", - launch_surface="web", - deployed_job_id=1118, - trial_evaluation_identifier="train", - is_manual_generation_strategy=True, - warm_started_from=None, - num_custom_trials=0, - ) - - expected_dict = { - **AxClientCreatedRecord.from_ax_client(ax_client=ax_client).flatten(), - "unique_identifier": "foo", - "owner": "bar", - "product_surface": "Axolotl", - "launch_surface": "web", - "deployed_job_id": 1118, - "trial_evaluation_identifier": "train", - "is_manual_generation_strategy": True, - "warm_started_from": None, - "num_custom_trials": 0, - # Extra fields - "scheduler_max_pending_trials": -1, - "scheduler_total_trials": None, - "support_tier": "", - } - - self.assertEqual(asdict(record), expected_dict) - - def test_optimization_created_record_from_experiment_without_generation_strategy( - self, - ) -> None: - experiment = get_branin_experiment() - experiment.experiment_type = "NOTEBOOK" - time_created_str = "2023-04-07T16:01:23.45" - experiment._time_created = datetime.strptime( - time_created_str, "%Y-%m-%dT%H:%M:%S.%f" - ) - - record = OptimizationCreatedRecord.from_experiment( - experiment=experiment, - generation_strategy=None, - unique_identifier=get_unique_identifier(experiment=experiment), - owner="bar", - product_surface="foo", - launch_surface="web", - deployed_job_id=1118, - trial_evaluation_identifier="train", - is_manual_generation_strategy=True, - warm_started_from=None, - num_custom_trials=0, - ) - - expected_dict = { - **asdict(ExperimentCreatedRecord.from_experiment(experiment=experiment)), - "unique_identifier": "branin_test_experiment_2023-04-07 16:01:23", - "owner": "bar", - "product_surface": "foo", - "launch_surface": "web", - "deployed_job_id": 1118, - "trial_evaluation_identifier": "train", - "is_manual_generation_strategy": True, - "warm_started_from": None, - "num_custom_trials": 0, - "arms_per_trial": -1, - "early_stopping_strategy_cls": None, - "global_stopping_strategy_cls": None, - "scheduler_max_pending_trials": -1, - "scheduler_total_trials": None, - "support_tier": "", - "transformed_dimensionality": None, - "num_requested_bayesopt_trials": None, - "num_requested_initialization_trials": None, - "num_requested_other_trials": None, - "max_parallelism": None, - "generation_strategy_name": None, - } - self.maxDiff = None - self.assertEqual(asdict(record), expected_dict) - - def test_optimization_created_record_from_experiment_with_generation_strategy( - self, - ) -> None: - ax_client = AxClient() - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - experiment_type="NOTEBOOK", - is_test=True, - ) - experiment = ax_client.experiment - generation_strategy = ax_client.generation_strategy - time_created_str = "2023-04-07T16:01:23.45" - experiment._time_created = datetime.strptime( - time_created_str, "%Y-%m-%dT%H:%M:%S.%f" - ) - - record = OptimizationCreatedRecord.from_experiment( - experiment=experiment, - generation_strategy=generation_strategy, - unique_identifier=get_unique_identifier(experiment=experiment), - owner="bar", - product_surface="foo", - launch_surface="web", - deployed_job_id=1118, - trial_evaluation_identifier="train", - is_manual_generation_strategy=True, - warm_started_from=None, - num_custom_trials=0, - ) - - expected_dict = { - **asdict(ExperimentCreatedRecord.from_experiment(experiment=experiment)), - **asdict( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=generation_strategy - ) - ), - "unique_identifier": "test_experiment_2023-04-07 16:01:23", - "owner": "bar", - "product_surface": "foo", - "launch_surface": "web", - "deployed_job_id": 1118, - "trial_evaluation_identifier": "train", - "is_manual_generation_strategy": True, - "warm_started_from": None, - "num_custom_trials": 0, - "arms_per_trial": -1, - "early_stopping_strategy_cls": None, - "global_stopping_strategy_cls": None, - "scheduler_max_pending_trials": -1, - "scheduler_total_trials": None, - "support_tier": "", - "transformed_dimensionality": 2, - } - - self.maxDiff = None - self.assertEqual(asdict(record), expected_dict) - - def test_optimization_completed_record_from_scheduler(self) -> None: - scheduler = Scheduler( - experiment=get_branin_experiment(), - generation_strategy=get_generation_strategy(), - options=SchedulerOptions( - total_trials=0, - tolerated_trial_failure_rate=0.2, - init_seconds_between_polls=10, - ), - ) - - record = OptimizationCompletedRecord.from_scheduler( - scheduler=scheduler, - unique_identifier="foo", - deployed_job_id=1118, - estimated_early_stopping_savings=19, - estimated_global_stopping_savings=98, - ) - - expected_dict = { - **SchedulerCompletedRecord.from_scheduler(scheduler=scheduler).flatten(), - "unique_identifier": "foo", - "deployed_job_id": 1118, - "estimated_early_stopping_savings": 19, - "estimated_global_stopping_savings": 98, - } - - self.assertDictsAlmostEqual( - asdict(record), expected_dict, consider_nans_equal=True - ) - - def test_optimization_completed_record_from_ax_client(self) -> None: - ax_client = AxClient() - ax_client.create_experiment( - name="test_experiment", - parameters=[ - {"name": "x", "type": "range", "bounds": [-5.0, 10.0]}, - {"name": "y", "type": "range", "bounds": [0.0, 15.0]}, - ], - objectives={"branin": ObjectiveProperties(minimize=True)}, - is_test=True, - ) - - record = OptimizationCompletedRecord.from_ax_client( - ax_client=ax_client, - unique_identifier="foo", - deployed_job_id=1118, - estimated_early_stopping_savings=19, - estimated_global_stopping_savings=98, - ) - expected_dict = { - **AxClientCompletedRecord.from_ax_client(ax_client=ax_client).flatten(), - "unique_identifier": "foo", - "deployed_job_id": 1118, - "estimated_early_stopping_savings": 19, - "estimated_global_stopping_savings": 98, - # Extra fields - "improvement_over_baseline": float("nan"), - "num_metric_fetch_e_encountered": -1, - "num_trials_bad_due_to_err": -1, - } - - self.assertDictsAlmostEqual( - asdict(record), expected_dict, consider_nans_equal=True - ) diff --git a/ax/telemetry/tests/test_scheduler.py b/ax/telemetry/tests/test_scheduler.py deleted file mode 100644 index 49e7865179a..00000000000 --- a/ax/telemetry/tests/test_scheduler.py +++ /dev/null @@ -1,280 +0,0 @@ -#!/usr/bin/env fbpython -# Copyright (c) Meta Platforms, Inc. and affiliates. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -# pyre-strict - -from typing import cast -from unittest import mock - -import numpy as np - -from ax.core.experiment import Experiment -from ax.core.objective import Objective -from ax.core.optimization_config import OptimizationConfig -from ax.metrics.branin import BraninMetric -from ax.modelbridge.cross_validation import compute_model_fit_metrics_from_modelbridge -from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy -from ax.modelbridge.registry import Models -from ax.runners.synthetic import SyntheticRunner -from ax.service.scheduler import get_fitted_model_bridge, Scheduler, SchedulerOptions -from ax.telemetry.experiment import ExperimentCompletedRecord, ExperimentCreatedRecord -from ax.telemetry.generation_strategy import GenerationStrategyCreatedRecord -from ax.telemetry.scheduler import SchedulerCompletedRecord, SchedulerCreatedRecord -from ax.utils.common.constants import Keys -from ax.utils.common.testutils import TestCase -from ax.utils.testing.core_stubs import get_branin_experiment, get_branin_search_space -from ax.utils.testing.modeling_stubs import get_generation_strategy - - -NUM_SOBOL = 5 - - -class TestScheduler(TestCase): - def test_scheduler_created_record_from_scheduler(self) -> None: - scheduler = Scheduler( - experiment=get_branin_experiment(), - generation_strategy=get_generation_strategy(), - options=SchedulerOptions( - total_trials=0, - tolerated_trial_failure_rate=0.2, - init_seconds_between_polls=10, - ), - ) - - record = SchedulerCreatedRecord.from_scheduler(scheduler=scheduler) - - expected = SchedulerCreatedRecord( - experiment_created_record=ExperimentCreatedRecord.from_experiment( - experiment=scheduler.experiment - ), - generation_strategy_created_record=( - GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=scheduler.standard_generation_strategy - ) - ), - scheduler_total_trials=0, - scheduler_max_pending_trials=10, - arms_per_trial=1, - early_stopping_strategy_cls=None, - global_stopping_strategy_cls=None, - transformed_dimensionality=2, - ) - self.assertEqual(record, expected) - - flat = record.flatten() - expected_dict = { - **ExperimentCreatedRecord.from_experiment( - experiment=scheduler.experiment - ).__dict__, - **GenerationStrategyCreatedRecord.from_generation_strategy( - generation_strategy=scheduler.standard_generation_strategy - ).__dict__, - "scheduler_total_trials": 0, - "scheduler_max_pending_trials": 10, - "arms_per_trial": 1, - "early_stopping_strategy_cls": None, - "global_stopping_strategy_cls": None, - "transformed_dimensionality": 2, - } - self.assertEqual(flat, expected_dict) - - def test_scheduler_completed_record_from_scheduler(self) -> None: - scheduler = Scheduler( - experiment=get_branin_experiment(), - generation_strategy=get_generation_strategy(), - options=SchedulerOptions( - total_trials=0, - tolerated_trial_failure_rate=0.2, - init_seconds_between_polls=10, - ), - ) - - with mock.patch.object( - scheduler, "get_improvement_over_baseline", return_value=5.0 - ): - record = SchedulerCompletedRecord.from_scheduler(scheduler=scheduler) - expected = SchedulerCompletedRecord( - experiment_completed_record=ExperimentCompletedRecord.from_experiment( - experiment=scheduler.experiment - ), - best_point_quality=float("nan"), - model_fit_quality=None, # nan because no model has been fit - model_std_quality=None, - model_fit_generalization=None, - model_std_generalization=None, - improvement_over_baseline=5.0, - num_metric_fetch_e_encountered=0, - num_trials_bad_due_to_err=0, - ) - self._compare_scheduler_completed_records(record, expected) - - flat = record.flatten() - expected_dict = { - **ExperimentCompletedRecord.from_experiment( - experiment=scheduler.experiment - ).__dict__, - "best_point_quality": float("nan"), - "model_fit_quality": None, - "model_std_quality": None, - "model_fit_generalization": None, - "model_std_generalization": None, - "improvement_over_baseline": 5.0, - "num_metric_fetch_e_encountered": 0, - "num_trials_bad_due_to_err": 0, - } - self.assertDictsAlmostEqual(flat, expected_dict, consider_nans_equal=True) - - def test_scheduler_raise_exceptions(self) -> None: - scheduler = Scheduler( - experiment=get_branin_experiment(), - generation_strategy=get_generation_strategy(), - options=SchedulerOptions( - total_trials=0, - tolerated_trial_failure_rate=0.2, - init_seconds_between_polls=10, - ), - ) - - with mock.patch.object( - scheduler, - "get_improvement_over_baseline", - side_effect=Exception("test_exception"), - ): - record = SchedulerCompletedRecord.from_scheduler(scheduler=scheduler) - flat = record.flatten() - self.assertTrue(np.isnan(flat["improvement_over_baseline"])) - - def test_scheduler_model_fit_metrics_logging(self) -> None: - # set up for model fit metrics - branin_experiment = Experiment( - name="branin_test_experiment", - search_space=get_branin_search_space(), - runner=SyntheticRunner(), - optimization_config=OptimizationConfig( - objective=Objective( - metric=BraninMetric(name="branin", param_names=["x1", "x2"]), - minimize=True, - ), - ), - is_test=True, - ) - branin_experiment._properties[Keys.IMMUTABLE_SEARCH_SPACE_AND_OPT_CONF] = True - generation_strategy = GenerationStrategy( - steps=[ - GenerationStep( - model=Models.SOBOL, num_trials=NUM_SOBOL, max_parallelism=NUM_SOBOL - ), - GenerationStep(model=Models.BOTORCH_MODULAR, num_trials=-1), - ] - ) - - # starting proper tests - scheduler = Scheduler( - experiment=branin_experiment, - generation_strategy=generation_strategy, - options=SchedulerOptions(), - ) - # Trying to attain a record without any trials yields an error in ModelFitRecord - # and a warning in SchedulerCompletedRecord. - - scheduler.run_n_trials(max_trials=NUM_SOBOL + 1) - - # end-to-end test with Scheduler - record = SchedulerCompletedRecord.from_scheduler(scheduler=scheduler) - model_bridge = get_fitted_model_bridge(scheduler) - - fit_metrics = compute_model_fit_metrics_from_modelbridge( - model_bridge=model_bridge, - generalization=False, - untransform=False, - ) - # checking fit metrics - r2 = fit_metrics.get("coefficient_of_determination") - self.assertIsInstance(r2, dict) - r2 = cast(dict[str, float], r2) - self.assertTrue("branin" in r2) - r2_branin = r2["branin"] - self.assertIsInstance(r2_branin, float) - - std = fit_metrics.get("std_of_the_standardized_error") - self.assertIsInstance(std, dict) - std = cast(dict[str, float], std) - self.assertTrue("branin" in std) - std_branin = std["branin"] - self.assertIsInstance(std_branin, float) - - model_std_quality = 1 / std_branin - - # check generalization metrics - gen_metrics = compute_model_fit_metrics_from_modelbridge( - model_bridge=model_bridge, - generalization=True, - untransform=False, - ) - r2_gen = gen_metrics.get("coefficient_of_determination") - r2_gen = cast(dict[str, float], r2_gen) - r2_gen_branin = r2_gen["branin"] - gen_std = gen_metrics.get("std_of_the_standardized_error") - gen_std = cast(dict[str, float], gen_std) - gen_std_branin = gen_std["branin"] - model_std_generalization = 1 / gen_std_branin - - expected = SchedulerCompletedRecord( - experiment_completed_record=ExperimentCompletedRecord.from_experiment( - experiment=scheduler.experiment - ), - best_point_quality=float("nan"), - model_fit_quality=r2_branin, - model_std_quality=model_std_quality, - model_fit_generalization=r2_gen_branin, - model_std_generalization=model_std_generalization, - improvement_over_baseline=float("nan"), - num_metric_fetch_e_encountered=0, - num_trials_bad_due_to_err=0, - ) - self._compare_scheduler_completed_records(record, expected) - - flat = record.flatten() - expected_dict = { - **ExperimentCompletedRecord.from_experiment( - experiment=scheduler.experiment - ).__dict__, - "best_point_quality": float("nan"), - "model_fit_quality": r2_branin, - "model_std_quality": model_std_quality, - "model_fit_generalization": r2_gen_branin, - "model_std_generalization": model_std_generalization, - "improvement_over_baseline": float("nan"), - "num_metric_fetch_e_encountered": 0, - "num_trials_bad_due_to_err": 0, - } - self.assertDictsAlmostEqual(flat, expected_dict, consider_nans_equal=True) - - def _compare_scheduler_completed_records( - self, record: SchedulerCompletedRecord, expected: SchedulerCompletedRecord - ) -> None: - self.assertEqual( - record.experiment_completed_record, expected.experiment_completed_record - ) - numeric_fields = [ - "best_point_quality", - "model_fit_quality", - "model_std_quality", - "model_fit_generalization", - "model_std_generalization", - "improvement_over_baseline", - "num_metric_fetch_e_encountered", - "num_trials_bad_due_to_err", - ] - for field in numeric_fields: - rec_field = getattr(record, field) - exp_field = getattr(expected, field) - if rec_field is None: - self.assertIsNone(exp_field, msg=field) - elif np.isnan(rec_field): - self.assertTrue(np.isnan(exp_field), msg=field) - else: - self.assertAlmostEqual(rec_field, exp_field, msg=field) diff --git a/sphinx/source/telemetry.rst b/sphinx/source/telemetry.rst deleted file mode 100644 index a43be7b1cbf..00000000000 --- a/sphinx/source/telemetry.rst +++ /dev/null @@ -1,56 +0,0 @@ -.. role:: hidden - :class: hidden-section - -ax.telemetry -============ - -.. automodule:: ax.telemetry -.. currentmodule:: ax.telemetry - -AxClient -~~~~~~~~ - -.. automodule:: ax.telemetry.ax_client - :members: - :undoc-members: - :show-inheritance: - -Common -~~~~~~ - -.. automodule:: ax.telemetry.common - :members: - :undoc-members: - :show-inheritance: - -Experiment -~~~~~~~~~~ - -.. automodule:: ax.telemetry.experiment - :members: - :undoc-members: - :show-inheritance: - -Generation Strategy -~~~~~~~~~~~~~~~~~~~ - -.. automodule:: ax.telemetry.generation_strategy - :members: - :undoc-members: - :show-inheritance: - -Optimization -~~~~~~~~~~~~ - -.. automodule:: ax.telemetry.optimization - :members: - :undoc-members: - :show-inheritance: - -Scheduler -~~~~~~~~~ - -.. automodule:: ax.telemetry.scheduler - :members: - :undoc-members: - :show-inheritance: