Addressed reviews

Author: JyotinderSingh

keras_remote/utils/packager.py

@@ -1,14 +1,36 @@
+"""Packaging utilities for serializing functions, args, and working directories.
+
+Handles zipping the user's working directory, serializing the function
+payload with cloudpickle, and extracting/replacing Data objects in
+arbitrarily nested arg structures.
+"""
+
 import os
 import zipfile
+from collections.abc import Callable
+from typing import Any
 
 import cloudpickle
 
 from keras_remote.data import Data
 
+# Type alias for a position path through nested args, e.g. ("arg", 0, "key").
+PositionPath = tuple[str | int, ...]
+
+
+def zip_working_dir(
+  base_dir: str, output_path: str, exclude_paths: set[str] | None = None
+) -> None:
+  """Zip a directory into a ZIP archive, excluding common non-source files.
 
-def zip_working_dir(base_dir, output_path, exclude_paths=None):
-  """Zips the base_dir into output_path, excluding .git, __pycache__,
-  and any paths in exclude_paths."""
+  Excludes ``.git``, ``__pycache__``, and any paths in *exclude_paths*
+  (which may be files or directories).
+
+  Args:
+      base_dir: Root directory to zip.
+      output_path: Destination path for the ZIP file.
+      exclude_paths: Absolute paths to skip during archiving.
+  """
   exclude_paths = exclude_paths or set()
   normalized_excludes = {os.path.normpath(p) for p in exclude_paths}
 
@@ -30,10 +52,28 @@ def zip_working_dir(base_dir, output_path, exclude_paths=None):
         zipf.write(file_path, archive_name)
 
 
