Skip to content

Commit 8f5d28e

Browse files
fix: pythonkernel class trying to call non-existent self.post() metho… (#31)
* fix: pythonkernel class trying to call non-existent self.post() method attr error * fix: replace entire aexit method body with just await self.shutdown
1 parent 9d13bac commit 8f5d28e

File tree

1 file changed

+22
-41
lines changed

1 file changed

+22
-41
lines changed

airtbench/kernel.py

Lines changed: 22 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -399,46 +399,7 @@ async def __aexit__(
399399
exc_val: BaseException | None,
400400
exc_tb: types.TracebackType | None,
401401
) -> None:
402-
if self._kernel_id is not None:
403-
try:
404-
logger.debug("Shutting down kernel")
405-
try:
406-
await self._post(f"api/kernels/{self._kernel_id}/shutdown")
407-
except Exception as e:
408-
logger.warning(f"Failed to gracefully shutdown kernel via API: {e}")
409-
self._kernel_id = None
410-
except Exception:
411-
logger.exception("Failed to shutdown kernel")
412-
413-
if self._container is not None and self.cleanup:
414-
try:
415-
logger.debug(f"Stopping container {self._container.id[:12]}")
416-
await self._container.stop(t=5)
417-
418-
logger.debug(f"Removing container {self._container.id[:12]}")
419-
await self._container.delete(force=self.force_remove)
420-
self._container = None
421-
422-
# Log success
423-
logger.debug("Container successfully removed")
424-
except Exception as e:
425-
logger.exception(f"Failed to stop/remove container: {e}")
426-
427-
if self._client is not None:
428-
try:
429-
logger.debug("Closing Docker client")
430-
await self._client.close()
431-
self._client = None
432-
except Exception:
433-
logger.exception("Failed to close Docker client")
434-
435-
if self._temp_dir is not None:
436-
try:
437-
logger.debug(f"Cleaning up temporary directory {self._temp_dir}")
438-
shutil.rmtree(self._temp_dir)
439-
self._temp_dir = None
440-
except Exception:
441-
logger.exception("Failed to clean up temporary directory")
402+
await self.shutdown()
442403

443404
async def get_container_logs(self) -> str:
444405
"""Get the logs of the container."""
@@ -687,7 +648,7 @@ async def get_kernel_state(self) -> KernelState:
687648
response.raise_for_status()
688649
kernel_info = await response.json()
689650

690-
return t.cast(KernelState, kernel_info["execution_state"])
651+
return t.cast("KernelState", kernel_info["execution_state"])
691652

692653
async def busy(self) -> bool:
693654
"""Check if the kernel is busy executing code."""
@@ -724,3 +685,23 @@ async def restart(self) -> None:
724685
response.raise_for_status()
725686

726687
logger.debug(f"Kernel {self._kernel_id} restarted")
688+
689+
690+
async def cleanup_routine() -> None:
691+
"""Perform cleanup of Docker resources."""
692+
try:
693+
client = aiodocker.Docker()
694+
# Clean up any dangling containers
695+
containers = await client.containers.list(all=True)
696+
for container in containers:
697+
container_info = await container.show()
698+
if container_info.get("State", {}).get("Status") == "exited":
699+
try:
700+
await container.delete(force=True)
701+
logger.debug(f"Cleaned up exited container {container_info['Id'][:12]}")
702+
except Exception as e:
703+
logger.debug(f"Could not clean up container: {e}")
704+
await client.close()
705+
logger.debug("Cleanup routine completed")
706+
except Exception as e:
707+
logger.warning(f"Cleanup routine failed: {e}")

0 commit comments

Comments
 (0)