Replaces print statements with logging (#28)

Author: JyotinderSingh

keras_remote/backend/execution.py

@@ -13,15 +13,14 @@
 
 import cloudpickle
 
+from absl import logging
+
 from keras_remote.constants import get_default_zone, zone_to_region
 from keras_remote.infra import container_builder
 from keras_remote.backend import gke_client
-from keras_remote.infra import infra
 from keras_remote.utils import packager
 from keras_remote.utils import storage
 
-logger = infra.logger
-
 
 @dataclass
 class JobContext:
@@ -155,7 +154,7 @@ def _prepare_artifacts(
     ctx: JobContext, tmpdir: str, caller_frame_depth: int = 3
 ) -> None:
     """Phase 1: Package function payload and working directory context."""
-    logger.info("Packaging function and context...")
+    logging.info("Packaging function and context...")
 
     # Get caller directory
     frame = inspect.stack()[caller_frame_depth]
@@ -170,28 +169,28 @@ def _prepare_artifacts(
     packager.save_payload(
         ctx.func, ctx.args, ctx.kwargs, ctx.env_vars, ctx.payload_path
     )
-    logger.info(f"Payload serialized to {ctx.payload_path}")
+    logging.info("Payload serialized to %s", ctx.payload_path)
 
     # Zip working directory
     ctx.context_path = os.path.join(tmpdir, "context.zip")
     packager.zip_working_dir(caller_path, ctx.context_path)
-    logger.info(f"Context packaged to {ctx.context_path}")
+    logging.info("Context packaged to %s", ctx.context_path)
 
     # Find requirements.txt
     ctx.requirements_path = _find_requirements(caller_path)
     if ctx.requirements_path:
-        logger.info(f"Found requirements.txt: {ctx.requirements_path}")
+        logging.info("Found requirements.txt: %s", ctx.requirements_path)
     else:
-        logger.info("No requirements.txt found")
+        logging.info("No requirements.txt found")
 
 
 def _build_container(ctx: JobContext) -> None:
     """Phase 2: Build or get cached container image."""
     if ctx.container_image:
         ctx.image_uri = ctx.container_image
-        logger.info(f"Using custom container: {ctx.image_uri}")
+        logging.info("Using custom container: %s", ctx.image_uri)
     else:
-        logger.info("Building container image...")
+        logging.info("Building container image...")
         ctx.image_uri = container_builder.get_or_build_container(
             base_image="python:3.12-slim",
             requirements_path=ctx.requirements_path,
@@ -203,7 +202,7 @@ def _build_container(ctx: JobContext) -> None:
 
 def _upload_artifacts(ctx: JobContext) -> None:
     """Phase 3: Upload artifacts to Cloud Storage."""
-    logger.info(f"Uploading artifacts to Cloud Storage (job: {ctx.job_id})...")
+    logging.info("Uploading artifacts to Cloud Storage (job: %s)...", ctx.job_id)
     storage.upload_artifacts(
         bucket_name=ctx.bucket_name,
         job_id=ctx.job_id,
@@ -215,7 +214,7 @@ def _upload_artifacts(ctx: JobContext) -> None:
 
 def _download_result(ctx: JobContext) -> dict:
     """Phase 6: Download and deserialize result from Cloud Storage."""
-    logger.info("Downloading result...")
+    logging.info("Downloading result...")
     result_path = storage.download_result(
         ctx.bucket_name, ctx.job_id, project=ctx.project
     )
@@ -226,14 +225,14 @@ def _download_result(ctx: JobContext) -> dict:
 
 def _cleanup_and_return(ctx: JobContext, result_payload: dict) -> Any:
     """Phase 7: Cleanup Cloud Storage artifacts and handle result."""
-    logger.info("Cleaning up artifacts...")
+    logging.info("Cleaning up artifacts...")
     storage.cleanup_artifacts(ctx.bucket_name, ctx.job_id, project=ctx.project)
 
     if result_payload["success"]:
-        logger.info("Remote execution completed successfully")
+        logging.info("Remote execution completed successfully")
         return result_payload["result"]
     else:
-        logger.error(f"Remote execution failed:\n{result_payload['traceback']}")
+        logging.error("Remote execution failed:\n%s", result_payload['traceback'])
         raise result_payload["exception"]
 
 
@@ -264,7 +263,7 @@ def execute_remote(ctx: JobContext, backend: BackendClient) -> Any:
         _upload_artifacts(ctx)
 
         # Phase 4: Submit job (backend-specific)
-        logger.info(f"Submitting job to {backend.__class__.__name__}...")
+        logging.info("Submitting job to %s...", backend.__class__.__name__)
         job = backend.submit_job(ctx)
 
         # Phase 5: Wait for completion (with cleanup on failure)

keras_remote/backend/gke_client.py

@@ -6,11 +6,10 @@
 from kubernetes import client, config
 from kubernetes.client.rest import ApiException
 
+from absl import logging
+
 from keras_remote.core.accelerators import TpuConfig
 from keras_remote.core import accelerators
-from keras_remote.infra import infra
-
-logger = infra.logger
 
 
 def submit_k8s_job(
@@ -58,10 +57,10 @@ def submit_k8s_job(
 
     try:
         created_job = batch_v1.create_namespaced_job(namespace=namespace, body=job)
-        logger.info(f"Submitted K8s job: {job_name}")
-        logger.info(f"View job with: kubectl get job {job_name} -n {namespace}")
-        logger.info(
-            f"View logs with: kubectl logs -l job-name={job_name} -n {namespace}"
+        logging.info("Submitted K8s job: %s", job_name)
+        logging.info("View job with: kubectl get job %s -n %s", job_name, namespace)
+        logging.info(
+            "View logs with: kubectl logs -l job-name=%s -n %s", job_name, namespace
         )
         return created_job
     except ApiException as e:
@@ -124,7 +123,7 @@ def wait_for_job(job, namespace="default", timeout=3600, poll_interval=10):
 
         # Check completion conditions
         if job_status.status.succeeded and job_status.status.succeeded >= 1:
-            print(f"[REMOTE] Job {job_name} completed successfully")
+            logging.info("Job %s completed successfully", job_name)
             return "success"
 
         if job_status.status.failed and job_status.status.failed >= 1:
@@ -137,7 +136,7 @@ def wait_for_job(job, namespace="default", timeout=3600, poll_interval=10):
 
         # Job still running
         if not logged_running:
-            logger.info(f"Job {job_name} running...")
+            logging.info("Job %s running...", job_name)
             logged_running = True
 
         time.sleep(poll_interval)
@@ -160,13 +159,13 @@ def cleanup_job(job_name, namespace="default"):
             namespace=namespace,
             body=client.V1DeleteOptions(propagation_policy="Foreground"),
         )
-        logger.info(f"Deleted K8s job: {job_name}")
+        logging.info("Deleted K8s job: %s", job_name)
     except ApiException as e:
         if e.status == 404:
             # Job already deleted
             pass
         else:
-            logger.warning(f"Failed to delete job {job_name}: {e.reason}")
+            logging.warning("Failed to delete job %s: %s", job_name, e.reason)
 
 
 def _parse_accelerator(accelerator):
@@ -339,7 +338,7 @@ def _print_pod_logs(core_v1, job_name, namespace):
                 logs = core_v1.read_namespaced_pod_log(
                     pod.metadata.name, namespace, tail_lines=100
                 )
-                print(f"[REMOTE] Pod {pod.metadata.name} logs:\n{logs}")
+                logging.info("Pod %s logs:\n%s", pod.metadata.name, logs)
 
 
 def _check_pod_scheduling(core_v1, job_name, namespace):

keras_remote/infra/container_builder.py

@@ -12,11 +12,10 @@
 from google.cloud import artifactregistry_v1
 from google.cloud import storage
 from google.cloud.devtools import cloudbuild_v1
+from absl import logging
+
 from keras_remote.constants import zone_to_ar_location, get_default_zone
 from keras_remote.core import accelerators
-from keras_remote.infra import infra
-
-logger = infra.logger
 
 REMOTE_RUNNER_FILE_NAME = "remote_runner.py"
 # Paths relative to this file's location (keras_remote/infra/)
@@ -54,19 +53,18 @@ def get_or_build_container(base_image, requirements_path, accelerator_type, proj
 
     # Check if image exists
     if _image_exists(image_uri, project):
-        logger.info(f"Using cached container: {image_uri}")
+        logging.info("Using cached container: %s", image_uri)
         ar_url = (
             "https://console.cloud.google.com/artifacts"
             f"/docker/{project}/{ar_location}"
             f"/keras-remote/base?project={project}"
         )
-        logger.info(f"View image: {ar_url}")
+        logging.info("View image: %s", ar_url)
         return image_uri
 
     # Build new image
-    logger.info(
-        f"Building new container (requirements changed): "
-        f"{image_uri}"
+    logging.info(
+        "Building new container (requirements changed): %s", image_uri
     )
     return _build_and_push(
         base_image, requirements_path, accelerator_type,
@@ -138,7 +136,7 @@ def _image_exists(image_uri, project):
     except google_exceptions.NotFound:
         return False
     except Exception:
-        logger.warning("Unexpected error checking image existence",
+        logging.warning("Unexpected error checking image existence",
                        exc_info=True)
         return False
 
@@ -212,26 +210,26 @@ def _build_and_push(base_image, requirements_path, accelerator_type,
             )
         )
 
-        logger.info("Submitting build to Cloud Build...")
+        logging.info("Submitting build to Cloud Build...")
         operation = build_client.create_build(project_id=project, build=build_config)
 
         # Get build ID from the operation metadata
         build_id = operation.metadata.build.id if hasattr(operation, "metadata") else None
         if build_id:
-            logger.info(f"Build ID: {build_id}")
-            logger.info(f"View build: https://console.cloud.google.com/cloud-build/builds/{build_id}?project={project}")
+            logging.info("Build ID: %s", build_id)
+            logging.info("View build: https://console.cloud.google.com/cloud-build/builds/%s?project=%s", build_id, project)
 
-        logger.info("Building container image (this may take 5-10 minutes)...")
+        logging.info("Building container image (this may take 5-10 minutes)...")
         result = operation.result(timeout=1200)  # 20 minute timeout
 
         if result.status == cloudbuild_v1.Build.Status.SUCCESS:
-            logger.info(f"Container built successfully: {image_uri}")
+            logging.info("Container built successfully: %s", image_uri)
             ar_url = (
                 "https://console.cloud.google.com/artifacts"
                 f"/docker/{project}/{ar_location}"
                 f"/keras-remote/base?project={project}"
             )
-            logger.info(f"View image: {ar_url}")
+            logging.info("View image: %s", ar_url)
             return image_uri
         else:
             raise RuntimeError(f"Build failed with status: {result.status}")
@@ -298,7 +296,7 @@ def _upload_build_source(tarball_path, bucket_name, project):
     blob.upload_from_filename(tarball_path)
 
     gcs_uri = f"gs://{bucket_name}/{blob_name}"
-    logger.info(f"Uploaded build source to {gcs_uri}")
-    logger.info(f"View source: https://console.cloud.google.com/storage/browser/{bucket_name}?project={project}")
+    logging.info("Uploaded build source to %s", gcs_uri)
+    logging.info("View source: https://console.cloud.google.com/storage/browser/%s?project=%s", bucket_name, project)
 
     return gcs_uri

keras_remote/runner/remote_runner.py

@@ -12,6 +12,7 @@
 import traceback
 import zipfile
 
+from absl import logging
 import cloudpickle
 from google.cloud import storage
 
@@ -26,7 +27,9 @@ def main():
     """
 
     if len(sys.argv) < 4:
-        print("Usage: remote_runner.py <context_gcs> <payload_gcs> <result_gcs>")
+        logging.error(
+            "Usage: remote_runner.py <context_gcs> <payload_gcs> <result_gcs>"
+        )
         sys.exit(1)
 
     run_gcs_mode()
@@ -44,7 +47,7 @@ def run_gcs_mode():
     payload_gcs = sys.argv[2]
     result_gcs = sys.argv[3]
 
-    print(f"[REMOTE] Starting GCS execution mode", flush=True)
+    logging.info("Starting GCS execution mode")
 
     # Define local paths using tempfile
     context_path = os.path.join(TEMP_DIR, "context.zip")
@@ -56,7 +59,7 @@ def run_gcs_mode():
         storage_client = storage.Client()
 
         # Download artifacts from Cloud Storage
-        print(f"[REMOTE] Downloading artifacts...", flush=True)
+        logging.info("Downloading artifacts...")
         _download_from_gcs(storage_client, context_gcs, context_path)
         _download_from_gcs(storage_client, payload_gcs, payload_path)
 
@@ -72,7 +75,7 @@ def run_gcs_mode():
         sys.path.insert(0, workspace_dir)
 
         # Load and deserialize payload
-        print("[REMOTE] Loading function payload", flush=True)
+        logging.info("Loading function payload")
         with open(payload_path, "rb") as f:
             payload = cloudpickle.load(f)
 
@@ -81,19 +84,19 @@ def run_gcs_mode():
         kwargs = payload["kwargs"]
         env_vars = payload.get("env_vars", {})
         if env_vars:
-            print(f"[REMOTE] Setting {len(env_vars)} environment variables", flush=True)
+            logging.info("Setting %d environment variables", len(env_vars))
             os.environ.update(env_vars)
 
         # Execute function and capture result
-        print(f"[REMOTE] Executing {func.__name__}()", flush=True)
+        logging.info("Executing %s()", func.__name__)
         result = None
         exception = None
 
         try:
             result = func(*args, **kwargs)
-            print(f"[REMOTE] Function completed successfully", flush=True)
+            logging.info("Function completed successfully")
         except BaseException as e:
-            print(f"[REMOTE] ERROR: {type(e).__name__}: {e}", flush=True)
+            logging.error("%s: %s", type(e).__name__, e)
             traceback.print_exc()
             sys.stdout.flush()
             sys.stderr.flush()
@@ -107,21 +110,21 @@ def run_gcs_mode():
             "success": exception is None,
             "result": result if exception is None else None,
             "exception": exception,
-            "traceback": traceback.format_exc() if exception else None
+            "traceback": traceback.format_exc() if exception else None,
         }
 
         with open(result_path, "wb") as f:
             cloudpickle.dump(result_payload, f)
 
         # Upload result to Cloud Storage
-        print(f"[REMOTE] Uploading result...", flush=True)
+        logging.info("Uploading result...")
         _upload_to_gcs(storage_client, result_path, result_gcs)
 
-        print("[REMOTE] Execution complete", flush=True)
+        logging.info("Execution complete")
         sys.exit(0 if exception is None else 1)
 
     except Exception as e:
-        print(f"[REMOTE] FATAL ERROR: {e}", flush=True)
+        logging.fatal("%s", e)
         traceback.print_exc()
         sys.exit(1)