From 272b99dec9440592ce0ead393a1f15a6a0ce2f8d Mon Sep 17 00:00:00 2001
From: Bernie Beckerman <bbeckerman@meta.com>
Date: Wed, 7 Aug 2024 11:39:26 -0700
Subject: [PATCH] pin numpy to < 2.0 until tensorboard cuts a release

Summary: The latest TB version hasn't incorporated numpy2.0-compatible changes (https://github.com/tensorflow/tensorboard/issues/6869#issuecomment-2273718763), and recently unpinned botorch (D60735276) means Ax CI is now failing. This pins numpy to < 2.0 to prevent issues in CI.

Differential Revision: D60917879
---
 ax/metrics/tensorboard.py | 11 +++++++++++
 setup.py                  |  3 +++
 2 files changed, 14 insertions(+)

diff --git a/ax/metrics/tensorboard.py b/ax/metrics/tensorboard.py
index 63d2b20e538..784de6e8b0c 100644
--- a/ax/metrics/tensorboard.py
+++ b/ax/metrics/tensorboard.py
@@ -228,3 +228,14 @@ def _get_event_multiplexer_for_trial(
         "TensorboardMetric, please install tensorboard."
     )
     pass
+except AttributeError as e:
+    if (
+        "`np.string_` was removed in the NumPy 2.0 release. Use `np.bytes_` instead."
+        in repr(e)
+    ):
+        logger.warning(
+            "tensorboard package not compatible with numpy 2.0+. "
+            "Please install numpy < 2.0."
+        )
+    else:
+        raise e
diff --git a/setup.py b/setup.py
index 0b236fe2299..db0a1ceaa93 100644
--- a/setup.py
+++ b/setup.py
@@ -53,6 +53,9 @@
 NOTEBOOK_REQUIRES = ["jupyter"]
 
 UNITTEST_MINIMAL_REQUIRES = [
+    # Replace with `tensorboard >= x.x` once tb cuts a release.
+    # https://github.com/tensorflow/tensorboard/issues/6869#issuecomment-2273718763
+    "numpy<2.0",
     "tensorboard",  # For tensorboard unit tests.
     "torchvision",  # For torchvision unit tests.
     "torchx",  # For torchx unit tests.