Address reviews

Author: JyotinderSingh

keras_remote/backend/execution.py

@@ -19,6 +19,7 @@
 from keras_remote.constants import get_default_zone, zone_to_region
 from keras_remote.credentials import ensure_credentials
 from keras_remote.infra import container_builder
+from keras_remote.infra.infra import get_default_project
 from keras_remote.utils import packager, storage
 
 
@@ -73,9 +74,7 @@ def from_params(
     if not zone:
       zone = get_default_zone()
     if not project:
-      project = os.environ.get("KERAS_REMOTE_PROJECT") or os.environ.get(
-        "GOOGLE_CLOUD_PROJECT"
-      )
+      project = get_default_project()
       if not project:
         raise ValueError(
           "project must be specified or set KERAS_REMOTE_PROJECT"

keras_remote/data.py

@@ -101,7 +101,9 @@ def __repr__(self):
     return f"Data({self._raw_path!r})"
 
 
-def _make_data_ref(gcs_uri, is_dir, mount_path=None):
+def _make_data_ref(
+  gcs_uri: str, is_dir: bool, mount_path: str | None = None
+) -> dict[str, object]:
   """Create a serializable data reference dict.
 
   These dicts replace Data objects in the payload before serialization.
@@ -115,6 +117,6 @@ def _make_data_ref(gcs_uri, is_dir, mount_path=None):
   }
 
 
-def is_data_ref(obj):
+def is_data_ref(obj: object) -> bool:
   """Check if an object is a serialized data reference."""
   return isinstance(obj, dict) and obj.get("__data_ref__") is True

keras_remote/infra/infra.py

@@ -5,5 +5,8 @@
 logger = logging.getLogger("keras_remote")
 
 
-def get_default_project():
-  return os.environ.get("KERAS_REMOTE_PROJECT")
+def get_default_project() -> str | None:
+  """Get project ID from KERAS_REMOTE_PROJECT or GOOGLE_CLOUD_PROJECT."""
+  return os.environ.get("KERAS_REMOTE_PROJECT") or os.environ.get(
+    "GOOGLE_CLOUD_PROJECT"
+  )

keras_remote/utils/storage.py

@@ -1,22 +1,24 @@
 """Cloud Storage operations for keras_remote."""
 
+from __future__ import annotations
+
 import os
 import tempfile
 
 from absl import logging
 from google.cloud import storage
 
-
-def _get_project():
-  """Get project ID from environment or default gcloud config."""
-  return os.environ.get("KERAS_REMOTE_PROJECT") or os.environ.get(
-    "GOOGLE_CLOUD_PROJECT"
-  )
+from keras_remote.data import Data
+from keras_remote.infra.infra import get_default_project
 
 
 def upload_artifacts(
-  bucket_name, job_id, payload_path, context_path, project=None
-):
+  bucket_name: str,
+  job_id: str,
+  payload_path: str,
+  context_path: str,
+  project: str | None = None,
+) -> None:
   """Upload execution artifacts to Cloud Storage.
 
   Args:
@@ -26,7 +28,7 @@ def upload_artifacts(
       context_path: Local path to context.zip
       project: GCP project ID (optional, uses env vars if not provided)
   """
-  project = project or _get_project()
+  project = project or get_default_project()
 
   client = storage.Client(project=project)
   bucket = client.bucket(bucket_name)
@@ -55,7 +57,9 @@ def upload_artifacts(
   )
 
 
-def download_result(bucket_name, job_id, project=None):
+def download_result(
+  bucket_name: str, job_id: str, project: str | None = None
+) -> str:
   """Download result from Cloud Storage.
 
   Args:
@@ -66,7 +70,7 @@ def download_result(bucket_name, job_id, project=None):
   Returns:
       Local path to downloaded result file
   """
-  project = project or _get_project()
+  project = project or get_default_project()
   client = storage.Client(project=project)
   bucket = client.bucket(bucket_name)
 
@@ -80,15 +84,17 @@ def download_result(bucket_name, job_id, project=None):
   return local_path
 
 
-def cleanup_artifacts(bucket_name, job_id, project=None):
+def cleanup_artifacts(
+  bucket_name: str, job_id: str, project: str | None = None
+) -> None:
   """Clean up job artifacts from Cloud Storage.
 
   Args:
       bucket_name: Name of the GCS bucket
       job_id: Unique job identifier
       project: GCP project ID (optional, uses env vars if not provided)
   """
-  project = project or _get_project()
+  project = project or get_default_project()
   client = storage.Client(project=project)
   bucket = client.bucket(bucket_name)
 
@@ -108,7 +114,12 @@ def cleanup_artifacts(bucket_name, job_id, project=None):
     )
 
 
-def upload_data(bucket_name, data, project=None, namespace_prefix="default"):
+def upload_data(
+  bucket_name: str,
+  data: Data,
+  project: str | None = None,
+  namespace_prefix: str = "default",
+) -> str:
   """Upload a Data object to GCS with content-based caching.
 
   For GCS Data: returns the original URI (no upload).
@@ -130,7 +141,7 @@ def upload_data(bucket_name, data, project=None, namespace_prefix="default"):
   content_hash = data.content_hash()
   cache_prefix = f"{namespace_prefix}/data-cache/{content_hash}"
 
-  project = project or _get_project()
+  project = project or get_default_project()
   client = storage.Client(project=project)
   bucket = client.bucket(bucket_name)
 
@@ -178,7 +189,7 @@ def upload_data(bucket_name, data, project=None, namespace_prefix="default"):
   return f"gs://{bucket_name}/{cache_prefix}"
 
 
-def _compute_total_size(path):
+def _compute_total_size(path: str) -> int:
   """Compute total size in bytes of a file or directory."""
   if os.path.isfile(path):
     return os.path.getsize(path)
@@ -189,7 +200,9 @@ def _compute_total_size(path):
   return total
 
 
-def _upload_directory(bucket, local_dir, gcs_prefix):
+def _upload_directory(
+  bucket: storage.Bucket, local_dir: str, gcs_prefix: str
+) -> None:
   """Upload a local directory to GCS preserving structure."""
   for root, _dirs, files in os.walk(local_dir):
     for fname in files:

keras_remote/utils/storage_test.py

@@ -9,9 +9,9 @@
 from absl.testing import absltest, parameterized
 
 from keras_remote.data import Data
+from keras_remote.infra.infra import get_default_project
 from keras_remote.utils.storage import (
   _compute_total_size,
-  _get_project,
   _upload_directory,
   cleanup_artifacts,
   download_result,
@@ -140,7 +140,7 @@ def test_resolves_project(self, kr_project, gc_project, expected):
     if gc_project:
       env["GOOGLE_CLOUD_PROJECT"] = gc_project
     with mock.patch.dict(os.environ, env, clear=True):
-      self.assertEqual(_get_project(), expected)
+      self.assertEqual(get_default_project(), expected)
 
 
 class TestUploadData(_GcsTestBase):