Adds parallel data upload/download and remote-side deduplication

Author: JyotinderSingh

keras_remote/runner/remote_runner.py

@@ -15,6 +15,7 @@
 import cloudpickle
 from absl import logging
 from google.cloud import storage
+from google.cloud.storage import transfer_manager
 
 # Base temp directory for remote execution artifacts
 TEMP_DIR = tempfile.gettempdir()
@@ -147,6 +148,7 @@ def resolve_volumes(volume_refs, storage_client):
 def resolve_data_refs(args, kwargs, storage_client):
   """Recursively resolve data ref dicts in args/kwargs to local paths."""
   counter = 0
+  resolved_uris = {}
 
   def _resolve(obj):
     nonlocal counter
@@ -155,14 +157,20 @@ def _resolve(obj):
       # Volume-mounted data refs are handled by Kubernetes, skip download
       if obj.get("mount_path") is not None:
         return obj["mount_path"]
+      gcs_uri = obj["gcs_uri"]
+      if gcs_uri in resolved_uris:
+        return resolved_uris[gcs_uri]
       local_dir = os.path.join(DATA_DIR, str(counter))
       counter += 1
       _download_data(obj, local_dir, storage_client)
       # Return file path for single files, directory path otherwise
       if not obj["is_dir"]:
         files = [f for f in os.listdir(local_dir) if f != ".cache_marker"]
         if len(files) == 1:
-          return os.path.join(local_dir, files[0])
+          path = os.path.join(local_dir, files[0])
+          resolved_uris[gcs_uri] = path
+          return path
+      resolved_uris[gcs_uri] = local_dir
       return local_dir
     # Recurse into containers to find nested data refs
     if isinstance(obj, dict):
@@ -187,17 +195,27 @@ def _download_data(ref, target_dir, storage_client):
   bucket = storage_client.bucket(bucket_name)
 
   blobs = bucket.list_blobs(prefix=prefix + "/")
-  count = 0
+  blob_names = []
   for blob in blobs:
     if blob.name.endswith("/") or blob.name.endswith(".cache_marker"):
       continue
-    rel_path = blob.name[len(prefix) + 1 :]
-    local_path = os.path.join(target_dir, rel_path)
-    os.makedirs(os.path.dirname(local_path), exist_ok=True)
-    blob.download_to_filename(local_path)
-    count += 1
-
-  logging.info("Downloaded %d files from %s to %s", count, gcs_uri, target_dir)
+    blob_names.append(blob.name[len(prefix) + 1 :])
+
+  if not blob_names:
+    return
+
+  transfer_manager.download_many_to_path(
+    bucket,
+    blob_names,
+    destination_directory=target_dir,
+    blob_name_prefix=prefix + "/",
+    worker_type=transfer_manager.THREAD,
+    raise_exception=True,
+  )
+
+  logging.info(
+    "Downloaded %d files from %s to %s", len(blob_names), gcs_uri, target_dir
+  )
 
 
 def _download_from_gcs(client, gcs_path, local_path):

keras_remote/runner/remote_runner_test.py

@@ -81,6 +81,15 @@ def test_parses_gcs_path(self):
 
 
 class TestDownloadData(absltest.TestCase):
+  def setUp(self):
+    super().setUp()
+    self.mock_download = self.enterContext(
+      mock.patch(
+        "keras_remote.runner.remote_runner.transfer_manager"
+        ".download_many_to_path",
+      )
+    )
+
   def test_downloads_files_skips_marker(self):
     tmp = _make_temp_path(self)
     target = tmp / "output"
@@ -91,9 +100,6 @@ def test_downloads_files_skips_marker(self):
 
     blob_data = MagicMock()
     blob_data.name = "prefix/hash/train.csv"
-    blob_data.download_to_filename = MagicMock(
-      side_effect=lambda p: pathlib.Path(p).write_text("train")
-    )
 
     blob_marker = MagicMock()
     blob_marker.name = "prefix/hash/.cache_marker"
@@ -115,9 +121,18 @@ def test_downloads_files_skips_marker(self):
 
     _download_data(ref, str(target), mock_client)
 
-    blob_data.download_to_filename.assert_called_once()
-    blob_marker.download_to_filename.assert_not_called()
-    blob_dir.download_to_filename.assert_not_called()
+    self.mock_download.assert_called_once()
+    blob_names = self.mock_download.call_args[0][1]
+    self.assertEqual(blob_names, ["train.csv"])
+    self.assertEqual(
+      self.mock_download.call_args.kwargs["destination_directory"],
+      str(target),
+    )
+    self.assertEqual(
+      self.mock_download.call_args.kwargs["blob_name_prefix"],
+      "prefix/hash/",
+    )
+    self.assertTrue(self.mock_download.call_args.kwargs["raise_exception"])
 
   def test_creates_subdirectories(self):
     tmp = _make_temp_path(self)
@@ -129,12 +144,6 @@ def test_creates_subdirectories(self):
 
     blob = MagicMock()
     blob.name = "prefix/hash/sub/deep.csv"
-    blob.download_to_filename = MagicMock(
-      side_effect=lambda p: (
-        pathlib.Path(p).parent.mkdir(parents=True, exist_ok=True)
-        or pathlib.Path(p).write_text("data")
-      )
-    )
     mock_bucket.list_blobs.return_value = [blob]
 
     ref = {
@@ -145,10 +154,27 @@ def test_creates_subdirectories(self):
 
     _download_data(ref, str(target), mock_client)
 
-    # The call should include the nested path
-    call_path = blob.download_to_filename.call_args[0][0]
-    self.assertIn("sub", call_path)
-    self.assertTrue(call_path.endswith("deep.csv"))
+    blob_names = self.mock_download.call_args[0][1]
+    self.assertEqual(blob_names, ["sub/deep.csv"])
+
+  def test_empty_listing_is_noop(self):
+    tmp = _make_temp_path(self)
+    target = tmp / "output"
+
+    mock_client = MagicMock()
+    mock_bucket = MagicMock()
+    mock_client.bucket.return_value = mock_bucket
+    mock_bucket.list_blobs.return_value = []
+
+    ref = {
+      "__data_ref__": True,
+      "gcs_uri": "gs://bucket/prefix/hash",
+      "is_dir": True,
+    }
+
+    _download_data(ref, str(target), mock_client)
+
+    self.mock_download.assert_not_called()
 
 
 class TestResolveDataRefs(absltest.TestCase):
@@ -200,16 +226,6 @@ def test_nested_refs_in_list(self):
 
   def test_single_file_returns_file_path(self):
     tmp = _make_temp_path(self)
-    mock_client = MagicMock()
-    mock_bucket = MagicMock()
-    mock_client.bucket.return_value = mock_bucket
-
-    blob = MagicMock()
-    blob.name = "prefix/hash/config.json"
-    blob.download_to_filename = MagicMock(
-      side_effect=lambda p: pathlib.Path(p).write_text("{}")
-    )
-    mock_bucket.list_blobs.return_value = [blob]
 
     ref = {
       "__data_ref__": True,
@@ -218,14 +234,55 @@ def test_single_file_returns_file_path(self):
       "mount_path": None,
     }
 
-    with mock.patch(
-      "keras_remote.runner.remote_runner.DATA_DIR",
-      str(tmp / "data"),
+    def fake_dl(ref, target_dir, client):
+      os.makedirs(target_dir, exist_ok=True)
+      pathlib.Path(os.path.join(target_dir, "config.json")).write_text("{}")
+
+    with (
+      mock.patch(
+        "keras_remote.runner.remote_runner.DATA_DIR",
+        str(tmp / "data"),
+      ),
+      mock.patch(
+        "keras_remote.runner.remote_runner._download_data",
+        side_effect=fake_dl,
+      ),
     ):
-      args, _ = resolve_data_refs((ref,), {}, mock_client)
+      args, _ = resolve_data_refs((ref,), {}, MagicMock())
 
     self.assertTrue(args[0].endswith("config.json"))
 
+  def test_duplicate_uri_downloaded_once(self):
+    tmp = _make_temp_path(self)
+
+    ref = {
+      "__data_ref__": True,
+      "gcs_uri": "gs://b/cache/hash",
+      "is_dir": True,
+      "mount_path": None,
+    }
+
+    def fake_dl(r, target_dir, client):
+      os.makedirs(target_dir, exist_ok=True)
+
+    with (
+      mock.patch(
+        "keras_remote.runner.remote_runner.DATA_DIR",
+        str(tmp / "data"),
+      ),
+      mock.patch(
+        "keras_remote.runner.remote_runner._download_data",
+        side_effect=fake_dl,
+      ) as mock_dl,
+    ):
+      args, kwargs = resolve_data_refs((ref, ref), {"d": ref}, MagicMock())
+
+    # Downloaded only once despite three references
+    mock_dl.assert_called_once()
+    # All resolved paths point to the same directory
+    self.assertEqual(args[0], args[1])
+    self.assertEqual(args[0], kwargs["d"])
+
   def test_non_ref_dict_preserved(self):
     mock_client = MagicMock()
     args, kwargs = resolve_data_refs(({"key": "value"},), {"x": 1}, mock_client)

keras_remote/utils/storage.py

@@ -7,6 +7,7 @@
 
 from absl import logging
 from google.cloud import storage
+from google.cloud.storage import transfer_manager
 
 from keras_remote.data import Data
 from keras_remote.infra.infra import get_default_project
@@ -205,9 +206,23 @@ def _upload_directory(
   bucket: storage.Bucket, local_dir: str, gcs_prefix: str
 ) -> None:
   """Upload a local directory to GCS preserving structure."""
+  filenames = []
   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).replace(os.sep, "/")
-      blob = bucket.blob(f"{gcs_prefix}/{rel_path}")
-      blob.upload_from_filename(local_path)
+      filenames.append(rel_path)
+
+  if not filenames:
+    return
+
+  logging.info("Uploading %d files to GCS...", len(filenames))
+
+  transfer_manager.upload_many_from_filenames(
+    bucket,
+    filenames,
+    source_directory=local_dir,
+    blob_name_prefix=f"{gcs_prefix}/",
+    worker_type=transfer_manager.THREAD,
+    raise_exception=True,
+  )

keras_remote/utils/storage_test.py

@@ -202,33 +202,33 @@ def track_blob(name):
     self.assertIn(marker_name, blobs)
     blobs[marker_name].upload_from_string.assert_called_once_with("")
 
-  def test_cache_miss_uploads_directory(self):
+  @mock.patch(
+    "keras_remote.utils.storage.transfer_manager.upload_many_from_filenames",
+  )
+  def test_cache_miss_uploads_directory(self, mock_upload):
     tmp = _make_temp_path(self)
     d_dir = tmp / "dataset"
     d_dir.mkdir()
     (d_dir / "train.csv").write_text("train")
     (d_dir / "val.csv").write_text("val")
     d = Data(str(d_dir))
+    content_hash = d.content_hash()
 
     mock_bucket = self.mock_gcs.bucket.return_value
-    blobs = {}
-
-    def track_blob(name):
-      b = MagicMock()
-      blobs[name] = b
-      if name.endswith(".cache_marker"):
-        b.exists.return_value = False
-      return b
-
-    mock_bucket.blob.side_effect = track_blob
+    marker_blob = MagicMock()
+    marker_blob.exists.return_value = False
+    mock_bucket.blob.return_value = marker_blob
 
     result = upload_data("jobs-bucket", d, project="proj")
 
-    self.assertIn("gs://jobs-bucket/default/data-cache/", result)
-    # Both files + marker should have blobs
-    blob_names = list(blobs.keys())
-    csv_blobs = [n for n in blob_names if n.endswith(".csv")]
-    self.assertEqual(len(csv_blobs), 2)
+    expected_prefix = f"default/data-cache/{content_hash}"
+    self.assertEqual(result, f"gs://jobs-bucket/{expected_prefix}")
+    # Directory upload via transfer_manager
+    mock_upload.assert_called_once()
+    filenames = sorted(mock_upload.call_args[0][1])
+    self.assertEqual(filenames, ["train.csv", "val.csv"])
+    # Marker written after upload
+    marker_blob.upload_from_string.assert_called_once_with("")
 
   def test_custom_namespace(self):
     tmp = _make_temp_path(self)
@@ -273,7 +273,16 @@ def test_empty_directory(self):
     self.assertEqual(_compute_total_size(str(d)), 0)
 
 
-class TestUploadDirectory(_GcsTestBase):
+class TestUploadDirectory(absltest.TestCase):
+  def setUp(self):
+    super().setUp()
+    self.mock_upload = self.enterContext(
+      mock.patch(
+        "keras_remote.utils.storage.transfer_manager"
+        ".upload_many_from_filenames",
+      )
+    )
+
   def test_preserves_structure(self):
     tmp = _make_temp_path(self)
     d = tmp / "dataset"
@@ -283,21 +292,27 @@ def test_preserves_structure(self):
     (sub / "b.csv").write_text("b")
 
     mock_bucket = MagicMock()
-    uploaded = {}
 
-    def track_blob(name):
-      b = MagicMock()
-      uploaded[name] = b
-      return b
+    _upload_directory(mock_bucket, str(d), "prefix/hash")
 
-    mock_bucket.blob.side_effect = track_blob
+    self.mock_upload.assert_called_once()
+    call_kwargs = self.mock_upload.call_args
+    filenames = sorted(call_kwargs[0][1])  # second positional arg
+    self.assertEqual(filenames, ["a.csv", "sub/b.csv"])
+    self.assertEqual(call_kwargs.kwargs["source_directory"], str(d))
+    self.assertEqual(call_kwargs.kwargs["blob_name_prefix"], "prefix/hash/")
+    self.assertTrue(call_kwargs.kwargs["raise_exception"])
+
+  def test_empty_directory_is_noop(self):
+    tmp = _make_temp_path(self)
+    d = tmp / "empty_dataset"
+    d.mkdir()
+
+    mock_bucket = MagicMock()
 
     _upload_directory(mock_bucket, str(d), "prefix/hash")
 
-    self.assertIn("prefix/hash/a.csv", uploaded)
-    self.assertIn("prefix/hash/sub/b.csv", uploaded)
-    for blob in uploaded.values():
-      blob.upload_from_filename.assert_called_once()
+    self.mock_upload.assert_not_called()
 
 
 if __name__ == "__main__":