Integrate Data API into execution pipeline (#66)

Author: JyotinderSingh

keras_remote/backend/execution.py

@@ -18,6 +18,7 @@
 from keras_remote.backend import gke_client, pathways_client
 from keras_remote.constants import get_default_zone, zone_to_region
 from keras_remote.credentials import ensure_credentials
+from keras_remote.data import _make_data_ref
 from keras_remote.infra import container_builder
 from keras_remote.infra.infra import get_default_project
 from keras_remote.utils import packager, storage
@@ -47,6 +48,9 @@ class JobContext:
   region: str = field(init=False)
   display_name: str = field(init=False)
 
+  # Data volumes {mount_path: Data}
+  volumes: Optional[dict] = None
+
   # Artifact paths (set during prepare phase)
   payload_path: Optional[str] = None
   context_path: Optional[str] = None
@@ -69,6 +73,7 @@ def from_params(
     zone: Optional[str],
     project: Optional[str],
     env_vars: dict,
+    volumes: Optional[dict] = None,
   ) -> "JobContext":
     """Factory method with default resolution for zone/project."""
     if not zone:
@@ -90,13 +95,14 @@ def from_params(
       container_image=container_image,
       zone=zone,
       project=project,
+      volumes=volumes,
     )
 
 
 class BaseK8sBackend:
   """Base class for Kubernetes-based backends."""
 
-  def __init__(self, cluster: Optional[str] = None, namespace: str = "default"):
+  def __init__(self, cluster: str, namespace: str = "default"):
     self.cluster = cluster
     self.namespace = namespace
 
@@ -202,6 +208,14 @@ def _find_requirements(start_dir: str) -> Optional[str]:
   return None
 
 
+def _maybe_exclude(data_path, caller_path, exclude_paths):
+  """Add data_path to exclude_paths if it's inside the caller directory."""
+  data_abs = os.path.normpath(data_path)
+  caller_abs = os.path.normpath(caller_path)
+  if data_abs.startswith(caller_abs + os.sep) or data_abs == caller_abs:
+    exclude_paths.add(data_abs)
+
+
 def _prepare_artifacts(
   ctx: JobContext, tmpdir: str, caller_frame_depth: int = 3
 ) -> None:
@@ -211,21 +225,58 @@ def _prepare_artifacts(
   # Get caller directory
   frame = inspect.stack()[caller_frame_depth]
   module = inspect.getmodule(frame[0])
-  if module:
+  caller_path: str
+  if module and module.__file__:
     caller_path = os.path.dirname(os.path.abspath(module.__file__))
   else:
     caller_path = os.getcwd()
 
-  # Serialize function + args
+  # Process Data objects
+  exclude_paths: set[str] = set()
+  ref_map = {}  # id(Data) -> ref dict (for arg replacement)
+  volume_refs = []  # list of ref dicts (for volumes)
+
+  # Process volumes
+  if ctx.volumes:
+    for mount_path, data_obj in ctx.volumes.items():
+      gcs_uri = storage.upload_data(ctx.bucket_name, data_obj, ctx.project)
+      volume_refs.append(
+        _make_data_ref(gcs_uri, data_obj.is_dir, mount_path=mount_path)
+      )
+      if not data_obj.is_gcs:
+        _maybe_exclude(data_obj.path, caller_path, exclude_paths)
+
+  # Process Data in function args
+  data_refs = packager.extract_data_refs(ctx.args, ctx.kwargs)
+  for data_obj, _position in data_refs:
+    gcs_uri = storage.upload_data(ctx.bucket_name, data_obj, ctx.project)
+    ref_map[id(data_obj)] = _make_data_ref(gcs_uri, data_obj.is_dir)
+    if not data_obj.is_gcs:
+      _maybe_exclude(data_obj.path, caller_path, exclude_paths)
+
+  # Replace Data with refs in args/kwargs
+  if ref_map:
+    ctx.args, ctx.kwargs = packager.replace_data_with_refs(
+      ctx.args, ctx.kwargs, ref_map
+    )
+
+  # Serialize function + args (with volume refs)
   ctx.payload_path = os.path.join(tmpdir, "payload.pkl")
   packager.save_payload(
-    ctx.func, ctx.args, ctx.kwargs, ctx.env_vars, ctx.payload_path
+    ctx.func,
+    ctx.args,
+    ctx.kwargs,
+    ctx.env_vars,
+    ctx.payload_path,
+    volumes=volume_refs or None,
   )
   logging.info("Payload serialized to %s", ctx.payload_path)
 
-  # Zip working directory
+  # Zip working directory (excluding Data paths)
   ctx.context_path = os.path.join(tmpdir, "context.zip")
-  packager.zip_working_dir(caller_path, ctx.context_path)
+  packager.zip_working_dir(
+    caller_path, ctx.context_path, exclude_paths=exclude_paths
+  )
   logging.info("Context packaged to %s", ctx.context_path)
 
   # Find requirements.txt
@@ -257,6 +308,8 @@ def _build_container(ctx: JobContext) -> None:
 
 def _upload_artifacts(ctx: JobContext) -> None:
   """Phase 3: Upload artifacts to Cloud Storage."""
+  if ctx.payload_path is None or ctx.context_path is None:
+    raise ValueError("payload_path and context_path must be set before upload")
   logging.info("Uploading artifacts to Cloud Storage (job: %s)...", ctx.job_id)
   storage.upload_artifacts(
     bucket_name=ctx.bucket_name,

keras_remote/core/core.py

@@ -9,6 +9,7 @@
 )
 from keras_remote.constants import DEFAULT_CLUSTER_NAME
 from keras_remote.core import accelerators
+from keras_remote.data import Data
 
 
 def run(
@@ -20,6 +21,7 @@ def run(
   cluster=None,
   backend=None,
   namespace="default",
+  volumes=None,
 ):
   """Execute function on remote TPU/GPU.
 
@@ -33,7 +35,25 @@ def run(
     cluster: GKE cluster name (default: from KERAS_REMOTE_CLUSTER)
     backend: Backend to use ('gke' or 'pathways')
     namespace: Kubernetes namespace (default: 'default')
+    volumes: Dict mapping absolute mount paths to Data objects, e.g.
+      ``{"/data": Data("./dataset/")}``. Data is downloaded to these
+      paths on the pod before function execution.
   """
+  # Validate volumes
+  if volumes is not None:
+    if not isinstance(volumes, dict):
+      raise TypeError(f"volumes must be a dict, got {type(volumes).__name__}")
+    for mount_path, data_obj in volumes.items():
+      if not isinstance(mount_path, str) or not mount_path.startswith("/"):
+        raise ValueError(
+          f"Volume mount path must be an absolute path "
+          f"(start with '/'), got: {mount_path!r}"
+        )
+      if not isinstance(data_obj, Data):
+        raise TypeError(
+          f"Volume value for {mount_path!r} must be a Data "
+          f"instance, got {type(data_obj).__name__}"
+        )
 
   def decorator(func):
     @functools.wraps(func)
@@ -79,6 +99,7 @@ def wrapper(*args, **kwargs):
           cluster,
           namespace,
           env_vars,
+          volumes,
         )
       elif resolved_backend == "pathways":
         return _execute_on_pathways(
@@ -92,6 +113,7 @@ def wrapper(*args, **kwargs):
           cluster,
           namespace,
           env_vars,
+          volumes,
         )
       else:
         raise ValueError(
@@ -114,6 +136,7 @@ def _execute_on_gke(
   cluster,
   namespace,
   env_vars,
+  volumes,
 ):
   """Execute function on GKE cluster with GPU/TPU nodes."""
   # Get GKE-specific defaults
@@ -123,7 +146,15 @@ def _execute_on_gke(
     namespace = os.environ.get("KERAS_REMOTE_GKE_NAMESPACE", "default")
 
   ctx = JobContext.from_params(
-    func, args, kwargs, accelerator, container_image, zone, project, env_vars
+    func,
+    args,
+    kwargs,
+    accelerator,
+    container_image,
+    zone,
+    project,
+    env_vars,
+    volumes=volumes,
   )
   return execute_remote(ctx, GKEBackend(cluster=cluster, namespace=namespace))
 
@@ -139,6 +170,7 @@ def _execute_on_pathways(
   cluster,
   namespace,
   env_vars,
+  volumes,
 ):
   """Execute function on GKE cluster via ML Pathways."""
   if not cluster:
@@ -147,7 +179,15 @@ def _execute_on_pathways(
     namespace = os.environ.get("KERAS_REMOTE_GKE_NAMESPACE", "default")
 
   ctx = JobContext.from_params(
-    func, args, kwargs, accelerator, container_image, zone, project, env_vars
+    func,
+    args,
+    kwargs,
+    accelerator,
+    container_image,
+    zone,
+    project,
+    env_vars,
+    volumes=volumes,
   )
   return execute_remote(
     ctx, PathwaysBackend(cluster=cluster, namespace=namespace)

keras_remote/core/core_test.py

@@ -33,7 +33,7 @@ def func():
 
       func()
       call_args = mock_exec.call_args
-      env_vars = call_args[0][-1]  # last positional arg
+      env_vars = call_args[0][-2]  # last positional arg
       self.assertEqual(env_vars, {"MY_VAR": "my_val"})
 
   def test_wildcard_pattern(self):
@@ -53,7 +53,7 @@ def func():
         pass
 
       func()
-      env_vars = mock_exec.call_args[0][-1]
+      env_vars = mock_exec.call_args[0][-2]
       self.assertIn("PREFIX_A", env_vars)
       self.assertIn("PREFIX_B", env_vars)
       self.assertNotIn("OTHER", env_vars)
@@ -70,7 +70,7 @@ def func():
         pass
 
       func()
-      env_vars = mock_exec.call_args[0][-1]
+      env_vars = mock_exec.call_args[0][-2]
       self.assertEqual(env_vars, {})
 
   def test_none_capture(self):
@@ -81,7 +81,7 @@ def func():
         pass
 
       func()
-      env_vars = mock_exec.call_args[0][-1]
+      env_vars = mock_exec.call_args[0][-2]
       self.assertEqual(env_vars, {})
 
   def test_mixed_exact_and_wildcard(self):
@@ -104,7 +104,7 @@ def func():
         pass
 
       func()
-      env_vars = mock_exec.call_args[0][-1]
+      env_vars = mock_exec.call_args[0][-2]
       self.assertEqual(
         env_vars, {"EXACT_VAR": "exact", "WILD_A": "a", "WILD_B": "b"}
       )

keras_remote/runner/remote_runner.py

@@ -18,6 +18,7 @@
 
 # Base temp directory for remote execution artifacts
 TEMP_DIR = tempfile.gettempdir()
+DATA_DIR = os.path.join(TEMP_DIR, "data")
 
 
 def main():
@@ -87,6 +88,12 @@ def run_gcs_mode():
       logging.info("Setting %d environment variables", len(env_vars))
       os.environ.update(env_vars)
 
+    # Resolve Data references
+    volumes = payload.get("volumes", [])
+    if volumes:
+      resolve_volumes(volumes, storage_client)
+    args, kwargs = resolve_data_refs(args, kwargs, storage_client)
+
     # Execute function and capture result
     logging.info("Executing %s()", func.__name__)
     result = None
@@ -129,6 +136,70 @@ def run_gcs_mode():
     sys.exit(1)
 
 
+def resolve_volumes(volume_refs, storage_client):
+  """Download volume data to their specified mount paths."""
+  for ref in volume_refs:
+    mount_path = ref["mount_path"]
+    logging.info("Resolving volume: %s -> %s", ref["gcs_uri"], mount_path)
+    _download_data(ref, mount_path, storage_client)
+
+
+def resolve_data_refs(args, kwargs, storage_client):
+  """Recursively resolve data ref dicts in args/kwargs to local paths."""
+  counter = 0
+
+  def _resolve(obj):
+    nonlocal counter
+    # Data ref that needs downloading (no mount_path means not volume-mounted)
+    if isinstance(obj, dict) and obj.get("__data_ref__"):
+      # Volume-mounted data refs are handled by Kubernetes, skip download
+      if obj.get("mount_path") is not None:
+        return obj["mount_path"]
+      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])
+      return local_dir
+    # Recurse into containers to find nested data refs
+    if isinstance(obj, dict):
+      return {k: _resolve(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+      return type(obj)(_resolve(item) for item in obj)
+    return obj
+
+  resolved_args = tuple(_resolve(a) for a in args)
+  resolved_kwargs = {k: _resolve(v) for k, v in kwargs.items()}
+  return resolved_args, resolved_kwargs
+
+
+def _download_data(ref, target_dir, storage_client):
+  """Download data from a GCS URI to a local directory."""
+  os.makedirs(target_dir, exist_ok=True)
+  gcs_uri = ref["gcs_uri"]
+
+  parts = gcs_uri.replace("gs://", "").split("/", 1)
+  bucket_name = parts[0]
+  prefix = parts[1].rstrip("/") if len(parts) > 1 else ""
+  bucket = storage_client.bucket(bucket_name)
+
+  blobs = bucket.list_blobs(prefix=prefix + "/")
+  count = 0
+  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)
+
+
 def _download_from_gcs(client, gcs_path, local_path):
   """Download file from GCS.