Add Data class and content-addressed storage upload

Introduces the Data class for declaring data dependencies (local paths
or GCS URIs) and upload_data() for content-hash-based caching in GCS.

Author: JyotinderSingh

keras_remote/__init__.py

@@ -6,3 +6,4 @@
 os.environ.setdefault("GRPC_ENABLE_FORK_SUPPORT", "0")
 
 from keras_remote.core.core import run as run
+from keras_remote.data import Data as Data

keras_remote/data.py

@@ -0,0 +1,120 @@
+"""Data class for declaring data dependencies in remote functions.
+
+Wraps local file/directory paths or GCS URIs. On the remote side, Data
+resolves to a plain filesystem path — the user's function only sees paths.
+"""
+
+import hashlib
+import os
+
+
+class Data:
+  """A reference to data that should be available on the remote pod.
+
+  Wraps a local file/directory path or a GCS URI. When passed as a function
+  argument or used in the ``volumes`` decorator parameter, Data is resolved
+  to a plain filesystem path on the remote side. The user's function code
+  never needs to know about Data — it just receives paths.
+
+  Args:
+      path: Local file/directory path (absolute or relative) or GCS URI
+            (``gs://bucket/prefix``).
+
+  Examples::
+
+      # Local directory
+      Data("./my_dataset/")
+
+      # Local file
+      Data("./config.json")
+
+      # GCS URI
+      Data("gs://my-bucket/datasets/imagenet/")
+  """
+
+  def __init__(self, path: str):
+    self._raw_path = path
+    if self.is_gcs:
+      self._resolved_path = path
+    else:
+      self._resolved_path = os.path.abspath(os.path.expanduser(path))
+      if not os.path.exists(self._resolved_path):
+        raise FileNotFoundError(
+          f"Data path does not exist: {path} "
+          f"(resolved to {self._resolved_path})"
+        )
+
+  @property
+  def path(self) -> str:
+    return self._resolved_path
+
+  @property
+  def is_gcs(self) -> bool:
+    return self._raw_path.startswith("gs://")
+
+  @property
+  def is_dir(self) -> bool:
+    if self.is_gcs:
+      return self._raw_path.endswith("/")
+    return os.path.isdir(self._resolved_path)
+
+  def content_hash(self) -> str:
+    """SHA-256 hash of all file contents, sorted by relative path.
+
+    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.
+    """
+    if self.is_gcs:
+      raise ValueError("Cannot compute content hash for GCS URI")
+
+    h = hashlib.sha256()
+    if os.path.isdir(self._resolved_path):
+      h.update(b"dir:")
+      for root, dirs, files in os.walk(self._resolved_path, followlinks=False):
+        dirs.sort()
+        for fname in sorted(files):
+          fpath = os.path.join(root, fname)
+          relpath = os.path.relpath(fpath, self._resolved_path)
+          h.update(relpath.encode("utf-8"))
+          with open(fpath, "rb") as f:
+            while True:
+              chunk = f.read(65536)  # 64 KB chunks
+              if not chunk:
+                break
+              h.update(chunk)
+    else:
+      h.update(b"file:")
+      h.update(os.path.basename(self._resolved_path).encode("utf-8"))
+      with open(self._resolved_path, "rb") as f:
+        while True:
+          chunk = f.read(65536)
+          if not chunk:
+            break
+          h.update(chunk)
+    return h.hexdigest()
+
+  def __repr__(self):
+    return f"Data({self._raw_path!r})"
+
+
+def _make_data_ref(gcs_uri, is_dir, mount_path=None):
+  """Create a serializable data reference dict.
+
+  These dicts replace Data objects in the payload before serialization.
+  The remote runner identifies them by the __data_ref__ key.
+  """
+  return {
+    "__data_ref__": True,
+    "gcs_uri": gcs_uri,
+    "is_dir": is_dir,
+    "mount_path": mount_path,
+  }
+
+
+def is_data_ref(obj):
+  """Check if an object is a serialized data reference."""
+  return isinstance(obj, dict) and obj.get("__data_ref__") is True

keras_remote/data_test.py

