a
    d!                     @   sX  d dl mZ d dlmZ d dlmZmZ d dlmZm	Z	 d dl
mZmZmZmZmZmZmZmZmZmZ d dlmZ d dlmZmZmZmZmZmZ d dlmZm Z m!Z!m"Z"m#Z# d d	l$m%Z% d d
l&Z'd dl(m)Z) d dl*m+Z+ d dl,m-Z-m.Z. d dl/m0Z0 d dl1m2Z2m3Z3 G dd de%Z4G dd de%Z5G dd deZ6edddG dd dZ7eG dd dZ8G dd deZ9G dd de#Z:e"ee6 d d!d"Z;e"ee<ef ed#d$d%Z=eeee-ed&d'd(Z>eee"e?d)ed*d+d,Z@d[e"ed)e?d
d-d.d/ZAe"d
d d0d1ZBe"ee<ef d
d#d2d3ZCe:ee<ef d
d4d5d6ZDe"ee<ef d
d#d7d8ZEe"ee<ef d
d#d9d:ZFe"ee<ef d
d#d;d<ZGee?ef e?e?ee?ef d=d>d?ZHG d@dA dAeZIG dBdC dCeIe!ZJG dDdE dEeIe ZKe"dFdGdHdIZLd
dJdKdLZMd
dJdMdNZNe"d
d dOdPZOe"d
d dQdRZPedSd
dTdUdVZQd\ee<edXdYdZZRd
S )]    )Sized)deepcopy)	dataclassfield)partialwraps)
AnyCallableDict	GeneratorIterableIteratorListOptionalTupleUnion)apply_to_collection)DatasetDistributedSamplerget_worker_infoRandomSamplerSamplerSequentialSampler)_BaseDataLoaderIter_MultiProcessingDataLoaderIter_SingleProcessDataLoaderIter
DataLoaderIterableDataset)	TypedDictN)	_Stateful)_collect_states_on_rank_zero)_FaultTolerantModeAutoRestartBatchKeys)MisconfigurationException)_collect_rng_states_set_rng_statesc                   @   sR   e Zd ZU eeef ed< eeef ed< eed< eed< eed< ee ed< dS )_IteratorStateDictdataset_statesampler_state	worker_idnum_workersnum_batches_fetchednameN)	__name__
__module____qualname__r
   intr   __annotations__r   str r3   r3   q/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/pytorch_lightning/utilities/auto_restart.pyr&   '   s   
r&   c                   @   s2   e Zd ZU eeef ed< eed< ee	 ed< dS )_MergedIteratorStateDictstatelatest_worker_idrepresent_map_datasetN)
r-   r.   r/   r
   r2   r   r1   r0   r   boolr3   r3   r3   r4   r5   0   s   
r5   c                       s  e Zd ZdZdeeef ee dd fddZ	ee
dddZd ee dd	d
dZeedddZee
 dddZe
dddZedddZd!ee eeeeef f dddZeee
f ddddZd"ee edddZeeeee
f f ddddZ  ZS )#FastForwardSamplera2  This FastForwardSampler wraps a :class:`torch.utils.data.Sampler` and records the number of iterations
    performed during an epoch.

    It maintains a state, saved with :meth:`state_dict`, that can be reloaded with
    :meth:`load_state_dict`. If the sampler is used in a multiprocessing context, the ``FastForwardSampler`` will record
    the state of the current worker.
    When reloading, the ``FastForwardSampler`` will "fast-forward" the wrapped sampler by iterating through all the
    samples seen in the last iterations (for the current worker).
    N)sampler	attr_namereturnc                    s<   t  jd d || _d| _d| _d| _d | _d | _|| _d S )N)Zdata_sourceFr   )	super__init___sampler
restarting_current_iteration_counter_dataloader_batch_size_cached_state_dictZ
_attr_name)selfr;   r<   	__class__r3   r4   r?   A   s    zFastForwardSampler.__init__)keyr=   c                 C   s"   || j v r| j | S t| j|d S N)__dict__getattrr@   )rF   rI   r3   r3   r4   __getattr__K   s    

