Ë
    g^(h-  ã                   óÖ   — d dl Z d dlZd dlmZmZ d dlmZ d dlmZ d dl	m
Z
mZmZ d dlmZ d dlmZ d dlmZmZmZ erd d	lmZ d
dgZ G d„ de«      Z G d„ d
e«      Z G d„ de«      Zy)é    N)Ú	GeneratorÚSequence)Úcontextmanager)ÚPath)ÚOptionalÚTYPE_CHECKINGÚUnion)Ú	url_to_fs)ÚStreamTransformExtension)ÚFileSystemBaseÚFileSystemReaderÚFileSystemWriter)ÚAbstractFileSystemÚFsspecWriterÚFsspecReaderc            
       ó  — e Zd Zdd„Zedeeej                  f   dede	e
j                  ddf   fd„«       Zdeeej                  f   dedeeej                  f   fd„Zdeeej                  f   deeej                  f   fd	„Zdeeej                  f   d
eeej                  f   ddfd„Zdeeej                  f   ddfd„Zedeeej                  f   defd„«       Zdeeej                  f   defd„Zdeeej                  f   ddfd„Zy)Ú
FileSystemÚreturnNc                 ó   — d | _         y ©N)Úfs)Úselfs    úm/var/www/skyplay_api_hub/venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_fsspec_filesystem.pyÚ__init__zFileSystem.__init__    s	   € Ø04ˆ�ó    ÚpathÚmodec              #   ó*  ‡K  — | j                   €J ‚t        j                  |«      }| j                   j                  |‰«      5 }	 |–— 	 d d d «       y #  t	        ˆfd„dD «       «      r	 | j                  |«       ‚ #  Y ‚ xY w‚ xY w# 1 sw Y   y xY w­w)Nc              3   ó&   •K  — | ]  }|‰v –— Œ
 y ­wr   © )Ú.0Úchr   s     €r   ú	<genexpr>z+FileSystem.create_stream.<locals>.<genexpr>1   s   øè ø€ Ò2 b�r˜T”zÑ2ùs   ƒzw+a)r   ÚosÚfspathÚopenÚanyÚrm_file)r   r   r   Ústreams     ` r   Úcreate_streamzFileSystem.create_stream#   s™   øè ø€ ð �w‰wÐ"Ð"Ð"Ü�y‰y˜‹ˆð
 �W‰W�\‰\˜$ Ó%ð 		¨ðØ“÷		ð 		øðÜÓ2¨EÔ2Ô2ðØŸ™ TÔ*ð øðØØú�ú÷		ð 		üsH   ƒA BÁBÁAÁ
	BÁBÁ*A<Á;BÁ<B Á>BÂBÂBÂBÚsuffixc                 óB   — t         j                  j                  ||«      S r   )r$   r   Újoin)r   r   r+   s      r   Úconcat_pathzFileSystem.concat_path8   s   € ô �w‰w�|‰|˜D &Ó)Ð)r   c                 ó.   — t        |fi |¤Ž\  | _        }|S r   )r
   r   )r   r   ÚkwargsÚ_s       r   Ú	init_pathzFileSystem.init_path=   s   € ô ˜tÑ. vÑ.‰
ˆŒ�Øˆr   Únew_pathc                 ó<   — | j                   j                  ||«       y r   )r   Úrename)r   r   r3   s      r   r5   zFileSystem.renameC   s   € ð 	�‰�‰�t˜XÕ&r   c                 ó>   — | j                   j                  |d¬«       y )NT)Úexist_ok)r   Úmakedirs©r   r   s     r   ÚmkdirzFileSystem.mkdirH   s   € Ø�‰×Ñ˜¨ÐÕ-r   Úcheckpoint_idc                 ó\   — t        |t        «      ry	 t        |«       y# t        $ r Y yw xY w)NFT)Ú
isinstancer   r
   Ú
ValueError©Úclsr;   s     r   Úvalidate_checkpoint_idz!FileSystem.validate_checkpoint_idK   s5   € ä�m¤TÔ*Øð	Ü�mÔ$ð øô ò 	Ùð	ús   “ Ÿ	+ª+c                 ó8   — | j                   j                  |«      S r   )r   Úexistsr9   s     r   rC   zFileSystem.existsW   s   € Ø�w‰w�~‰~˜dÓ#Ð#r   c                 ó:   — | j                   j                  |«       y r   )r   Úrmr9   s     r   r(   zFileSystem.rm_fileZ   s   € Ø�‰�
‰
�4Õr   )r   N)Ú__name__Ú
__module__Ú__qualname__r   r   r	   Ústrr$   ÚPathLiker   ÚioÚIOBaser*   r.   r2   r5   r:   ÚclassmethodÚboolrA   rC   r(   r    r   r   r   r      s‡  „ ó5ð ðØ˜#˜rŸ{™{Ð*Ñ+ðØ36ðà	�2—9‘9˜d DÐ(Ñ	)òó ðð(*Ø˜#˜rŸ{™{Ð*Ñ+ð*Ø58ð*à	ˆs�B—K‘KÐÑ	 ó*ð
Ø˜#˜rŸ{™{Ð*Ñ+ðà	ˆs�B—K‘KÐÑ	 óð'Ø˜#˜rŸ{™{Ð*Ñ+ð'Ø7<¸SÀ"Ç+Á+Ð=MÑ7Nð'à	ó'ð
.˜%  R§[¡[Ð 0Ñ1ð .°dó .ð ð	°5¸¸b¿k¹kÐ9IÑ3Jð 	Ètò 	ó ð	ð$˜5  b§k¡kÐ!1Ñ2ð $°tó $ð˜E # r§{¡{Ð"2Ñ3ð ¸ô r   r   c                   ó´   ‡ — e Zd ZdZ	 	 	 	 	 	 ddeeej                  f   dedede	de	ded	e
ee      d
dfˆ fd„Zedeeej                  f   d
efd„«       Zˆ xZS )r   a`  
    Basic implementation of StorageWriter using FFspec.

    This implementation makes the following assumptions and simplifications:

    * The checkpoint path is an empty or non-existing directory.
    * File creation is atomic

    The checkpoint consist of one file per write request plus
    a `.metadata` file with the serialized metadata.

    Nr   Úsingle_file_per_rankÚ
