fix numpy mkl test

Author: djherbis

Dockerfile.tmpl

@@ -32,6 +32,9 @@ RUN uv pip uninstall --system google-cloud-bigquery-storage
 # to avoid affecting the larger build, we'll post-install it.
 RUN uv pip install --no-build-isolation --system "git+https://github.com/Kaggle/learntools"
 
+# DO_NOT_MERGE: Wait for Colab to fix mkl in image.
+RUN uv pip install --system --force-reinstall -i https://pypi.anaconda.org/intel/simple numpy
+
 # b/328788268 We install an incompatible pair of libs (shapely<2, libpysal==4.9.2) so we can't put this one in the requirements.txt
 RUN uv pip install --system  "libpysal==4.9.2"
 

tests/test_numpy.py

@@ -3,9 +3,10 @@
 from distutils.version import StrictVersion
 
 import numpy as np
-from numpy.distutils.system_info import get_info
+import io
+from contextlib import redirect_stdout
 
-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")) 
@@ -18,5 +19,13 @@ def test_array(self):
     # Numpy must be linked to the MKL. (Occasionally, a third-party package will muck up the installation
     # and numpy will be reinstalled with an OpenBLAS backing.)
     def test_mkl(self):
-        # This will throw an exception if the MKL is not linked correctly or return an empty dict.
-        self.assertTrue(get_info("blas_mkl"))
+        try:
+            from numpy.distutils.system_info import get_info
+            # This will throw an exception if the MKL is not linked correctly or return an empty dict.
+            self.assertTrue(get_info("blas_mkl"))
+        except:
+            # Fallback to check if mkl is present via show_config()
+            config_out = io.StringIO()
+            with redirect_stdout(config_out):
+                np.show_config()
+            self.assertIn("mkl_rt", config_out.getvalue())