zFastForwardSampler.__getattr__)dataloader_batch_sizer=   c                 C   s
   || _ dS )zSetup the ``FastForwardSampler``.

        This is required only when the provided dataset subclassed
        :class:`torch.utils.data.Dataset`.
        N)rD   )rF   rN   r3   r3   r4   setupP   s    zFastForwardSampler.setupr=   c                 C   s   t  }|r|jS dS Nr   r   idrF   Zworker_infor3   r3   r4   r)   X   s    zFastForwardSampler.worker_idc                 C   s   t | j| _d| _d| _| S rQ   )iterr@   sampler_iterrB   rC   rF   r3   r3   r4   __iter__]   s    zFastForwardSampler.__iter__c                 C   s   | j d ur| | j  | j| jk r<t| j |  jd7  _q| j d urLd | _ |  jd7  _|  jd7  _d}zt| jW S  ty   d}Y n0 d| _d| _d | _ d| _|rtd S )N   FTr   )rE   _load_non_random_staterC   rB   nextrV   StopIterationrA   )rF   Z
has_raisedr3   r3   r4   __next__c   s(    



zFastForwardSampler.__next__c                 C   s   t | jtsJ t| jS rJ   )
isinstancer@   r   lenrW   r3   r3   r4   __len__   s    zFastForwardSampler.__len__)num_batches_processedr=   c                 C   s   | j d| |iiS )zoReturns the state of the sampler in the current worker.

        The worker id indexes the state dict.
        current_iteration)r)   _compute_current_iteration)rF   ra   r3   r3   r4   
state_dict   s    zFastForwardSampler.state_dictrd   r=   c                 C   s   t |}|| _d| _dS )a  Loads the saved state for the wrapped sampler.

        If the ``state_dict`` contains multiple states, it means there were multiple workers. The state will be cached
        and fully reloaded (fast-forward) the first time :meth:`__iter__` is called.
        TN)r   rE   rA   rF   rd   r3   r3   r4   load_state_dict   s    z"FastForwardSampler.load_state_dictc                 C   s0   |dur|}n| j }| jr,|dur,|| j9 }|S )a  This function is used to compute the effective iteration.

        As DataLoader can perform ``prefecthing`` or training can fail while processing a batch, the current iteration
        needs to be computed using the ``num_batches_processed`` processed information.
        N)rB   rD   )rF   ra   rb   r3   r3   r4   rc      s    
