Skip to content

Filter stale trials from Scatter, ArmEffects using built in filtering logic #4396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mpolson64 wants to merge 1 commit into from
Closed

Filter stale trials from Scatter, ArmEffects using built in filtering logic #4396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 13 additions & 15 deletions ax/analysis/plotly/arm_effects.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@
MARGIN_REDUCUTION,
MULTIPLE_CANDIDATE_TRIALS_LEGEND,
SINGLE_CANDIDATE_TRIAL_LEGEND,
STALE_FAIL_REASON,
trial_index_to_color,
truncate_label,
X_TICKER_SCALING_FACTOR,
Expand All @@ -44,7 +43,7 @@
from ax.core.arm import Arm
from ax.core.base_trial import sort_by_trial_index_and_arm_name
from ax.core.experiment import Experiment
from ax.core.trial_status import STALE_ABANDONED_CANDIDATE_STATUSES, TrialStatus
from ax.core.trial_status import TrialStatus
from ax.exceptions.core import UserInputError
from ax.generation_strategy.generation_strategy import GenerationStrategy
from plotly import graph_objects as go
Expand Down Expand Up @@ -138,7 +137,16 @@ def __init__(
self.use_model_predictions = use_model_predictions
self.relativize = relativize
self.trial_index = trial_index
self.trial_statuses = trial_statuses

# By default, include all trials except those that are abandoned or stale.
if trial_statuses is not None:
self.trial_statuses: list[TrialStatus] | None = [*trial_statuses]
elif self.trial_index is not None:
self.trial_statuses: list[TrialStatus] | None = None
else:
self.trial_statuses: list[TrialStatus] | None = [
*{*TrialStatus} - {TrialStatus.ABANDONED, TrialStatus.STALE}
]
self.additional_arms = additional_arms
self.labels: Mapping[str, str] = labels or {}
self.show_cumulative_best = show_cumulative_best
Expand Down Expand Up @@ -328,19 +336,9 @@ def _prepare_figure(
candidate_trials = df[df["trial_status"] == TrialStatus.CANDIDATE.name][
"trial_index"
].unique()
# Filter out undesired trials like STALE and ABANDONED trials from plot.
status_filter = ~df["trial_status"].isin(
[ts.name for ts in STALE_ABANDONED_CANDIDATE_STATUSES]
)
# Also filter out failed trials that failed with STALE_FAIL_REASON.
stale_failed_filter = ~(
(df["trial_status"] == TrialStatus.FAILED.name)
& (df["fail_reason"].notna())
& (df["fail_reason"] == STALE_FAIL_REASON)
)
trials = df[status_filter & stale_failed_filter]["trial_index"].unique()

# Check if candidate_trial is NaN and handle it
trials = df["trial_index"].unique()

trial_indices = list(trials)
trial_indices.extend(candidate_trials)

Expand Down
26 changes: 12 additions & 14 deletions ax/analysis/plotly/scatter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
MARGIN_REDUCUTION,
MULTIPLE_CANDIDATE_TRIALS_LEGEND,
SINGLE_CANDIDATE_TRIAL_LEGEND,
STALE_FAIL_REASON,
trial_index_to_color,
truncate_label,
Z_SCORE_95_CI,
Expand All @@ -45,7 +44,7 @@
)
from ax.core.arm import Arm
from ax.core.experiment import Experiment
from ax.core.trial_status import STALE_ABANDONED_CANDIDATE_STATUSES, TrialStatus
from ax.core.trial_status import TrialStatus
from ax.exceptions.core import UserInputError
from ax.generation_strategy.generation_strategy import GenerationStrategy
from ax.utils.common.logger import get_logger
Expand Down Expand Up @@ -123,7 +122,15 @@ def __init__(
self.use_model_predictions = use_model_predictions
self.relativize = relativize
self.trial_index = trial_index
self.trial_statuses = trial_statuses
# By default, include all trials except those that are abandoned or stale.
if trial_statuses is not None:
self.trial_statuses: list[TrialStatus] | None = [*trial_statuses]
elif self.trial_index is not None:
self.trial_statuses: list[TrialStatus] | None = None
else:
self.trial_statuses: list[TrialStatus] | None = [
*{*TrialStatus} - {TrialStatus.ABANDONED, TrialStatus.STALE}
]
self.additional_arms = additional_arms
self.labels: dict[str, str] = {**labels} if labels is not None else {}
self.show_pareto_frontier = show_pareto_frontier
Expand Down Expand Up @@ -334,17 +341,8 @@ def _prepare_figure(
candidate_trials = df[df["trial_status"] == TrialStatus.CANDIDATE.name][
"trial_index"
].unique()
# Filter out undesired trials like STALE and ABANDONED trials from plot.
status_filter = ~df["trial_status"].isin(
[ts.name for ts in STALE_ABANDONED_CANDIDATE_STATUSES]
)
# Also filter out failed trials that failed with STALE_FAIL_REASON.
stale_failed_filter = ~(
(df["trial_status"] == TrialStatus.FAILED.name)
& (df["fail_reason"].notna())
& (df["fail_reason"] == STALE_FAIL_REASON)
)
trials = df[status_filter & stale_failed_filter]["trial_index"].unique()

trials = df["trial_index"].unique()

trials_list = trials.tolist()
trial_indices = trials_list.copy()
Expand Down