Filters JAX packages from requirements.txt

Author: JyotinderSingh

keras_remote/infra/container_builder.py

@@ -2,6 +2,7 @@
 
 import hashlib
 import os
+import re
 import shutil
 import string
 import tarfile
@@ -28,6 +29,58 @@
 )
 _RUNNER_DIR = os.path.join(os.path.dirname(__file__), os.pardir, "runner")
 
+# JAX-related packages managed by the Dockerfile template.
+# User requirements containing these are filtered out to prevent overriding
+# the accelerator-specific JAX installation (e.g., jax[tpu], jax[cuda12]).
+_JAX_PACKAGE_NAMES = frozenset({"jax", "jaxlib", "libtpu", "libtpu-nightly"})
+_PACKAGE_NAME_RE = re.compile(r"^([a-zA-Z0-9]([a-zA-Z0-9._-]*[a-zA-Z0-9])?)")
+_KEEP_MARKER = "# kr:keep"
+
+
+def _filter_jax_requirements(requirements_content):
+  """Remove JAX-related packages from requirements content.
+
+  Strips lines that would override the accelerator-specific JAX installation
+  managed by the Dockerfile template. Logs a warning for each filtered line.
+
+  To preserve a JAX line, append ``# kr:keep`` to it in requirements.txt.
+
+  Args:
+      requirements_content: Raw text of a requirements.txt file.
+
+  Returns:
+      Filtered requirements text with JAX-related lines removed.
+  """
+  filtered_lines = []
+  for line in requirements_content.splitlines(keepends=True):
+    stripped = line.strip()
+    # Preserve blanks, comments, and pip flags (-e, --index-url, etc.)
+    if not stripped or stripped.startswith("#") or stripped.startswith("-"):
+      filtered_lines.append(line)
+      continue
+
+    # Allow users to bypass the filter with an inline marker.
+    if _KEEP_MARKER in line:
+      filtered_lines.append(line)
+      continue
+
+    m = _PACKAGE_NAME_RE.match(stripped)
+    if m:
+      # PEP 503 normalization: lowercase, collapse [-_.] to '-'
+      normalized = re.sub(r"[-_.]+", "-", m.group(1)).lower()
+      if normalized in _JAX_PACKAGE_NAMES:
+        logging.warning(
+          "Filtered '%s' from requirements — JAX is installed "
+          "automatically with the correct accelerator backend. "
+          "To override, add '# kr:keep' to the line.",
+          stripped,
+        )
+        continue
+
+    filtered_lines.append(line)
+
+  return "".join(filtered_lines)
+
 
 def get_or_build_container(
   base_image,
@@ -108,7 +161,7 @@ def _hash_requirements(requirements_path, category, base_image):
 
   if requirements_path and os.path.exists(requirements_path):
     with open(requirements_path, "r") as f:
-      content += f.read()
+      content += _filter_jax_requirements(f.read())
 
   # Include remote_runner.py in the hash so container rebuilds when it changes
   remote_runner_path = os.path.join(_RUNNER_DIR, REMOTE_RUNNER_FILE_NAME)
@@ -195,9 +248,12 @@ def _build_and_push(
     with open(dockerfile_path, "w") as f:
       f.write(dockerfile_content)
 
-    # Copy requirements.txt if it exists
+    # Copy requirements.txt (with JAX-related packages filtered out)
     if requirements_path and os.path.exists(requirements_path):
-      shutil.copy(requirements_path, os.path.join(tmpdir, "requirements.txt"))
+      with open(requirements_path, "r") as f:
+        filtered = _filter_jax_requirements(f.read())
+      with open(os.path.join(tmpdir, "requirements.txt"), "w") as f:
+        f.write(filtered)
 
     # Copy remote_runner.py
     remote_runner_src = os.path.join(_RUNNER_DIR, REMOTE_RUNNER_FILE_NAME)

keras_remote/infra/container_builder_test.py

@@ -9,6 +9,7 @@
 from google.api_core import exceptions as google_exceptions
 
 from keras_remote.infra.container_builder import (
+  _filter_jax_requirements,
   _generate_dockerfile,
   _hash_requirements,
   _image_exists,
@@ -23,6 +24,64 @@ def _make_temp_path(test_case):
   return pathlib.Path(td.name)
 
 
+class TestFilterJaxRequirements(parameterized.TestCase):
+  @parameterized.named_parameters(
+    dict(testcase_name="bare_jax", line="jax\n"),
+    dict(testcase_name="jax_with_tpu_extras", line="jax[tpu]>=0.4.6\n"),
+    dict(testcase_name="jax_cuda", line="jax[cuda12]==0.4.30\n"),
+    dict(testcase_name="jax_cpu", line="jax[cpu]\n"),
+    dict(testcase_name="jaxlib", line="jaxlib>=0.4.6\n"),
+    dict(testcase_name="libtpu", line="libtpu\n"),
+    dict(testcase_name="libtpu_nightly_hyphen", line="libtpu-nightly\n"),
+    dict(testcase_name="libtpu_nightly_underscore", line="libtpu_nightly\n"),
+    dict(testcase_name="jax_uppercase", line="JAX\n"),
+    dict(testcase_name="jax_mixed_case", line="Jax[tpu]\n"),
+  )
+  def test_filters_jax_packages(self, line):
+    self.assertEqual(_filter_jax_requirements(line), "")
+
+  @parameterized.named_parameters(
+    dict(testcase_name="numpy", line="numpy==1.26\n"),
+    dict(testcase_name="keras", line="keras\n"),
+    dict(testcase_name="scipy", line="scipy>=1.12\n"),
+    dict(testcase_name="comment", line="# jax should be here\n"),
+    dict(testcase_name="blank", line="\n"),
+    dict(testcase_name="pip_flag", line="-e git+https://foo\n"),
+    dict(testcase_name="index_url", line="--index-url https://pypi.org\n"),
+  )
+  def test_preserves_non_jax_packages(self, line):
+    self.assertEqual(_filter_jax_requirements(line), line)
+
+  @parameterized.named_parameters(
+    dict(testcase_name="jax_keep", line="jax==0.4.30  # kr:keep\n"),
+    dict(testcase_name="jaxlib_keep", line="jaxlib  # kr:keep\n"),
+    dict(testcase_name="libtpu_keep", line="libtpu-nightly  # kr:keep\n"),
+  )
+  def test_kr_keep_overrides_filter(self, line):
+    self.assertEqual(_filter_jax_requirements(line), line)
+
+  def test_mixed_requirements(self):
+    content = (
+      "numpy==1.26\njax[tpu]>=0.4.6\nscipy\n"
+      "jaxlib\nkeras\njax==0.4.30  # kr:keep\n"
+    )
+    result = _filter_jax_requirements(content)
+    self.assertEqual(
+      result, "numpy==1.26\nscipy\nkeras\njax==0.4.30  # kr:keep\n"
+    )
+
+  def test_empty_string(self):
+    self.assertEqual(_filter_jax_requirements(""), "")
+
+  def test_only_jax_packages(self):
+    self.assertEqual(_filter_jax_requirements("jax\njaxlib\nlibtpu\n"), "")
+
+  def test_preserves_comments_and_blanks(self):
+    content = "# ML deps\nnumpy\n\njax\n# end\n"
+    result = _filter_jax_requirements(content)
+    self.assertEqual(result, "# ML deps\nnumpy\n\n# end\n")
+
+
 class TestHashRequirements(parameterized.TestCase):
   def test_deterministic(self):
     tmp_path = _make_temp_path(self)
@@ -81,6 +140,17 @@ def test_returns_hex_string(self):
     h = _hash_requirements(str(req), "gpu", "python:3.12-slim")
     self.assertRegex(h, r"^[0-9a-f]{64}$")
 
+  def test_jax_in_requirements_does_not_affect_hash(self):
+    tmp_path = _make_temp_path(self)
+    req_without_jax = tmp_path / "r1.txt"
+    req_without_jax.write_text("numpy==1.26\n")
+    req_with_jax = tmp_path / "r2.txt"
+    req_with_jax.write_text("numpy==1.26\njax[tpu]>=0.4.6\n")
+
+    h1 = _hash_requirements(str(req_without_jax), "tpu", "python:3.12-slim")
+    h2 = _hash_requirements(str(req_with_jax), "tpu", "python:3.12-slim")
+    self.assertEqual(h1, h2)
+
 
 class TestGenerateDockerfile(parameterized.TestCase):
   @parameterized.named_parameters(