z-FastForwardSampler._compute_current_iterationc                 C   s   || j  d | _d S )Nrb   )r)   rB   rf   r3   r3   r4   rZ      s    z)FastForwardSampler._load_non_random_state)N)N)N)N)r-   r.   r/   __doc__r   r   r   r   r2   r?   r   rM   r0   rO   propertyr)   r   rX   r]   r`   r
   rd   rg   rc   rZ   __classcell__r3   r3   rG   r4   r:   6   s   
$
&r:   T)frozenZunsafe_hashc                   @   s   e Zd ZU dZeedZeee	f e
d< eedZeee	f e
d< dZee
d< dZee
d< dZee
d< d	Zee e
d
< eed dddZd	S )IteratorStatez4The state of an iterator in a single worker process.default_factoryr'   r(   r   r)   r*   r+   Nr,   re   c                 C   s   | f i |S rJ   r3   clsrd   r3   r3   r4   from_state_dict   s    zIteratorState.from_state_dict)r-   r.   r/   rh   r   dictr'   r
   r0   r   r1   r(   r)   r*   r+   r,   r   r2   classmethodr&   rq   r3   r3   r3   r4   rl      s   
rl   c                   @   s   e Zd ZU dZeedZeed< dZ	e
ed< dZee ed< ee eddd	d
Zeee
ef dddZeee
ef dddZeed dddZe
dddZdS )MergedIteratorStatea   This class is used to hold the current iterator state and lives on the iterator.

    It holds the current merged states from all worker processes. Once an iterator advances, it can store updates of the
    worker states in this merged iterator state.
    rm   r6   r   r7   Nr8   )generator_name	new_stater=   c                 C   sT   |d u | _ |j}|d u r$|| j|< n&|| jvr8i | j|< | j| }|||< || _d S rJ   )r8   r)   r6   r7   )rF   ru   rv   r7   r6   r3   r3   r4   update   s    



zMergedIteratorState.updaterP   c                    s    fdd j  D S )z;Returns the merged sampler states for all worker processes.c                    s   i | ]}d  j | jd  qS )r   )r6   r(   .0krW   r3   r4   
<dictcomp>       z6MergedIteratorState.sampler_states.<locals>.<dictcomp>r6   keysrW   r3   rW   r4   sampler_states   s    z"MergedIteratorState.sampler_statesc                    s    fdd j  D S )z;Returns the merged dataset states for all worker processes.c                    s   i | ]}| j | j| qS r3   r6   r'   rx   rW   r3   r4   r{      r|   z6MergedIteratorState.dataset_states.<locals>.<dictcomp>r}   rW   r3   rW   r4   dataset_states   s    z"MergedIteratorState.dataset_statesre   c                 C   sL   |d r$dd |d   D |d< ndd |d   D |d< | f i |S )Nr8   c                 S   s   i | ]\}}|t |qS r3   rl   rq   ry   r)   r6   r3   r3   r4   r{      s   z7MergedIteratorState.from_state_dict.<locals>.<dictcomp>r6   c                 S   s$   i | ]\}}|d d |  D qS )c                 S   s   i | ]\}}|t |qS r3   r   r   r3   r3   r4   r{      s   zBMergedIteratorState.from_state_dict.<locals>.<dictcomp>.<dictcomp>items)ry   sampler_nameZworker_stater3   r3   r4   r{      s   r   ro   r3   r3   r4   rq      s    


z#MergedIteratorState.from_state_dictc                 C   s
   t | jS rJ   )r_   r6   rW   r3   r3   r4   r`      s    zMergedIteratorState.__len__)r-   r.   r/   rh   r   rr   r6   r
   r1   r7   r0   r8   r   r9   r2   rl   rw   ri   r   r   r   rs   r5   rq   r`   r3   r3   r3   r4   rt      s   
rt   c                   @   s   e Zd ZdZeddddZeedddZee	e
eeef f d	d
dZedddZeee
f eeddddZeeeee
f f dddZdS )CaptureMapDataseta  This class is used to capture the state from the map-based state dataset.

    Note:
        We currently don't support restoring if we fail during the first `N = num_workers` batches, where
        `num_workers` is the number of workers spawned by the dataloader.
    Ndatasetr=   c                 C   s   || _ d | _d S rJ   )r   rE   rF   r   r3   r3   r4   r?      s    zCaptureMapDataset.__init__rP   c                 C   s   t  }|r|jS dS rQ   rR   rT   r3   r3   r4   r)      s    zCaptureMapDataset.worker_id)itemr=   c                 C   s:   | j d ur0| j| j v r*t| j | j d  d | _ | j| S NZ
rng_states)rE   r)   r%   r   )rF   r   r3   r3   r4   __getitem__  s
    
zCaptureMapDataset.__getitem__c                 C   s   t | jtsJ t| jS rJ   )r^   r   r   r_   rW   r3   r3   r4   r`   
  s    zCaptureMapDataset.__len__)rd   r7   r*   r=   c                 C   s   t t|||| _d S rJ   )_rotate_worker_indicesr   rE   )rF   rd   r7   r*   r3   r3   r4   rg     s    z!CaptureMapDataset.load_state_dictc                 C   s   | j dt iiS r   )r)   r$   rW   r3   r3   r4   rd     s    zCaptureMapDataset.state_dict)r-   r.   r/   rh   r   r?   ri   r0   r)   r   r   r
   r   r`   rg   r2   rd   r3   r3   r3   r4   r      s    r   c                       s   e Zd ZdZedd fddZeedddZe	e
ef dd	d
Ze	e
ef ddddZddddZddddZedddZe	e
ef dddZ  ZS )CaptureIterableDataseta  The ``CaptureIterableDataset`` is used to wrap an :class:`torch.utils.data.IterableDataset`.

    On ``__iter__`` function call,   the ``CaptureIterableDataset`` will wrap the wrapped dataset     generators into
    ``FastForwardSampler`` to keep track of progress. On ``__next__`` function call, the ``CaptureIterableDataset`` will
    return a dictionary containing     user data and metadata containing the ``FastForwardSampler`` samplers state_dict.
    Nr   c                    s*   t    t|| _d | _d | _d| _d S )NF)r>   r?   r   r   samplers_state_dictZ_has_wrappedr   rG   r3   r4   r?     s
    

zCaptureIterableDataset.__init__rP   c                 C   s   | j jS rJ   )r   r;   rW   r3   r3   r4   r;   %  s    zCaptureIterableDataset.samplerc                 C   s"   | j d usJ dd | j  D S )Nc                 S   s   i | ]\}}||  qS r3   rd   ry   rz   vr3   r3   r4   r{   +  r|   z5CaptureIterableDataset.state_dict.<locals>.<dictcomp>)r   r   rW   r3   r3   r4   rd   )  s    z!CaptureIterableDataset.state_dictre   c                 C   s   t || _d S rJ   )r   r   rf   r3   r3   r4   rg   -  s    z&CaptureIterableDataset.load_state_dictc                 C   s   i | _ | jj}dd | D }| D ]P\}}t|tr<q(t||d}| jd urb|| j|  || j |< t	|||< q(| 
  d S )Nc                 S   s$   i | ]\}}t |ttfr||qS r3   )r^   r   r   r   r3   r3   r4   r{   7  r|   zCCaptureIterableDataset._wrap_generator_samplers.<locals>.<dictcomp>)r<   )r   r   rK   r   r^   r   r:   r   rg   rU   reset_on_epoch)rF   Zdataset_dictZdataset_sampler_generatorsZgenerator_attr_name	generatorr;   r3   r3   r4   _wrap_generator_samplers0  s    


z/CaptureIterableDataset._wrap_generator_samplersc                 C   s
   d | _ d S rJ   )r   rW   r3   r3   r4   r   N  s    z%CaptureIterableDataset.reset_on_epochc                 C   s,   t | j| _t| jtr td|   | S )NzPyTorch Lightning Fault-Tolerant feature does not support `__iter__` returning a generator. Please use the `__next__` function to fetch the next batch and use a sampler for doing your iterations.)rU   r   	iter_datar^   r   r#   r   rW   r3   r3   r4   rX   Q  s    zCaptureIterableDataset.__iter__c                 C   s
   t | jS rJ   )r[   r   rW   r3   r3   r4   r]   `  s    zCaptureIterableDataset.__next__)r-   r.   r/   rh   r   r?   ri   r   r;   r
   r2   r   rd   rg   r   r   r   rX   r]   rj   r3   r3   rG   r4   r     s   r   
