model's bug fix

Author: calderjo

Dockerfile.tmpl

@@ -55,6 +55,9 @@ RUN uv pip install --system --force-reinstall -i https://pypi.anaconda.org/intel
 # newer daal4py requires tbb>=2022, but libpysal is downgrading it for some reason
 RUN uv pip install --system "tbb>=2022" "libpysal==4.9.2"
 
+# b/404590350: Ray and torchtune have conflicting tune cli, we will prioritize torchtune.
+RUN uv pip install --system --force-reinstall --no-deps torchtune
+
 # Adding non-package dependencies:
 
 ADD clean-layer.sh  /tmp/clean-layer.sh

kaggle_requirements.txt

@@ -130,9 +130,11 @@ tensorflow-io
 tensorflow-text
 tensorflow_decision_forests
 timm
+torchao
 torchinfo
 torchmetrics
 torchtune
+triton
 tsfresh
 vtk
 wandb

tests/test_torchtune.py

@@ -1,9 +1,10 @@
 import unittest
-
 import subprocess
 
 class TestTorchtune(unittest.TestCase):
     def test_help(self):
-        ret_code = subprocess.run(["tune", "--help"])
-        self.assertEqual(0, ret_code.returncode)
-        self.assertIsNone(ret_code.stderr)
+        result = subprocess.run(["tune", "--help"], stdout=subprocess.PIPE)
+
+        self.assertEqual(0, result.returncode)
+        self.assertIsNone(result.stderr)
+        self.assertIn("Download a model from the Hugging Face Hub or Kaggle Model Hub.", result.stdout.decode("utf-8"))