@@ -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