dataloaderr=   c                 C   s(   t | jtr| jS t | jtr$| jS dS )z]If the ``DataLoader`` is wrapping a mapping based Dataset, return the ``FastForwardSampler``.N)r^   r;   r:   Zbatch_samplerr   r3   r3   r4   _find_fast_forward_samplersd  s    r   )r   rd   r=   c                 C   s   t | }t|dd}|d |kr8td| d|d  t|trt|jdkrRqBt|d trt|d d D ]}t|j qp|j	|_
|j
| dd	 |S )
zThis function is used to cycle back the DataLoader ``_MultiProcessingDataLoaderIter`` workers and call the
    reset function.

    Returns:
        iterator: Return the iterator generated from the provided ``DataLoader``.
    Z_num_workersr   r*   zThe provided `num_workers` z= doesn't match the one used while generating the checkpoint: Zprevious_workerrY   T)Z
first_iter)rU   rL   r#   r^   r   r[   Z_worker_queue_idx_cycler0   rangeZ_original_resetZ_reset)r   rd   Ziter_dataloaderr*   _r3   r3   r4   _cycle_to_next_worker_and_resetm  s"    
r   )samplesr   
collate_fnfault_tolerant_moder=   c           	      C   s   || }d}|j r| }nlt|dd}t }|r8|jnd}|durx| }||vrx|rp|jdkrptd| dd|i}|du r|i i}d|tj|iS )a  A collate_fn function that adds the state dict of a :class:`CaptureIterableDataset` or
    :class:`CaptureMapDataset` used in the worker processes. This function gets executed within the worker
    processes. The structure will be:

    .. code-block:: python

        {
            "data": ...,  # data returned by Dataset
            "__pl_restart_meta": {"sampler_name0": state_dict0, "sampler_name1": state_dict1},
        }
    Nrd   r   rY   zThe state_dict returned by z1 needs to be indexed by `worker_id` integer keys.data)	is_automaticrd   rL   r   rS   r*   r#   r"   PL_RESTART_META)	r   r   r   r   r   metadataZstate_dict_fninfor)   r3   r3   r4   _capture_metadata_collate  s$    

