From 4e22769cbb4ee9a6a327d8097d4c2ac68984a35a Mon Sep 17 00:00:00 2001
From: Johnny Chavez <64660690+calderjo@users.noreply.github.com>
Date: Tue, 1 Oct 2024 23:01:52 -0700
Subject: [PATCH] Fix build due to numpy (#1430)

numpy 2.0 and greater is causing issues with handful of packages
---
 Dockerfile.tmpl      | 10 +++++++---
 tests/test_fastai.py |  7 +++++++
 tests/test_numpy.py  |  8 +++++++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index b58d26f2..e4bde450 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -210,6 +210,8 @@ ADD patches/keras_internal.py \
 # b/350573866: xgboost v2.1.0 breaks learntools
 RUN apt-get install -y libfreetype6-dev && \
     apt-get install -y libglib2.0-0 libxext6 libsm6 libxrender1 libfontconfig1 --fix-missing && \
+    rm -rf /opt/conda/lib/python3.10/site-packages/numpy* && \
+    pip install "numpy==1.26.4" && \
     pip install gensim \
         textblob \
         wordcloud \
@@ -425,8 +427,7 @@ RUN pip install bleach \
         webencodings \
         widgetsnbextension \
         # Require pyarrow newer than https://github.com/advisories/GHSA-5wvp-7f3h-6wmm
-        {{ if eq .Accelerator "gpu" }} pyarrow {{ else }} "pyarrow>=14.0.1" {{ end }} \
-        fastai
+        {{ if eq .Accelerator "gpu" }} pyarrow {{ else }} "pyarrow>=14.0.1" {{ end }} 
 
 RUN python -m spacy download en_core_web_sm && python -m spacy download en_core_web_lg && \
     apt-get update && apt-get install -y ffmpeg && \
@@ -470,7 +471,8 @@ RUN pip install wandb \
         Rtree \
         accelerate && \
         apt-get -y install libspatialindex-dev && \
-        rm -rf /opt/conda/lib/python3.10/site-packages/numpy* && \
+    # b/370860329: newer versions are not capable with current tensorflow
+    rm -rf /opt/conda/lib/python3.10/site-packages/numpy* && \
     pip install "numpy==1.26.4" && \
     pip install pytorch-ignite \
         qgrid \
@@ -501,6 +503,8 @@ RUN pip install wandb \
         timm \
         torchinfo && \
     pip install git+https://github.com/facebookresearch/segment-anything.git && \
+    # b/370860329: newer versions are not capable with current tensorflow
+    pip install --no-dependencies fastai fastdownload && \
     # b/343971718: remove duplicate aiohttp installs, and reinstall it
     rm -rf /opt/conda/lib/python3.10/site-packages/aiohttp* && \
     mamba install --force-reinstall -y aiohttp && \
diff --git a/tests/test_fastai.py b/tests/test_fastai.py
index edfd402e..0de1f82f 100644
--- a/tests/test_fastai.py
+++ b/tests/test_fastai.py
@@ -5,6 +5,13 @@
 from fastai.tabular.all import *
 
 class TestFastAI(unittest.TestCase):
+    # Basic import
+    def test_basic(self):
+        import fastai
+        import fastcore
+        import fastprogress
+        import fastdownload
+
     def test_has_version(self):
         self.assertGreater(len(fastai.__version__), 2)
     
diff --git a/tests/test_numpy.py b/tests/test_numpy.py
index 18f74b8c..071c3d30 100644
--- a/tests/test_numpy.py
+++ b/tests/test_numpy.py
@@ -1,9 +1,15 @@
 import unittest
 
+from distutils.version import StrictVersion
+
 import numpy as np
 from numpy.distutils.system_info import get_info
 
-class TestNumpy(unittest.TestCase):    
+class TestNumpy(unittest.TestCase):   
+    def test_version(self):
+        # b/370860329: newer versions are not capable with current tensorflow
+        self.assertEqual(StrictVersion(np.__version__), StrictVersion("1.26.4")) 
+
     def test_array(self):
         array = np.array([1, 3])