o
    iÓ  ã                   @  s„   d dl mZ d dlmZ d dlmZ d dlmZmZm	Z	 e	dej
edƒ\ZZer/d dlmZ n
e	dej
edƒ\ZZG d	d
„ d
ƒZdS )é    )Úannotations)ÚTYPE_CHECKING)ÚSmartCacheDataset)Ú
IgniteInfoÚmin_versionÚoptional_importzignite.engineÚEvents)ÚEnginer	   c                   @  sB   e Zd ZdZddd„Zdd
d„Zddd„Zddd„Zddd„ZdS )ÚSmartCacheHandlerzš
    Attach SmartCache logic to the engine in Ignite.
    Mainly include the `start`, `update_cache`, and `shutdown` functions of SmartCacheDataset.

    Úsmartcacherr   ÚreturnÚNonec                 C  s   t |tƒs	tdƒ‚|| _dS )zÕ
        Args:
            smartcacher: predefined SmartCacheDataset, will attach it to the engine.

        Raises:
            TypeError: When ``smartcacher`` is not a ``monai.data.SmartCacheDataset``.

        z3smartcacher must be a monai.data.SmartCacheDataset.N)Ú
isinstancer   Ú	TypeErrorr   )Úselfr   © r   úc/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/handlers/smartcache_handler.pyÚ__init__!   s   
	
zSmartCacheHandler.__init__Úenginer	   c                 C  s4   |  tj| j¡ |  tj| j¡ |  tj| j¡ dS )zg
        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        N)Úadd_event_handlerr   ÚSTARTEDÚstartedÚEPOCH_COMPLETEDÚepoch_completedÚ	COMPLETEDÚ	completed©r   r   r   r   r   Úattach.   s   zSmartCacheHandler.attachc                 C  ó   | j  ¡  dS )zÝCallback for train or validation/evaluation started Event.
        Start the replacement thread of SmartCacheDataset.

        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        N)r   Ústartr   r   r   r   r   7   ó   zSmartCacheHandler.startedc                 C  r   )zÞCallback for train or validation/evaluation epoch completed Event.
        Update cache content with replacement data.

        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        N)r   Úupdate_cacher   r   r   r   r   @   r    z!SmartCacheHandler.epoch_completedc                 C  r   )zÞCallback for train or validation/evaluation completed Event.
        Stop the replacement thread of SmartCacheDataset.

        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        N)r   Úshutdownr   r   r   r   r   I   r    zSmartCacheHandler.completedN)r   r   r   r   )r   r	   r   r   )	Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r   r   r   r   r   r   r   r   r
      s    


	
		r
   N)Ú
__future__r   Útypingr   Ú
monai.datar   Úmonai.utilsr   r   r   ÚOPT_IMPORT_VERSIONr   Ú_Úignite.enginer	   r
   r   r   r   r   Ú<module>   s   