r   z)pl.utilities.fetching.AbstractDataFetcher)fnitdlr+   data_fetcherr=   c                    s&   t td fdd}|S )NrP   c                     s   j }  }|d |tj  }}d7 t| trLfdd| D }nFt| trt}|d usjJ tj	|
|t| d dg} | |S )Nr   rY   c              	      s0   g | ](\}}t  j|t| d  |dqS )r   )r*   r(   r+   r)   r,   )rl   r*   listr~   )ry   Zsampler_iter_nameiterator_state)r   r+   r3   r4   
<listcomp>  s   z7_next_data_wrapper.<locals>.wrapper.<locals>.<listcomp>r   r*   r(   r'   r)   r+   )r   r"   r   r^   r   r   r   r   rl   r*   rd   r   r~   _store_dataloader_iter_state)r   combined_batchbatchr6   
ff_samplerr   r   r   r   r+   r3   r4   wrapper  s*    


	z#_next_data_wrapper.<locals>.wrapper)r   r   )r   r   r   r+   r   r   r3   r   r4   _next_data_wrapper  s    #r   )r   iteratorr   r+   r=   c                 C   s:   t  jsdS t| jttfs"J t|j|| |||_dS )aV  Patches the iterator of a PyTorch dataloader by injecting logic for fault-tolerant training when it is
    necessary to remove the sampler state dict from provided data batch.

    The custom data has this format:
    .. code-block:: python
        {
            "batch": ...,  # data returned by DataLoader
            "__pl_restart_meta": {
                "sampler0": {
                    0: {"current_iteration": ...},
                    1: {"current_iteration": ...},
                },
                "sampler1": ...,
            },
        }
    Each sampler in the worker process tracks the current iteration. We return all of them to the main process
    as part of the sample and then a special collate function :func:`_capture_metadata_collate`
    will extract the current iteration as part of the metadata returned by a custom batch.
    N)	r!   detect_current_moder   r^   r   r   r   r   
_next_data)r   r   r   r+   r3   r3   r4   patch_dataloader_iterator  s    
r   c                 C   sD   t  }| j}|jr(t|tr,|jtu r,dS tt| j||d| _dS )zeWrap default collate function to retrieve captured dataset state dict when fault tolerant is enabled.N)r   r   r   )	r!   r   r   
is_enabledr^   r   funcr   r   )r   r   r   r3   r3   r4   _add_capture_metadata_collate  s    r   c                 C   s^   |d d }t |ts t|}t| }|d us4J ||j | jj|j|d |jd d S )Nr6   r   r7   )r7   r*   )	r^   rl   rq   r   rg   r(   r   r'   r*   )r   rd   r   r   r3   r3   r4   3_reload_dataloader_state_dict_automatic_map_dataset  s    

r   )r   rd   r=   c                 C   s    |  dd |d  D  d S )Nc                 S   s   i | ]\}}||d  d qS )r   r(   r3   )ry   r   r6   r3   r3   r4   r{   .  r|   zL_reload_dataloader_state_dict_automatic_iterable_dataset.<locals>.<dictcomp>r6   )rg   r   )r   rd   r3   r3   r4   8_reload_dataloader_state_dict_automatic_iterable_dataset*  s    r   c                 C   s>   | j }t|trt| | nt|tr2t|| ntdd S )N3This shouldn't be happening. Please, open an issue.)r   r^   r   r   r   r   r#   )r   rd   r   r3   r3   r4   '_reload_dataloader_state_dict_automatic2  s    

r   c                    s    d } d | d } d |  dd }|rr|D ]<}t| |}t|tsbtd| d| d|||  q4t| jtsd S  fdd	 d  D }| jt||| d S )
Nr7   r6   r*   r(   zThe DataLoader attribute :z( should have a `load_state_dict` method.c                    s"   i | ]}| d  | d | qS r   r3   )ry   r)   r   r3   r4   r{   S  s   z8_reload_dataloader_state_dict_manual.<locals>.<dictcomp>)	getrL   r^   r   r#   rg   r   r~   r   )r   rd   r7   r*   r(   Zdataloader_attr_nameobjr'   r3   r   r4   $_reload_dataloader_state_dict_manual>  s"    



r   c                 C   sB   t  }|jsdS |jr$t| | n|jr6t| | ntddS )zCUtility to reload state_dict within dataloader for fault tolerance.Nr   )r!   r   r   r   r   	is_manualr   r#   )r   rd   r   r3   r3   r4   _reload_dataloader_state_dict[  s    r   )r6   r7   r*   r=   c                    sf   dkrS |d kr t dtkr4t d|d   fddtD }fdd|D S )	znThis function is used to rotate the worker indices based on the `latest_worker_id` the training failed
    on.r   rY   z=The `latest_worker_id` should be within [0, num_workers - 1].z4The `state` should contain `num_workers - 1` values.c                    s   g | ]} |  |fqS r3   r3   )ry   i)next_worker_idr*   r3   r4   r   w  r|   z*_rotate_worker_indices.<locals>.<listcomp>c                    s"   i | ]\}}| v r| | qS r3   r3   )ry   Zold_idZnew_idr6   r3   r4   r{   x  r|   z*_rotate_worker_indices.<locals>.<dictcomp>)r#   r_   r   )r6   r7   r*   Zold_to_new_worker_id_mapr3   )r   r*   r6   r4   r   m  s    r   c                       s|   e Zd ZdZeeef ddddZddddZed fd	d
