reduce redundancy

Author: JyotinderSingh

keras_remote/infra/container_builder.py

@@ -109,9 +109,15 @@ def get_or_build_container(
   cluster_name = cluster_name or get_default_cluster_name()
   category = accelerators.get_category(accelerator_type)
 
+  # Read and filter requirements once, reuse for hashing and building.
+  filtered_requirements = None
+  if requirements_path and os.path.exists(requirements_path):
+    with open(requirements_path, "r") as f:
+      filtered_requirements = _filter_jax_requirements(f.read())
+
   # Generate deterministic hash from requirements + base image + category
   requirements_hash = _hash_requirements(
-    requirements_path, category, base_image
+    filtered_requirements, category, base_image
   )
 
   # Use category for image name (e.g., 'tpu-hash', 'gpu-hash')
@@ -137,7 +143,7 @@ def get_or_build_container(
   logging.info("Building new container (requirements changed): %s", image_uri)
   return _build_and_push(
     base_image,
-    requirements_path,
+    filtered_requirements,
     category,
     project,
     image_uri,
@@ -147,12 +153,12 @@ def get_or_build_container(
 
 
 def _hash_requirements(
-  requirements_path: str | None, category: str, base_image: str
+  filtered_requirements: str | None, category: str, base_image: str
 ) -> str:
   """Create deterministic hash from requirements + category + remote_runner + base image.
 
   Args:
-      requirements_path: Path to requirements.txt (or None)
+      filtered_requirements: Pre-filtered requirements content (or None)
       category: Accelerator category ('cpu', 'gpu', 'tpu')
       base_image: Base Docker image (e.g., 'python:3.12-slim')
 
@@ -161,9 +167,8 @@ def _hash_requirements(
   """
   content = f"base_image={base_image}\ncategory={category}\n"
 
-  if requirements_path and os.path.exists(requirements_path):
-    with open(requirements_path, "r") as f:
-      content += _filter_jax_requirements(f.read())
+  if filtered_requirements:
+    content += filtered_requirements
 
   # 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)
@@ -218,7 +223,7 @@ def _image_exists(image_uri: str, project: str) -> bool:
 
 def _build_and_push(
   base_image: str,
-  requirements_path: str | None,
+  filtered_requirements: str | None,
   category: str,
   project: str,
   image_uri: str,
@@ -229,7 +234,7 @@ def _build_and_push(
 
   Args:
       base_image: Base Docker image
-      requirements_path: Path to requirements.txt (or None)
+      filtered_requirements: Pre-filtered requirements content (or None)
       category: Accelerator category ('cpu', 'gpu', 'tpu')
       project: GCP project ID
       image_uri: Target image URI
@@ -242,20 +247,18 @@ def _build_and_push(
     # Generate Dockerfile
     dockerfile_content = _generate_dockerfile(
       base_image=base_image,
-      requirements_path=requirements_path,
+      has_requirements=filtered_requirements is not None,
       category=category,
     )
 
     dockerfile_path = os.path.join(tmpdir, "Dockerfile")
     with open(dockerfile_path, "w") as f:
       f.write(dockerfile_content)
 
-    # Copy requirements.txt (with JAX-related packages filtered out)
-    if requirements_path and os.path.exists(requirements_path):
-      with open(requirements_path, "r") as f:
-        filtered = _filter_jax_requirements(f.read())
+    # Write pre-filtered requirements.txt
+    if filtered_requirements is not None:
       with open(os.path.join(tmpdir, "requirements.txt"), "w") as f:
-        f.write(filtered)
+        f.write(filtered_requirements)
 
     # Copy remote_runner.py
     remote_runner_src = os.path.join(_RUNNER_DIR, REMOTE_RUNNER_FILE_NAME)
@@ -267,7 +270,7 @@ def _build_and_push(
     with tarfile.open(tarball_path, "w:gz") as tar:
       tar.add(dockerfile_path, arcname="Dockerfile")
       tar.add(remote_runner_dst, arcname=REMOTE_RUNNER_FILE_NAME)
-      if requirements_path and os.path.exists(requirements_path):
+      if filtered_requirements is not None:
         tar.add(
           os.path.join(tmpdir, "requirements.txt"), arcname="requirements.txt"
         )
@@ -330,13 +333,13 @@ def _build_and_push(
 
 
 def _generate_dockerfile(
-  base_image: str, requirements_path: str | None, category: str
+  base_image: str, has_requirements: bool, category: str
 ) -> str:
   """Generate Dockerfile content based on configuration.
 
   Args:
       base_image: Base Docker image
-      requirements_path: Path to requirements.txt (or None)
+      has_requirements: Whether filtered requirements content is available
       category: Accelerator category ('cpu', 'gpu', 'tpu')
 
   Returns:
@@ -354,7 +357,7 @@ def _generate_dockerfile(
     jax_install = "RUN python3 -m pip install 'jax[cuda12]'"
 
   requirements_section = ""
-  if requirements_path and os.path.exists(requirements_path):
+  if has_requirements:
     requirements_section = (
       "COPY requirements.txt /tmp/requirements.txt\n"
       "RUN python3 -m pip install -r /tmp/requirements.txt\n"

keras_remote/infra/container_builder_test.py

@@ -1,7 +1,5 @@
 """Tests for keras_remote.infra.container_builder — hashing, Dockerfile gen, caching."""
 
-import pathlib
-import tempfile
 from unittest import mock
 from unittest.mock import MagicMock
 
@@ -17,13 +15,6 @@
 )
 
 
-def _make_temp_path(test_case):
-  """Create a temp directory that is cleaned up after the test."""
-  td = tempfile.TemporaryDirectory()
-  test_case.addCleanup(td.cleanup)
-  return pathlib.Path(td.name)
-
-
 class TestFilterJaxRequirements(parameterized.TestCase):
   @parameterized.named_parameters(
     dict(testcase_name="bare_jax", line="jax\n"),
@@ -84,71 +75,42 @@ def test_preserves_comments_and_blanks(self):
 
 class TestHashRequirements(parameterized.TestCase):
   def test_deterministic(self):
-    tmp_path = _make_temp_path(self)
-    req = tmp_path / "requirements.txt"
-    req.write_text("numpy==1.26\n")
-
-    h1 = _hash_requirements(str(req), "gpu", "python:3.12-slim")
-    h2 = _hash_requirements(str(req), "gpu", "python:3.12-slim")
+    h1 = _hash_requirements("numpy==1.26\n", "gpu", "python:3.12-slim")
+    h2 = _hash_requirements("numpy==1.26\n", "gpu", "python:3.12-slim")
     self.assertEqual(h1, h2)
 
   def test_different_requirements_different_hash(self):
-    tmp_path = _make_temp_path(self)
-    req1 = tmp_path / "r1.txt"
-    req1.write_text("numpy==1.26\n")
-    req2 = tmp_path / "r2.txt"
-    req2.write_text("scipy==1.12\n")
-
-    h1 = _hash_requirements(str(req1), "gpu", "python:3.12-slim")
-    h2 = _hash_requirements(str(req2), "gpu", "python:3.12-slim")
+    h1 = _hash_requirements("numpy==1.26\n", "gpu", "python:3.12-slim")
+    h2 = _hash_requirements("scipy==1.12\n", "gpu", "python:3.12-slim")
     self.assertNotEqual(h1, h2)
 
   def test_different_category_different_hash(self):
-    tmp_path = _make_temp_path(self)
-    req = tmp_path / "requirements.txt"
-    req.write_text("numpy\n")
-
-    h1 = _hash_requirements(str(req), "gpu", "python:3.12-slim")
-    h2 = _hash_requirements(str(req), "tpu", "python:3.12-slim")
+    h1 = _hash_requirements("numpy\n", "gpu", "python:3.12-slim")
+    h2 = _hash_requirements("numpy\n", "tpu", "python:3.12-slim")
     self.assertNotEqual(h1, h2)
 
   def test_different_base_image_different_hash(self):
-    tmp_path = _make_temp_path(self)
-    req = tmp_path / "requirements.txt"
-    req.write_text("numpy\n")
-
-    h1 = _hash_requirements(str(req), "gpu", "python:3.12-slim")
-    h2 = _hash_requirements(str(req), "gpu", "python:3.11-slim")
+    h1 = _hash_requirements("numpy\n", "gpu", "python:3.12-slim")
+    h2 = _hash_requirements("numpy\n", "gpu", "python:3.11-slim")
     self.assertNotEqual(h1, h2)
 
-  @parameterized.named_parameters(
-    dict(testcase_name="none", requirements_path=None),
-    dict(
-      testcase_name="nonexistent",
-      requirements_path="/nonexistent/path.txt",
-    ),
-  )
-  def test_missing_requirements_valid(self, requirements_path):
-    h = _hash_requirements(requirements_path, "cpu", "python:3.12-slim")
+  def test_missing_requirements_valid(self):
+    h = _hash_requirements(None, "cpu", "python:3.12-slim")
     self.assertIsInstance(h, str)
     self.assertLen(h, 64)
 
   def test_returns_hex_string(self):
-    tmp_path = _make_temp_path(self)
-    req = tmp_path / "r.txt"
-    req.write_text("keras\n")
-    h = _hash_requirements(str(req), "gpu", "python:3.12-slim")
+    h = _hash_requirements("keras\n", "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")
+    filtered_without_jax = _filter_jax_requirements("numpy==1.26\n")
+    filtered_with_jax = _filter_jax_requirements(
+      "numpy==1.26\njax[tpu]>=0.4.6\n"
+    )
+
+    h1 = _hash_requirements(filtered_without_jax, "tpu", "python:3.12-slim")
+    h2 = _hash_requirements(filtered_with_jax, "tpu", "python:3.12-slim")
     self.assertEqual(h1, h2)
 
 
@@ -176,7 +138,7 @@ class TestGenerateDockerfile(parameterized.TestCase):
   def test_jax_install(self, category, expected, not_expected):
     content = _generate_dockerfile(
       base_image="python:3.12-slim",
-      requirements_path=None,
+      has_requirements=False,
       category=category,
     )
     for s in expected:
@@ -185,13 +147,9 @@ def test_jax_install(self, category, expected, not_expected):
       self.assertNotIn(s, content)
 
   def test_with_requirements(self):
-    tmp_path = _make_temp_path(self)
-    req = tmp_path / "requirements.txt"
-    req.write_text("numpy\n")
-
     content = _generate_dockerfile(
       base_image="python:3.12-slim",
-      requirements_path=str(req),
+      has_requirements=True,
       category="cpu",
     )
     self.assertIn("COPY requirements.txt", content)
@@ -200,7 +158,7 @@ def test_with_requirements(self):
   def test_without_requirements(self):
     content = _generate_dockerfile(
       base_image="python:3.12-slim",
-      requirements_path=None,
+      has_requirements=False,
       category="cpu",
     )
     self.assertNotIn("COPY requirements.txt", content)
@@ -218,15 +176,15 @@ def test_without_requirements(self):
   def test_contains_expected_content(self, expected_substring):
     content = _generate_dockerfile(
       base_image="python:3.12-slim",
-      requirements_path=None,
+      has_requirements=False,
       category="cpu",
     )
     self.assertIn(expected_substring, content)
 
   def test_uses_base_image(self):
     content = _generate_dockerfile(
       base_image="python:3.11-bullseye",
-      requirements_path=None,
+      has_requirements=False,
       category="cpu",
     )
     self.assertIn("FROM python:3.11-bullseye", content)