sync_filesÚthread_countÚper_thread_copy_aheadÚ	overwriteÚ_extensionsr   c           	      ó”   •— t         ‰	| �  |||||||¬«       t        «       | _         | j                  j                  |fi |¤Ž| _        y)a=  
        Initialize the writer pointing to `path`.

        Args:
            path: directory where the checkpoint will be written to.
            single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True.
            sync_files : force files to be synced to permanent storage. Default to True.
            thread_count: Number of IO threads to use to write. Default to 1.
            per_thread_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb.
            overwrite: Whether to allow overwriting existing checkpoints. Defaults to True.
            _extensions: Extensions to apply to output streams (EXPERIMENTAL)

        N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be consistent in the case of a failure.
        )rT   rU   N©Úsuperr   r   r   r2   r   )
r   r   rP   rQ   rR   rS   rT   rU   r0   Ú	__class__s
            €r   r   zFsspecWriter.__init__m   sT   ø€ ô2 	‰ÑØØ ØØØ!ØØ#ð 	ô 	
ô “,ˆŒØ%�D—G‘G×%Ñ% dÑ5¨fÑ5ˆ�	r   r;   c                 ó,   — t         j                  |«      S r   ©r   rA   r?   s     r   rA   z#FsspecWriter.validate_checkpoint_id’   ó   € ä×0Ñ0°Ó?Ð?r   )TTé   i€–˜ TN)rF   rG   rH   Ú__doc__r	   rI   r$   rJ   rN   Úintr   r   r   r   rM   rA   Ú__classcell__©rY   s   @r   r   r   _   s½   ø„ ñð  &*ØØØ%/ØØDHñ#6à�C˜Ÿ™Ð$Ñ%ð#6ð #ð#6ð ð	#6ð
 ð#6ð  #ð#6ð ð#6ð ˜hÐ'?Ñ@ÑAð#6ð 
õ#6ðJ ð@°5¸¸b¿k¹kÐ9IÑ3Jð @Ètò @ó ô@r   c                   ó~   ‡ — e Zd Zdeeej                  f   ddfˆ fd„Zedeeej                  f   de	fd„«       Z
ˆ xZS )r   r   r   Nc                 ó†   •— t         ‰| �  |«       t        «       | _         | j                  j                  |fi |¤Ž| _        y r   rW   )r   r   r0   rY   s      €r   r   zFsspecReader.__init__˜   s6   ø€ Ü‰Ñ˜ÔÜ“,ˆŒØ%�D—G‘G×%Ñ% dÑ5¨fÑ5ˆ�	r   r;   c                 ó,   — t         j                  |«      S r   r[   r?   s     r   rA   z#FsspecReader.validate_checkpoint_id�   r\   r   )rF   rG   rH   r	   rI   r$   rJ   r   rM   rN   rA   r`   ra   s   @r   r   r   —   sX   ø„ ð6˜U 3¨¯©Ð#3Ñ4ð 6À4õ 6ð
 ð@°5¸¸b¿k¹kÐ9IÑ3Jð @Ètò @ó ô@r   )rK   r$   Úcollections.abcr   r   Ú
contextlibr   Úpathlibr   Útypingr   r   r	   Úfsspec.corer
   Ú'torch.distributed.checkpoint._extensionr   Ú'torch.distributed.checkpoint.filesystemr   r   r   Úfsspecr   Ú__all__r   r   r   r    r   r   ú<module>rn      so   ðó 
Û 	ß /Ý %Ý ß 1Ñ 1å !å L÷ñ ñ Ý)ð Øð€ô<�ô <ô@5@Ð#ô 5@ôp@Ð#õ @r   