feat(async-job): write fatal error message to k8s termination message path

Signed-off-by: Jon Burdo <jon@jonburdo.com>

Author: jonburdo

jobs/async-upload/job/entrypoint.py

@@ -1,7 +1,6 @@
 import asyncio
 import logging
 import os
-import sys
 
 from job.upload import perform_upload
 from .config import get_config
@@ -13,16 +12,32 @@
 from .download import perform_download
 
 # Configure logging
-log_level = os.getenv('LOGLEVEL', logging.INFO)
+log_level = os.getenv("LOGLEVEL", logging.INFO)
 logging.basicConfig(
-    level=log_level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    force=True
+    level=log_level,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    force=True,
 )
 logger = logging.getLogger(__name__)
 
 # Test logging configuration immediately
 logger.info("📝 Logging configuration initialized successfully")
 
+termination_message_path = os.environ.get(
+    "TERMINATION_MESSAGE_PATH", "/dev/termination-log"
+)
+
+
+def write_to_termination_message_path(message):
+    with open(termination_message_path, "w") as f:
+        f.write(message)
+
+
+def record_error(exc):
+    message = f"Unexpected error: {exc}"
+    write_to_termination_message_path(message)
+    logger.error(message)
+
 
 async def main() -> None:
     """
@@ -32,7 +47,10 @@ async def main() -> None:
     logger.info("🚀 Starting async upload job...")
     try:
         # Get complete configuration
-        config = get_config()
+        try:
+            config = get_config()
+        except Exception as e:
+            raise RuntimeError("Failed to get config") from e
 
         client = validate_and_get_model_registry_client(config)
 
@@ -42,19 +60,16 @@ async def main() -> None:
         # Download the model from the defined source
         perform_download(client, config)
 
-
         # Upload the model to the destination
         uri = perform_upload(config)
 
         await update_model_artifact_uri(uri, client, config)
 
-    except ValueError as e:
-        logger.error(f"Configuration error: {str(e)}")
-        raise
-    except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
+    except BaseException as e:
+        record_error(e)
         raise
-    logger.info("🏁 Job completed successfully")
+    else:
+        logger.info("🏁 Job completed successfully")
 
 
 if __name__ == "__main__":  # pragma: no cover