Check for Nans and Infs in TensorboardMetric

Summary: Raises a ValueError in bulk_fetch_trial_data if a Nan or an Inf is found. This will get wrapped up in a MetricFetchE and handled appropriately in the Scheduler (ex. INFO if we intend to try and fetch again, WARN if coming from a tracking metric, mark trial as ABANDONED if the metric is needed for the optimization https://fburl.com/code/eq37gghi).

Differential Revision: D60670356

ax/metrics/tensorboard.py

@@ -13,6 +13,8 @@
 from logging import Logger
 from typing import Any, Dict, List, Optional
 
+import numpy as np
+
 import pandas as pd
 from ax.core.base_trial import BaseTrial
 from ax.core.map_data import MapData, MapKeyInfo
@@ -166,6 +168,12 @@ def bulk_fetch_trial_data(
                         .reset_index()
                     )
 
+                    # If there are any NaNs or Infs in the data, raise an Exception
+                    if np.any(~np.isfinite(df["mean"])):
+                        raise ValueError(
+                            f"Found NaNs or Infs in data for {metric.name}."
+                        )
+
                     # Apply per-metric post-processing
                     # Apply cumulative "best" (min if lower_is_better)
                     if metric.cumulative_best: