Parallelize GCS data transfers and deduplicate downloads (#82)

Author: JyotinderSingh

keras_remote/runner/remote_runner.py

@@ -15,6 +15,9 @@
 import cloudpickle
 from absl import logging
 from google.cloud import storage
+from google.cloud.storage import transfer_manager
+
+_DOWNLOAD_BATCH_SIZE = 10000
 
 # Base temp directory for remote execution artifacts
 TEMP_DIR = tempfile.gettempdir()
@@ -147,6 +150,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 +159,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 +197,39 @@ def _download_data(ref, target_dir, storage_client):
   bucket = storage_client.bucket(bucket_name)
 
   blobs = bucket.list_blobs(prefix=prefix + "/")
-  count = 0
+  total_downloaded = 0
+  batch = []
   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
+    batch.append(blob.name[len(prefix) + 1 :])
+    if len(batch) >= _DOWNLOAD_BATCH_SIZE:
+      transfer_manager.download_many_to_path(
+        bucket,
+        batch,
+        destination_directory=target_dir,
+        blob_name_prefix=prefix + "/",
+        worker_type=transfer_manager.THREAD,
+        raise_exception=True,
+      )
+      total_downloaded += len(batch)
+      batch = []
+
+  if batch:
+    transfer_manager.download_many_to_path(
+      bucket,
+      batch,
+      destination_directory=target_dir,
+      blob_name_prefix=prefix + "/",
+      worker_type=transfer_manager.THREAD,
+      raise_exception=True,
+    )
+    total_downloaded += len(batch)
 
-  logging.info("Downloaded %d files from %s to %s", count, gcs_uri, target_dir)
+  if total_downloaded:
+    logging.info(
+      "Downloaded %d files from %s to %s", total_downloaded, gcs_uri, target_dir
+    )
 
 
 def _download_from_gcs(client, gcs_path, local_path):

keras_remote/runner/remote_runner_test.py

@@ -13,6 +13,7 @@
 from absl.testing import absltest
 
 from keras_remote.runner.remote_runner import (
+  _DOWNLOAD_BATCH_SIZE,
   _download_data,
   _download_from_gcs,
   _upload_to_gcs,
@@ -81,6 +82,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 +101,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 +122,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 +145,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 +155,57 @@ 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_large_listing_downloads_in_batches(self):
+    tmp = _make_temp_path(self)
+    target = tmp / "output"
+
+    mock_client = MagicMock()
+    mock_bucket = MagicMock()
+    mock_client.bucket.return_value = mock_bucket
+
+    num_blobs = _DOWNLOAD_BATCH_SIZE + 5
+    blobs = []
+    for i in range(num_blobs):
+      blob = MagicMock()
+      blob.name = f"prefix/hash/file_{i}.csv"
+      blobs.append(blob)
+    mock_bucket.list_blobs.return_value = blobs
+
+    ref = {
+      "__data_ref__": True,
+      "gcs_uri": "gs://bucket/prefix/hash",
+      "is_dir": True,
+    }
+
+    _download_data(ref, str(target), mock_client)
+
+    self.assertEqual(self.mock_download.call_count, 2)
+    first_batch = self.mock_download.call_args_list[0][0][1]
+    second_batch = self.mock_download.call_args_list[1][0][1]
+    self.assertEqual(len(first_batch), _DOWNLOAD_BATCH_SIZE)
+    self.assertEqual(len(second_batch), 5)
+
+  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 +257,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 +265,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,
+  )