Z	e
ddddZddddZed fddZ  ZS )_StatefulDataLoaderIterz;This mixin is used to make PyTorch DataLoaderIter stateful.N)r(   r=   c                 C   s(   t | ddd | _| j|| jf d S )N_sampler_state_idxr   rY   )rL   r   _sampler_stateappendrF   r(   r3   r3   r4   Z__accumulate_state~  s    z*_StatefulDataLoaderIter.__accumulate_staterP   c                 C   s$   dd | j j D }| | dS )z;This function is used to extract the sampler states if any.c                 S   s,   i | ]$\}}t |tr|d kr|| qS )r   )r^   r   rd   r   r3   r3   r4   r{     s   z@_StatefulDataLoaderIter._store_sampler_state.<locals>.<dictcomp>N)_loaderrK   r   )_StatefulDataLoaderIter__accumulate_stater   r3   r3   r4   _store_sampler_state  s    
z,_StatefulDataLoaderIter._store_sampler_statec                    s   t   }|   |S rJ   )r>   _next_indexr   )rF   ZindexesrG   r3   r4   r     s    
z#_StatefulDataLoaderIter._next_index)loaderr=   c                 C   s,   t | || _|j| _d| _g | _d| _d S rQ   )r   r   _lightning_fetcher_data_fetcherr+   r   r   rF   r   r3   r3   r4   _prepare_loader  s    z'_StatefulDataLoaderIter._prepare_loaderc                 C   s$   t | jjtr | jjjd | j_d S )Nr   )r^   r   r   r   keywordsrW   r3   r3   r4   __del__  s    z_StatefulDataLoaderIter.__del__c                    s   t   }|d |tj  }}|  jd7  _| jd\}}t| d }t	| j
j|||| jdg}|| jkstJ | j| | |S )Nr   rY   r   r   )r>   r   r"   r   r+   r   popr   r~   rl   r   r*   r   r   )rF   r   r   r6   r(   Zsampler_state_idxr)   rG   r3   r4   r     s     

z"_StatefulDataLoaderIter._next_data)r-   r.   r/   rh   r
   r0   r   r   r   r   r   r   r   r   rj   r3   r3   rG   r4   r   {  s   	r   c                       s"   e Zd Zed fddZ  ZS )$_SingleProcessDataLoaderIterStatefulr   c                    s   |  | t | d S rJ   r   r>   r?   r   rG   r3   r4   r?     s    
z-_SingleProcessDataLoaderIterStateful.__init__r-   r.   r/   r   r?   rj   r3   r3   rG   r4   r     s   r   c                       s"   e Zd Zed fddZ  ZS )&_MultiProcessingDataLoaderIterStatefulr   c                    s   |  | t | d S rJ   r   r   rG   r3   r4   r?     s    
z/_MultiProcessingDataLoaderIterStateful.__init__r   r3   r3   rG   r4   r     s   r   r   )rF   r=   c                 C   sB   t | dstd| jdkr$t| S t | dr6|   t| S d S )Nr   z_A stateful iterator should be used only when a DataFetcher has been attached to the DataLoader.r   check_worker_number_rationality)hasattrr#   r*   r   r   r   rW   r3   r3   r4   _get_iterator  s    


r   rP   c                   C   s*   t  jsdS ttds tjt_tt_dS )zSThis function is used to replace the DataLoader iterator by their stateful version.N_ori_get_iterator)r!   r   r   r   r   r   r   r3   r3   r3   r4   _patch_dataloader_get_iterators  s
    