@@ -0,0 +1,205 @@
+"""Tests for keras_remote.data — Data class and helpers."""
+
+import os
+import pathlib
+import tempfile
+
+from absl.testing import absltest
+
+from keras_remote.data import Data, _make_data_ref, is_data_ref
+
+
+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 TestDataConstructor(absltest.TestCase):
+  def test_local_file(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("a,b\n1,2\n")
+    d = Data(str(f))
+
+    self.assertEqual(d.path, str(f))
+    self.assertFalse(d.is_gcs)
+    self.assertFalse(d.is_dir)
+
+  def test_local_directory(self):
+    tmp = _make_temp_path(self)
+    d_dir = tmp / "dataset"
+    d_dir.mkdir()
+    (d_dir / "train.csv").write_text("data")
+    d = Data(str(d_dir))
+
+    self.assertEqual(d.path, str(d_dir))
+    self.assertFalse(d.is_gcs)
+    self.assertTrue(d.is_dir)
+
+  def test_gcs_uri_directory(self):
+    d = Data("gs://my-bucket/data/")
+    self.assertEqual(d.path, "gs://my-bucket/data/")
+    self.assertTrue(d.is_gcs)
+    self.assertTrue(d.is_dir)
+
+  def test_gcs_uri_file(self):
+    d = Data("gs://my-bucket/data/file.csv")
+    self.assertTrue(d.is_gcs)
+    self.assertFalse(d.is_dir)
+
+  def test_nonexistent_path_raises(self):
+    with self.assertRaises(FileNotFoundError) as cm:
+      Data("/nonexistent/path/to/data")
+    self.assertIn("/nonexistent/path/to/data", str(cm.exception))
+
+  def test_relative_path_resolved(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "file.txt"
+    f.write_text("content")
+    # Use relative-like path by going through expanduser
+    d = Data(str(f))
+    self.assertTrue(os.path.isabs(d.path))
+
+  def test_repr(self):
+    d = Data("gs://bucket/path/")
+    self.assertEqual(repr(d), "Data('gs://bucket/path/')")
+
+
+class TestContentHash(absltest.TestCase):
+  def test_deterministic_file_hash(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("hello,world\n")
+
+    d1 = Data(str(f))
+    d2 = Data(str(f))
+    self.assertEqual(d1.content_hash(), d2.content_hash())
+
+  def test_different_content_different_hash(self):
+    tmp = _make_temp_path(self)
+    f1 = tmp / "a.csv"
+    f1.write_text("content_a")
+    f2 = tmp / "b.csv"
+    f2.write_text("content_b")
+
+    self.assertNotEqual(
+      Data(str(f1)).content_hash(), Data(str(f2)).content_hash()
+    )
+
+  def test_deterministic_dir_hash(self):
+    tmp = _make_temp_path(self)
+    d = tmp / "dataset"
+    d.mkdir()
+    (d / "train.csv").write_text("train data")
+    (d / "val.csv").write_text("val data")
+
+    d1 = Data(str(d))
+    d2 = Data(str(d))
+    self.assertEqual(d1.content_hash(), d2.content_hash())
+
+  def test_dir_content_change_changes_hash(self):
+    tmp = _make_temp_path(self)
+    d = tmp / "dataset"
+    d.mkdir()
+    (d / "train.csv").write_text("original")
+
+    h1 = Data(str(d)).content_hash()
+    (d / "train.csv").write_text("modified")
+    h2 = Data(str(d)).content_hash()
+
+    self.assertNotEqual(h1, h2)
+
+  def test_file_vs_dir_different_hash(self):
+    """A single file and a directory containing only that file should
+    produce different hashes due to the type prefix."""
+    tmp = _make_temp_path(self)
+
+    # Single file
+    f = tmp / "file.csv"
+    f.write_text("same content")
+
+    # Directory containing only that file
+    d = tmp / "dir"
+    d.mkdir()
+    (d / "file.csv").write_text("same content")
+
+    file_hash = Data(str(f)).content_hash()
+    dir_hash = Data(str(d)).content_hash()
+    self.assertNotEqual(file_hash, dir_hash)
+
+  def test_empty_directory(self):
+    tmp = _make_temp_path(self)
+    d = tmp / "empty"
+    d.mkdir()
+
+    # Should not raise, should return a valid hash
+    h = Data(str(d)).content_hash()
+    self.assertIsInstance(h, str)
+    self.assertEqual(len(h), 64)  # SHA-256 hex digest
+
+  def test_gcs_uri_raises(self):
+    d = Data("gs://bucket/data/")
+    with self.assertRaises(ValueError):
+      d.content_hash()
+
+  def test_nested_directory_hash(self):
+    tmp = _make_temp_path(self)
+    d = tmp / "nested"
+    sub = d / "sub"
+    sub.mkdir(parents=True)
+    (d / "a.txt").write_text("a")
+    (sub / "b.txt").write_text("b")
+
+    h = Data(str(d)).content_hash()
+    self.assertIsInstance(h, str)
+    self.assertEqual(len(h), 64)
+
+  def test_path_included_in_hash(self):
+    """Files with same content but different names produce different
+    hashes."""
+    tmp = _make_temp_path(self)
+    d1 = tmp / "dir1"
+    d1.mkdir()
+    (d1 / "alpha.csv").write_text("same")
+
+    d2 = tmp / "dir2"
+    d2.mkdir()
+    (d2 / "beta.csv").write_text("same")
+
+    self.assertNotEqual(
+      Data(str(d1)).content_hash(), Data(str(d2)).content_hash()
+    )
+
+
+class TestMakeDataRef(absltest.TestCase):
+  def test_basic_ref(self):
+    ref = _make_data_ref("gs://b/prefix", True)
+    self.assertTrue(ref["__data_ref__"])
+    self.assertEqual(ref["gcs_uri"], "gs://b/prefix")
+    self.assertTrue(ref["is_dir"])
+    self.assertIsNone(ref["mount_path"])
+
+  def test_with_mount_path(self):
+    ref = _make_data_ref("gs://b/p", False, mount_path="/data")
+    self.assertEqual(ref["mount_path"], "/data")
+    self.assertFalse(ref["is_dir"])
+
+
+class TestIsDataRef(absltest.TestCase):
+  def test_valid_ref(self):
+    ref = {"__data_ref__": True, "gcs_uri": "gs://b/p", "is_dir": True}
+    self.assertTrue(is_data_ref(ref))
+
+  def test_plain_dict(self):
+    self.assertFalse(is_data_ref({"key": "value"}))
+
+  def test_non_dict(self):
+    self.assertFalse(is_data_ref("string"))
+    self.assertFalse(is_data_ref(42))
+    self.assertFalse(is_data_ref(None))
+
+
+if __name__ == "__main__":
+  absltest.main()