o
    Ti                     @   s   d dl Z G dd deZdS )    Nc                   @   sN   e Zd ZdddZdd ZdddZd	efd
dZdd	efddZdd Z	dS )CheckpointEngineNc                 C      d S N )selfconfig_paramsr   r   i/home/ubuntu/.local/lib/python3.10/site-packages/deepspeed/runtime/checkpoint_engine/checkpoint_engine.py__init__      zCheckpointEngine.__init__c                 C   r   r   r   r   tagr   r   r   create      zCheckpointEngine.createFc                 C   s   t j||d d S )N)exist_ok)osmakedirs)r   pathr   r   r   r   r      s   zCheckpointEngine.makedirsr   c                 C   r   r   r   )r   
state_dictr   r   r   r   save   r
   zCheckpointEngine.savec                 C   r   r   r   )r   r   map_locationr   r   r   load   r
   zCheckpointEngine.loadc                 C   r   r   r   r   r   r   r   commit   r   zCheckpointEngine.commitr   )F)
__name__
__module____qualname__r	   r   r   strr   r   r   r   r   r   r   r   	   s    

r   )r   objectr   r   r   r   r   <module>   s   