r   c                  C   s   t tdd} | r| t_t`dS )zUThis function is used to restore the DataLoader `get_iterator` with its original one.r   N)rL   r   r   r   )Zget_iteratorr3   r3   r4   "_teardown_dataloader_get_iterators  s    r   c                    s   t ttf | j}t|dd d u r(tddd |j D }|sHtd fdd|	 D }|srtd  d	t
|d
krtd| d	t|d tu r|jrtdnt|d turtdd S )Nr]   zFault-tolerance doesn't support an `IterableDataset` without `__next__` method implemented. Hint: We recommend you to move your logic from `__iter__` inside and rely on a sampler to perform the sample sampling.c                 S   s    i | ]\}}t |tr||qS r3   )r^   r   r   r3   r3   r4   r{     r|   z._validate_iterable_dataset.<locals>.<dictcomp>zRFault-tolerance doesn't support an IterableDataset without a sampler as attribute.c                    s   g | ]}t | v r|qS r3   )type)ry   r   SUPPORTED_SAMPLERSr3   r4   r     r|   z._validate_iterable_dataset.<locals>.<listcomp>Fault-tolerance supports only .rY   z@A single sampler is supported within an Iterable Dataset. Found r   @A `DistributedSampler` sampler shuffle attribute is set to True.&Only `SequentialSampler` is supported.)r   r   r   r   rL   AttributeErrorrK   r   	TypeErrorvaluesr_   
ValueErrorr   shuffle)r   r   r   r;   r3   r   r4   _validate_iterable_dataset  s$    

r   c                 C   sn   t ttf}t| dd }|d ur:t||vr:td| dt|tu rV|jrVtdnt|t u rjtdd S )Nr;   r   r   r   r   )r   r   r   rL   r   r   r   )r   r   r;   r3   r3   r4   _validate_map_dataset  s    

r   zpl.trainer.states.RunningStage)r   stager=   c                    s   t  jsdS ddlm}m  t| |r0| j}n| }g tt	 t
f dd fdd}t|t	 f| tdkr|tjjjjkrtdD ].} t| t	sJ t| jtrtnt}||  qdS )	zVThis function is used to validate that Fault-tolerance is possible with the user data.Nr   )CombinedLoaderCycleIteratorr   c                    s   t |  r| j} |  d S rJ   )r^   r   r   r   r   Z
dl_loadersr3   r4   flatten_dataloader  s    
z>_validate_fault_tolerant_automatic.<locals>.flatten_dataloaderrY   z2Fault-tolerance supports only a single dataloader.)r!   r   r   Z$pytorch_lightning.trainer.supportersr   r   r^   loadersr   r   r   r   r_   plZtrainerZstatesZRunningStageZTRAININGr   r   r   r   r   )r   r   r   Zdataloadersr   Zvalidator_fnr3   r   r4   "_validate_fault_tolerant_automatic  s     

 r   r6   )rd   rI   r=   c                    s"   t t d fdd t| t  S )zKThis utility collects the state across processes for a collection of state.)r6   r=   c                    s&   | v rt | S  fdd|  D S )Nc                    s   i | ]\}}|t |t qS r3   )r   r
   r   )r   r3   r4   r{   5  r|   zL_collect_states_on_rank_zero_over_collection.<locals>.fn.<locals>.<dictcomp>)r    r   r   r   rI   r3   r4   r   2  s    z8_collect_states_on_rank_zero_over_collection.<locals>.fn)r
   r   )rd   rI   r3   r   r4   ,_collect_states_on_rank_zero_over_collection/  s    r   )r   )r6   )Scollections.abcr   copyr   Zdataclassesr   r   	functoolsr   r   typingr   r	   r
   r   r   r   r   r   r   r   Z#lightning_utilities.core.apply_funcr   Ztorch.utils.datar   r   r   r   r   r   Ztorch.utils.data.dataloaderr   r   r   r   r   Ztyping_extensionsr   Zpytorch_lightningr   Z lightning_fabric.utilities.typesr   Z'pytorch_lightning.utilities.distributedr    Z!pytorch_lightning.utilities.enumsr!   r"   Z&pytorch_lightning.utilities.exceptionsr#   Z pytorch_lightning.utilities.seedr$   r%   r&   r5   r:   rl   rt   r   r   r   r2   r   r   r0   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r3   r3   r3   r4   <module>   s|   0 	t
6%N	%&2 #$<		!