-def save_payload(func, args, kwargs, env_vars, output_path, volumes=None):
-  """Uses cloudpickle to serialize the function, args, kwargs, and
-  env_vars."""
-  payload = {
+def save_payload(
+  func: Callable,
+  args: tuple,
+  kwargs: dict[str, Any],
+  env_vars: dict[str, str],
+  output_path: str,
+  volumes: list[dict[str, Any]] | None = None,
+) -> None:
+  """Serialize a function call payload with cloudpickle.
+
+  The resulting pickle file contains a dict with keys ``func``, ``args``,
+  ``kwargs``, ``env_vars``, and optionally ``volumes``.
+
+  Args:
+      func: The user function to execute remotely.
+      args: Positional arguments (Data objects should already be replaced).
+      kwargs: Keyword arguments.
+      env_vars: Environment variables to set on the remote pod.
+      output_path: Destination path for the pickle file.
+      volumes: Optional list of volume data-ref dicts.
+  """
+  payload: dict[str, Any] = {
     "func": func,
     "args": args,
     "kwargs": kwargs,
@@ -45,50 +85,89 @@ def save_payload(func, args, kwargs, env_vars, output_path, volumes=None):
     cloudpickle.dump(payload, f)
 
 
-def extract_data_refs(args, kwargs):
+def extract_data_refs(
+  args: tuple, kwargs: dict[str, Any]
+) -> list[tuple[Data, PositionPath]]:
   """Scan args and kwargs for Data objects at any nesting depth.
 
-  Returns list of (data_obj, position_path) tuples.
+  Returns a list of ``(data_obj, position_path)`` tuples. The position
+  path encodes where each Data object was found, e.g.
+  ``("arg", 0)`` or ``("kwarg", "config", "data")``.
+
+  Circular references are handled safely via an ``id()``-based visited
+  set.
   """
-  refs = []
+  refs: list[tuple[Data, PositionPath]] = []
   for i, arg in enumerate(args):
     _scan_for_data(arg, ("arg", i), refs)
   for key, val in kwargs.items():
     _scan_for_data(val, ("kwarg", key), refs)
   return refs
 
 
-def _scan_for_data(obj, path, refs):
+def _scan_for_data(
+  obj: Any,
+  path: PositionPath,
+  refs: list[tuple[Data, PositionPath]],
+  visited: set[int] | None = None,
+) -> None:
+  """Recursively collect Data objects from a nested structure."""
+  if visited is None:
+    visited = set()
+  obj_id = id(obj)
+  if obj_id in visited:
+    return
+  visited.add(obj_id)
   if isinstance(obj, Data):
     refs.append((obj, path))
-  elif isinstance(obj, (list, tuple)):
+  elif isinstance(obj, (list, tuple, set, frozenset)):
     for i, item in enumerate(obj):
-      _scan_for_data(item, path + (i,), refs)
+      _scan_for_data(item, path + (i,), refs, visited)
   elif isinstance(obj, dict):
     for key, val in obj.items():
-      _scan_for_data(val, path + (key,), refs)
+      _scan_for_data(val, path + (key,), refs, visited)
 
 
-def replace_data_with_refs(args, kwargs, ref_map):
-  """Replace Data objects with serializable ref dicts.
+def replace_data_with_refs(
+  args: tuple,
+  kwargs: dict[str, Any],
+  ref_map: dict[int, dict[str, Any]],
+) -> tuple[tuple, dict[str, Any]]:
+  """Replace Data objects in args/kwargs with serializable ref dicts.
 
   Args:
-      ref_map: dict mapping id(Data) -> ref dict
+      args: Positional arguments, possibly containing Data objects.
+      kwargs: Keyword arguments, possibly containing Data objects.
+      ref_map: Mapping from ``id(Data)`` to the replacement ref dict.
+
   Returns:
-      (new_args, new_kwargs) -- new tuples/dicts with Data replaced
+      ``(new_args, new_kwargs)`` with all matched Data objects replaced.
   """
   new_args = tuple(_replace_in_value(a, ref_map) for a in args)
   new_kwargs = {k: _replace_in_value(v, ref_map) for k, v in kwargs.items()}
   return new_args, new_kwargs
 
 
-def _replace_in_value(obj, ref_map):
-  if isinstance(obj, Data) and id(obj) in ref_map:
-    return ref_map[id(obj)]
+def _replace_in_value(
+  obj: Any,
+  ref_map: dict[int, dict[str, Any]],
+  visited: set[int] | None = None,
+) -> Any:
+  """Recursively replace Data objects with their ref dicts."""
+  if visited is None:
+    visited = set()
+  obj_id = id(obj)
+  if obj_id in visited:
+    return obj
+  visited.add(obj_id)
+  if isinstance(obj, Data) and obj_id in ref_map:
+    return ref_map[obj_id]
   elif isinstance(obj, list):
-    return [_replace_in_value(item, ref_map) for item in obj]
+    return [_replace_in_value(item, ref_map, visited) for item in obj]
   elif isinstance(obj, tuple):
-    return tuple(_replace_in_value(item, ref_map) for item in obj)
+    return tuple(_replace_in_value(item, ref_map, visited) for item in obj)
+  elif isinstance(obj, (set, frozenset)):
+    return [_replace_in_value(item, ref_map, visited) for item in obj]
   elif isinstance(obj, dict):
-    return {k: _replace_in_value(v, ref_map) for k, v in obj.items()}
+    return {k: _replace_in_value(v, ref_map, visited) for k, v in obj.items()}
   return obj

keras_remote/utils/packager_test.py

@@ -312,6 +312,34 @@ def test_no_data_objects(self):
     refs = extract_data_refs((1, "hello"), {"lr": 0.01})
     self.assertEqual(len(refs), 0)
 
+  def test_nested_in_set(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("data")
+    d = Data(str(f))
+
+    refs = extract_data_refs(({d, "other"},), {})
+    self.assertEqual(len(refs), 1)
+    self.assertIs(refs[0][0], d)
+
+  def test_nested_in_frozenset(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("data")
+    d = Data(str(f))
+
+    refs = extract_data_refs((frozenset({d}),), {})
+    self.assertEqual(len(refs), 1)
+    self.assertIs(refs[0][0], d)
+
+  def test_circular_reference_does_not_recurse(self):
+    """Circular structures in args should not cause infinite recursion."""
+    circular = {"key": "value"}
+    circular["self"] = circular
+
+    refs = extract_data_refs((circular,), {})
+    self.assertEqual(len(refs), 0)
+
 
 class TestReplaceDataWithRefs(absltest.TestCase):
   def test_replaces_direct_arg(self):
@@ -357,6 +385,38 @@ def test_preserves_non_data(self):
     self.assertEqual(new_args, (1, "hello", [1, 2]))
     self.assertEqual(new_kwargs, {"x": 3})
 
+  def test_replaces_in_set(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("data")
+    d = Data(str(f))
+    ref = {"__data_ref__": True, "gcs_uri": "gs://b/p"}
+    ref_map = {id(d): ref}
+
+    new_args, _ = replace_data_with_refs(({d},), {}, ref_map)
+    self.assertIsInstance(new_args[0], list)
+    self.assertIn(ref, new_args[0])
+
+  def test_replaces_in_frozenset(self):
+    tmp = _make_temp_path(self)
+    f = tmp / "data.csv"
+    f.write_text("data")
+    d = Data(str(f))
+    ref = {"__data_ref__": True, "gcs_uri": "gs://b/p"}
+    ref_map = {id(d): ref}
+
+    new_args, _ = replace_data_with_refs((frozenset({d}),), {}, ref_map)
+    self.assertIsInstance(new_args[0], list)
+    self.assertIn(ref, new_args[0])
+
+  def test_circular_reference_does_not_recurse(self):
+    """Circular structures should not cause infinite recursion."""
+    circular = [1, 2]
+    circular.append(circular)
+
+    new_args, _ = replace_data_with_refs((circular,), {}, {})
+    self.assertIsInstance(new_args[0], list)
+
 
 if __name__ == "__main__":
   absltest.main()