Add more logging (#32)

Closes #31

sagemaker_shim/cli.py

@@ -44,6 +44,8 @@ def cli() -> None:
 
 @cli.command(short_help="Start the model server")
 def serve() -> None:
+    logger.info("Starting the model server")
+
     with AuxiliaryData():
         uvicorn.run(
             app=app,
@@ -56,6 +58,8 @@ def serve() -> None:
             loop="asyncio",
         )
 
+    logger.info("Model server stopped")
+
 
 @cli.command(short_help="Invoke the model")
 @click.option(
@@ -74,6 +78,8 @@ def serve() -> None:
 )
 @async_to_sync
 async def invoke(tasks: str, file: str) -> None:
+    logger.info("Invoking the model")
+
     tasks_json: str | bytes
 
     if tasks and file:
@@ -105,6 +111,8 @@ async def invoke(tasks: str, file: str) -> None:
     else:
         raise click.UsageError("Empty task list provided")
 
+    logger.info("Model invocation complete")
+
 
 def set_memory_limits() -> None:
     max_memory_mb = int(

sagemaker_shim/models.py

@@ -209,6 +209,8 @@ def chown(
         targetpath: StrOrBytesPath,
         numeric_owner: bool,
     ) -> None:
+        logger.info(f"chown of extracted {targetpath=}")
+
         if self.proc_user.uid is None and self.proc_user.gid is None:
             # No user or group specified, use the default
             return super().chown(
@@ -583,25 +585,23 @@ def proc_env(self) -> dict[str, str]:
 
     async def invoke(self) -> InferenceResult:
         """Run the inference on a single case"""
+        logger.info(f"Awaiting lock for {self.pk=}")
 
         await asyncio.wait_for(lock.acquire(), timeout=1.0)
 
+        logger.info(f"Invoking {self.pk=}")
+
         try:
             inference_result = await self._invoke()
-
-            logger.info(
-                f"Inference for {self.pk=} complete, {inference_result=}"
-            )
-
             self.upload_inference_result(inference_result=inference_result)
-
-            return inference_result
         finally:
             lock.release()
 
-    async def _invoke(self) -> InferenceResult:
-        logger.info(f"Invoking {self.pk=}")
+        logger.info(f"Invocation {self.pk=} complete")
+
+        return inference_result
 
+    async def _invoke(self) -> InferenceResult:
         try:
             self.reset_io()
 
@@ -700,8 +700,8 @@ def upload_inference_result(
         )
 
         logger.info(
-            f"Uploading Inference Result to "
-            f"{self.output_bucket_name=} with {bucket_key=}"
+            f"Uploading {bucket_key=} in "
+            f"{self.output_bucket_name=} with {inference_result=}"
         )
 
         self._s3_client.upload_fileobj(