Address reviews

Author: JyotinderSingh

keras_remote/backend/pathways_client.py

@@ -2,6 +2,7 @@
 
 import time
 
+from absl import logging
 from kubernetes import client
 from kubernetes.client.rest import ApiException
 
@@ -13,9 +14,6 @@
 )
 from keras_remote.backend.log_streaming import LogStreamer
 from keras_remote.core import accelerators
-from keras_remote.infra import infra
-
-logger = infra.logger
 
 LWS_GROUP = "leaderworkerset.x-k8s.io"
 LWS_VERSION = "v1"
@@ -40,7 +38,7 @@ def _get_lws_version(group=LWS_GROUP):
     # If we didn't find the group, raise ApiException to fallback
     raise ApiException(status=404, reason=f"API group {group} not found")
   except ApiException:
-    logger.warning(
+    logging.warning(
       "Failed to retrieve LWS API version from cluster. Defaulting to '%s'",
       LWS_VERSION,
     )
@@ -108,8 +106,8 @@ def submit_pathways_job(
       plural=LWS_PLURAL,
       body=lws_manifest,
     )
-    logger.info(f"Submitted Pathways job (LWS): {job_name}")
-    logger.info(
+    logging.info(f"Submitted Pathways job (LWS): {job_name}")
+    logging.info(
       "View job with: kubectl get %s %s -n %s", LWS_PLURAL, job_name, namespace
     )
     return created_lws
@@ -150,11 +148,11 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
       try:
         pod = core_v1.read_namespaced_pod(leader_pod_name, namespace)
         if not logged_running:
-          logger.info(f"Found pod: {leader_pod_name}")
+          logging.info(f"Found pod: {leader_pod_name}")
           logged_running = True
 
         if pod.status.phase == "Succeeded":
-          logger.info(f"[REMOTE] Job {job_name} completed successfully")
+          logging.info(f"[REMOTE] Job {job_name} completed successfully")
           return "success"
 
         if pod.status.phase == "Failed":
@@ -163,7 +161,7 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
 
         elif pod.status.phase == "Pending":
           _check_pod_scheduling(core_v1, job_name, namespace)
-          logger.debug("Pod is Pending...")
+          logging.debug("Pod is Pending...")
 
         elif pod.status.phase == "Running":
           streamer.start(leader_pod_name)
@@ -183,7 +181,7 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
         # Check current state
         if container_status.state.terminated:
           if container_status.state.terminated.exit_code == 0:
-            logger.info(f"[REMOTE] Job {job_name} completed successfully")
+            logging.info(f"[REMOTE] Job {job_name} completed successfully")
             return "success"
           else:
             _print_pod_logs(core_v1, job_name, namespace)
@@ -195,7 +193,7 @@ def wait_for_job(job_id, namespace="default", timeout=3600, poll_interval=10):
         # Check last state (in case it restarted)
         if container_status.last_state.terminated:
           if container_status.last_state.terminated.exit_code == 0:
-            logger.info(
+            logging.info(
               f"[REMOTE] Job {job_name} completed successfully (restarted)"
             )
             return "success"
@@ -223,13 +221,13 @@ def cleanup_job(job_name, namespace="default"):
       plural=LWS_PLURAL,
       name=job_name,
     )
-    logger.info(f"Deleted LeaderWorkerSet: {job_name}")
+    logging.info(f"Deleted LeaderWorkerSet: {job_name}")
   except ApiException as e:
     if e.status == 404:
       # Job already deleted
       pass
     else:
-      logger.warning(
+      logging.warning(
         "Failed to delete LeaderWorkerSet %s: %s",
         job_name,
         e.reason,

keras_remote/data.py

@@ -6,6 +6,9 @@
 
 import hashlib
 import os
+import posixpath
+
+from absl import logging
 
 
 class Data:
@@ -20,6 +23,14 @@ class Data:
       path: Local file/directory path (absolute or relative) or GCS URI
             (``gs://bucket/prefix``).
 
+  .. note::
+
+      For GCS URIs, a trailing slash indicates a directory (prefix).
+      ``Data("gs://my-bucket/dataset/")`` is treated as a directory,
+      while ``Data("gs://my-bucket/dataset")`` is treated as a single
+      object. If you intend to reference a GCS directory, always
+      include the trailing slash.
+
   Examples::
 
       # Local directory
@@ -28,14 +39,20 @@ class Data:
       # Local file
       Data("./config.json")
 
-      # GCS URI
+      # GCS directory — trailing slash required
       Data("gs://my-bucket/datasets/imagenet/")
+
+      # GCS single object
+      Data("gs://my-bucket/datasets/weights.h5")
   """
 
   def __init__(self, path: str):
+    if not path:
+      raise ValueError("Data path must not be empty")
     self._raw_path = path
     if self.is_gcs:
       self._resolved_path = path
+      _warn_if_missing_trailing_slash(path)
     else:
       self._resolved_path = os.path.abspath(os.path.expanduser(path))
       if not os.path.exists(self._resolved_path):
@@ -63,10 +80,11 @@ def content_hash(self) -> str:
 
     Includes a type prefix ("dir:" or "file:") to prevent collisions
     between a single file and a directory containing only that file.
-    Symlinks are not followed (followlinks=False) to ensure
-    deterministic hashing and prevent circular symlink infinite
-    recursion. Users with symlinked data should pass the resolved
-    target path.
+
+    Symlinked directories are not recursed into (followlinks=False)
+    to prevent infinite recursion from circular symlinks. Symlinked
+    files are read and their resolved contents are hashed, so the
+    hash reflects the actual data visible at runtime.
     """
     if self.is_gcs:
       raise ValueError("Cannot compute content hash for GCS URI")
@@ -80,15 +98,18 @@ def content_hash(self) -> str:
           fpath = os.path.join(root, fname)
           relpath = os.path.relpath(fpath, self._resolved_path)
           h.update(relpath.encode("utf-8"))
+          h.update(b"\0")
           with open(fpath, "rb") as f:
             while True:
               chunk = f.read(65536)  # 64 KB chunks
               if not chunk:
                 break
               h.update(chunk)
+          h.update(b"\0")
     else:
       h.update(b"file:")
       h.update(os.path.basename(self._resolved_path).encode("utf-8"))
+      h.update(b"\0")
       with open(self._resolved_path, "rb") as f:
         while True:
           chunk = f.read(65536)
@@ -101,6 +122,23 @@ def __repr__(self):
     return f"Data({self._raw_path!r})"
 
 
+def _warn_if_missing_trailing_slash(path: str) -> None:
+  """Log a warning if a GCS path looks like a directory but has no trailing slash."""
+  if path.endswith("/"):
+    return
+  gcs_path = path.split("//", 1)[1]  # strip gs://
+  last_segment = posixpath.basename(gcs_path)
+  if last_segment and "." not in last_segment:
+    logging.warning(
+      "GCS path %r does not end with '/' but the last segment "
+      "(%r) has no file extension. If this is a directory "
+      "(prefix), add a trailing slash: %r",
+      path,
+      last_segment,
+      path + "/",
+    )
+
+
 def _make_data_ref(
   gcs_uri: str, is_dir: bool, mount_path: str | None = None
 ) -> dict[str, object]:

keras_remote/data_test.py

@@ -49,6 +49,10 @@ def test_gcs_uri_file(self):
     self.assertTrue(d.is_gcs)
     self.assertFalse(d.is_dir)
 
+  def test_empty_path_raises(self):
+    with self.assertRaises(ValueError):
+      Data("")
+
   def test_nonexistent_path_raises(self):
     with self.assertRaises(FileNotFoundError) as cm:
       Data("/nonexistent/path/to/data")
@@ -156,6 +160,46 @@ def test_nested_directory_hash(self):
     self.assertIsInstance(h, str)
     self.assertEqual(len(h), 64)
 
+  def test_filename_content_boundary(self):
+    """Filename/content collisions must produce different hashes.
+
+    Without a delimiter, file "a" with content "bc" and file "ab" with
+    content "c" would both hash the byte sequence "abc".
+    """
+    tmp = _make_temp_path(self)
+    d1 = tmp / "dir1"
+    d1.mkdir()
+    (d1 / "a").write_text("bc")
+
+    d2 = tmp / "dir2"
+    d2.mkdir()
+    (d2 / "ab").write_text("c")
+
+    self.assertNotEqual(
+      Data(str(d1)).content_hash(), Data(str(d2)).content_hash()
+    )
+
+  def test_file_boundary_across_entries(self):
+    """Consecutive file entries must not collide.
+
+    Without a delimiter between entries, two files ["x" -> "y", "z" -> ""]
+    and ["x" -> "", "yz" -> ""] would produce the same hash input.
+    """
+    tmp = _make_temp_path(self)
+    d1 = tmp / "dir1"
+    d1.mkdir()
+    (d1 / "x").write_text("y")
+    (d1 / "z").write_text("")
+
+    d2 = tmp / "dir2"
+    d2.mkdir()
+    (d2 / "x").write_text("")
+    (d2 / "yz").write_text("")
+
+    self.assertNotEqual(
+      Data(str(d1)).content_hash(), Data(str(d2)).content_hash()
+    )
+
   def test_path_included_in_hash(self):
     """Files with same content but different names produce different
     hashes."""

keras_remote/infra/infra.py

@@ -1,9 +1,5 @@
-import logging
 import os
 
-logging.basicConfig(level=logging.INFO, format="%(message)s")
-logger = logging.getLogger("keras_remote")
-
 
 def get_default_project() -> str | None:
   """Get project ID from KERAS_REMOTE_PROJECT or GOOGLE_CLOUD_PROJECT."""

keras_remote/utils/storage.py

@@ -139,6 +139,7 @@ def upload_data(
     return data.path
 
   content_hash = data.content_hash()
+  namespace_prefix = namespace_prefix.strip("/")
   cache_prefix = f"{namespace_prefix}/data-cache/{content_hash}"
 
   project = project or get_default_project()
@@ -207,6 +208,6 @@ def _upload_directory(
   for root, _dirs, files in os.walk(local_dir):
     for fname in files:
       local_path = os.path.join(root, fname)
-      rel_path = os.path.relpath(local_path, local_dir)
+      rel_path = os.path.relpath(local_path, local_dir).replace(os.sep, "/")
       blob = bucket.blob(f"{gcs_prefix}/{rel_path}")
       blob.upload_from_filename(local_path)