From 7a0dc361c528976a285c7845905042320a8280cd Mon Sep 17 00:00:00 2001
From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com>
Date: Fri, 20 Jun 2025 06:20:52 -0400
Subject: [PATCH 1/2] fix: pythonkernel class trying to call non-existent
 self.post() method attr error

---
 airtbench/kernel.py | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/airtbench/kernel.py b/airtbench/kernel.py
index 43177cb..3a7bd7c 100644
--- a/airtbench/kernel.py
+++ b/airtbench/kernel.py
@@ -403,9 +403,14 @@ async def __aexit__(
             try:
                 logger.debug("Shutting down kernel")
                 try:
-                    await self._post(f"api/kernels/{self._kernel_id}/shutdown")
+                    async with aiohttp.ClientSession() as session:
+                        async with session.post(
+                            f"{self._base_url}/api/kernels/{self._kernel_id}/shutdown",
+                            params={"token": self._token},
+                        ) as response:
+                            response.raise_for_status()
                 except Exception as e:
-                    logger.warning(f"Failed to gracefully shutdown kernel via API:  {e}")
+                    logger.warning(f"Failed to gracefully shutdown kernel via API: {e}")
                 self._kernel_id = None
             except Exception:
                 logger.exception("Failed to shutdown kernel")
@@ -687,7 +692,7 @@ async def get_kernel_state(self) -> KernelState:
             response.raise_for_status()
             kernel_info = await response.json()
 
-        return t.cast(KernelState, kernel_info["execution_state"])
+        return t.cast("KernelState", kernel_info["execution_state"])
 
     async def busy(self) -> bool:
         """Check if the kernel is busy executing code."""
@@ -724,3 +729,23 @@ async def restart(self) -> None:
             response.raise_for_status()
 
         logger.debug(f"Kernel {self._kernel_id} restarted")
+
+
+async def cleanup_routine() -> None:
+    """Perform cleanup of Docker resources."""
+    try:
+        client = aiodocker.Docker()
+        # Clean up any dangling containers
+        containers = await client.containers.list(all=True)
+        for container in containers:
+            container_info = await container.show()
+            if container_info.get("State", {}).get("Status") == "exited":
+                try:
+                    await container.delete(force=True)
+                    logger.debug(f"Cleaned up exited container {container_info['Id'][:12]}")
+                except Exception as e:
+                    logger.debug(f"Could not clean up container: {e}")
+        await client.close()
+        logger.debug("Cleanup routine completed")
+    except Exception as e:
+        logger.warning(f"Cleanup routine failed: {e}")

From 3a8096ad1c1bc36f2b982336a15a5d98ca68a8ee Mon Sep 17 00:00:00 2001
From: Ads Dawson <104169244+GangGreenTemperTatum@users.noreply.github.com>
Date: Fri, 20 Jun 2025 06:28:05 -0400
Subject: [PATCH 2/2] fix: replace entire aexit method body with just await
 self.shutdown

---
 airtbench/kernel.py | 46 +--------------------------------------------
 1 file changed, 1 insertion(+), 45 deletions(-)

diff --git a/airtbench/kernel.py b/airtbench/kernel.py
index 3a7bd7c..ebd88fb 100644
--- a/airtbench/kernel.py
+++ b/airtbench/kernel.py
@@ -399,51 +399,7 @@ async def __aexit__(
         exc_val: BaseException | None,
         exc_tb: types.TracebackType | None,
     ) -> None:
-        if self._kernel_id is not None:
-            try:
-                logger.debug("Shutting down kernel")
-                try:
-                    async with aiohttp.ClientSession() as session:
-                        async with session.post(
-                            f"{self._base_url}/api/kernels/{self._kernel_id}/shutdown",
-                            params={"token": self._token},
-                        ) as response:
-                            response.raise_for_status()
-                except Exception as e:
-                    logger.warning(f"Failed to gracefully shutdown kernel via API: {e}")
-                self._kernel_id = None
-            except Exception:
-                logger.exception("Failed to shutdown kernel")
-
-        if self._container is not None and self.cleanup:
-            try:
-                logger.debug(f"Stopping container {self._container.id[:12]}")
-                await self._container.stop(t=5)
-
-                logger.debug(f"Removing container {self._container.id[:12]}")
-                await self._container.delete(force=self.force_remove)
-                self._container = None
-
-                # Log success
-                logger.debug("Container successfully removed")
-            except Exception as e:
-                logger.exception(f"Failed to stop/remove container: {e}")
-
-        if self._client is not None:
-            try:
-                logger.debug("Closing Docker client")
-                await self._client.close()
-                self._client = None
-            except Exception:
-                logger.exception("Failed to close Docker client")
-
-        if self._temp_dir is not None:
-            try:
-                logger.debug(f"Cleaning up temporary directory {self._temp_dir}")
-                shutil.rmtree(self._temp_dir)
-                self._temp_dir = None
-            except Exception:
-                logger.exception("Failed to clean up temporary directory")
+        await self.shutdown()
 
     async def get_container_logs(self) -> str:
         """Get the logs of the container."""