o }oi@sddlZddlZddlZddlZddlmZddlmZddlm Z m Z m Z m Z m Z mZddlZddlmZddlmZddlmZmZddlmZdd lmZdd lmZdd lmZdd lm Z dd l!m"Z"ddl#m$Z$ddl%m&Z&m'Z'm(Z(ddl)m*Z*m+Z+GdddeZ,dS)N)deepcopy)Path)AnyDictIterableListOptionalUnion)proxy)get_filesystem)ModelCheckpoint_is_local_file_protocol)call)rank_zero_info)EMA)logging)AppState)AsyncFinalizableCheckpointIO)is_global_rank_zero) ckpt_to_dirinject_model_parallel_rankuninject_model_parallel_rank)import_multistorageclientis_multistorageclient_urlcseZdZdZdZ        dXded ed ed ed ed ededeffdd ZddZ ddZ de ee fddffdd Z deddffdd ZfddZfdd Zdeefd!d"ZdYd#eedefd$d%Zd&eddfd'd(Zd)d*deefd+d,Zd&eeefd-ee ddfd.d/Zd&eeefdeefd0d1ZdZd2d3Zed4eeefdefd5d6Zed4eeefdefd7d8Zed[d4eeefddfd9d:Z ed[d4eeefddfd;d<Z! d\d&ed)d*d=edefd>d?Z"d)d*d&eddffd@dA Z#d)d*d&edBefdCdDZ$ d[d)d*d&eddffdEdF Z%d&edefdGdHZ&dIe'edefdJdKZ(d&eeefdefdLdMZ)e*de'efdNdOZ+edPeeefddfdQdRZ,d)dSdTedUedefdVdWZ-Z.S)]NeMoModelCheckpointa?Light wrapper around Lightning's ModelCheckpoint to force a saved checkpoint on train_end. Extends Lightning's on_save_checkpoint func to save the .nemo file. Saves the .nemo file based on the best checkpoint saved (according to the monitor value). Also contains func to save the EMA copy of the model. z -unfinishedFT.nemoNalways_save_nemosave_nemo_on_train_endsave_best_modelpostfixn_resumemodel_parallel_size async_savesave_last_n_optim_statesc  s||_||_||_||_|jr|jstd||_d|_||_||_ d|_ g|_ d| vr4| d|_ nd|_ tjdi| |jdkrR|rTtd|dSdSdS)NzFound save_best_model is True and save_nemo_on_train_end is False. Set save_nemo_on_train_end to True to automatically save the best model.prefixrzChecking previous runs)rrrr$rwarningr previous_best_pathr"r#async_finalize_cbdeferred_ckpts_to_removepopr&super__init__ save_top_kdebugnemo_topk_check_previous_run) selfrrrr r!r"r#r$kwargs __class__r'^/home/ubuntu/.local/lib/python3.10/site-packages/nemo/utils/callbacks/nemo_model_checkpoint.pyr.0s,    zNeMoModelCheckpoint.__init__c sTzjjjjWn tytdwi_d_d_d_tfddjD}|D]Z}dt|vsBdt|vrFt|}t|}|ddd ksZ|d dd kr[q4| j t j d }|t j krt d ||d}|r||||d }t|j|<q4t jd krdSjdkrdnd}tjjj|d}jdur|drt |j}nt |jj}nt |j}td|}td|j} t|D]/} |d} j| | | rj| r| td| q|d_|d_jj_dS)z3 Check if there are previous runs. zQLightning's ModelCheckpoint was updated. NeMoModelCheckpoint will need an update.r%Nc3s|] }|s|VqdSN_is_ema_filepath).0pathr2r'r6 pszCNeMoModelCheckpoint.nemo_topk_check_previous_run..mp_ranktp_ranki -last.ckpt-lastz[A-z]minFTkeyreverserzNumber of models to delete: rRemoved checkpoint: ) best_k_modelskth_best_model_pathbest_model_scorebest_model_pathAttributeErrorlist_saved_checkpoint_pathsstrrfindmonitorlenresearchstartfloatmodesortedgetr"is_dirr/maxrr0_has_ema_ckptsranger,_del_model_without_trainer_fsexists_ema_format_filepath) r2 checkpoints checkpointindexmatchvalue_reverserImodels_to_delete ema_enabled_modelr'r<r6r1`s`             z0NeMoModelCheckpoint.nemo_topk_check_previous_runcsdtdtffdd fddjD_tjdkrHjdk}tjjj|d }|d _jj_ |d_ jj _ dSd _d_ d _ d_ dS) N ckpt_pathreturncs>tj|ptjt|ptj|d}|o| S)N.ckpt)osr;isfilerisdir removesuffixis_checkpoint_unfinished)rmrar<r'r6 __is_ckpt_oks zKNeMoModelCheckpoint._remove_invalid_entries_from_topk..__is_ckpt_okcsi|] \}}|r||qSr'r')r:kv) _NeMoModelCheckpoint__is_ckpt_okr'r6 szINeMoModelCheckpoint._remove_invalid_entries_from_topk..rrDrErr%) rPboolrIitemsrSrXrYrZrJ kth_valuerLrK)r2 reverse_arrbest_k_models_arrr')rxr2r6!_remove_invalid_entries_from_topks    z5NeMoModelCheckpoint._remove_invalid_entries_from_topk state_dictrncst||dS)z& Load the state dict. N)r-load_state_dictr)r2rr4r'r6rs  z#NeMoModelCheckpoint.load_state_dictstagecshtrtdt|jtjrtj t ||||j |j}||_|j |j|_dS)z' Setup the checkpoint. z)Removing unfinished checkpoints if any...N)rrr0r_remove_unfinished_checkpointsdirpathtorch distributedis_initializedbarrierr-setupstrategy broadcastrmlast_model_path)r2trainer pl_modulerr;r4r'r6rs    zNeMoModelCheckpoint.setupc st|||}|js |St}|jdur|jdkrtd||_|jdur4|jdkr4t |j }n|j }|j rt j |sBdS|j |jkrOtd|S|j |_t|}tj|ddd}d|vri|d}|j|d d tjrztj||}|j|jd td |j|j|d d ntjrtj||}|j|jd td |j|durtrtd|t|||S)z& Save the checkpoint. NrCz@always_save_nemo will slow down training for model_parallel > 1.z*Best model has not changed, skipping save.cpuF) map_location weights_onlyrTstrict save_pathzNew best .nemo model saved to: zNew .nemo model saved to: Removing old .nemo backup )r-on_save_checkpointrrr"rr(_format_nemo_checkpoint_namemodel_restore_pathrrLrrpr;rar)r0rrrloadrrrr_backup_existing_nemo_ckptsave_toinforr rm) r2rrrdoutput app_statemaybe_injected_best_model_pathold_state_dict backup_pathr4r'r6rsJ             z&NeMoModelCheckpoint.on_save_checkpointcs|jrdS|jrR|jdkrRd}t|jtr|j|jdkrd}t|jtr/|j|jdkr/d}|rR||}|j| ||j krKt d|jdnt ||t |||jr|jd|jdkrpt |d n$tj|jd dr|jd d|_|j|j|_|j|j|jr||}|j|d |durt rt !d |t"|#|dSdSdSdS) z3 Save the checkpoint on train end. NrFTzLast checkpoint z already savedz&SaveBestCheckpointConnector.resume_endr%z was told to save the best checkpoint at the end of training, but no saved checkpoints were found. Saving latest model instead.rorr)$ fast_dev_run save_lastval_check_interval isinstancerW global_stepint_monitor_candidatesrformat_checkpoint_nameCHECKPOINT_NAME_LASTrr0r-_save_last_checkpoint on_train_endrrrrLr(rpr;rrsplitr_checkpoint_connectorrestorerrrrrrr r)r2rrshould_save_last_checkpointmonitor_candidatesrr4r'r6rs@    z NeMoModelCheckpoint.on_train_endcCs|}|}|jr%|j}|j||ddr%||}|d7}|j||dds||kr+dS|jrCt|d|t|r=nt |||j |S)aSearch for an available name with version infix and rename existing checkpoint. NOTE: this behavior is slightly different from regular checkpoints. PTL creates new regular checkpoint with the first available name. Here, for backward compatibility, we create .nemo checkpoint as before and create a backup under the first available name. Args: trainer (Trainer): trainer instance. Returns: Path to the backup checkpoint or None, if no backup was created F)check_dist_ckptrCNz/ already exists, moving existing checkpoint to ) r_enable_version_counterSTARTING_VERSION file_existsis_global_zerorrrshutilmoverr)r2r base_pathavailable_path version_cntr'r'r6r3s"   z.NeMoModelCheckpoint._backup_existing_nemo_ckptverc Csn|durdn|jd|}t|jr!|jd|j||jStjtjtj |j|j||jS)Nr%rw/) CHECKPOINT_JOIN_CHARrrr&r rpr;abspath expanduserjoin)r2r version_infixr'r'r6rVs  "z0NeMoModelCheckpoint._format_nemo_checkpoint_namefilepathcCst|}t|r5tr3zt|}tj|ddtd|WdStd|dYdSdSt}|j durF|j dkrFt |}tsS|j durt|j dkrvz|j |td|WdStd |dYdSdSdS) NT) ignore_errorsz Removed distributed checkpoint: z(Tried to remove distributed checkpoint: z but failed.rCrrHzTried to remove checkpoint: )rrr[rrrmtreerrrr"rdata_parallel_rankr`r)r2r dist_ckptrr'r'r6r_^s*   z.NeMoModelCheckpoint._del_model_without_trainerrzlightning.pytorch.TrainercCs"d}|jD] }t|tr|}q|Sr7) callbacksrr)r2r ema_callbackcallbackr'r'r6 _ema_callback~s   z!NeMoModelCheckpoint._ema_callbackstorage_optionsc Cs||}d}t||jd}t||jkr||}td|d|jj|dd}||||dursz=NeMoModelCheckpoint._get_checkpoints_list..cSst|ddddS)Nz-step=rC-r)rr)xr'r'r6z;NeMoModelCheckpoint._get_checkpoints_list..)rFcsg|] }tj|qSr')rpr;r)r:rdrr'r6rs)rpr;dirnamelistdirrY)r2rrcr'rr6rs  z)NeMoModelCheckpoint._get_checkpoints_listcCs&t|d||jj||jjddS)Non_load_checkpointr)r_call_lightning_module_hookrload_model_state_dictlightning_modulestrict_loading)r2rrdr'r'r6rs  z,NeMoModelCheckpoint._load_current_state_dictrcCs8tt|}|d}|d}|d}t|tjS)aFormat the path to the unfinished checkpoint marker file. If the marker file exists, corresponding checkpoint is considered unfinished/incomplete. NOTE: Marker path for the EMA checkpoint part is the same as for the original checkpoint. Args: checkpoint_path: Path to the checkpoint file or dir. Does not need to exist. Returns: Path to the unfinished checkpoint marker file. rro-EMA)rPrrsrrUNFINISHED_CHECKPOINT_SUFFIX)rmarker_filepathr'r'r6(format_checkpoint_unfinished_marker_paths    z      z$NeMoModelCheckpoint._save_checkpointrcsfdd}|S)zLCreates a callback that can be used to finalize async (and sync) ckpt saves.c stdd__jr!jD] }|tqjddj s-dSt dddt dj sCJj d }td ||D] }j|dd qSdS) Nz"Finalize callback called for step z , filepath TrzAsync checkpoint save for step z (z) finalized successfully at rrzCheckpoints to remove: )override_async)rr0_last_global_step_saved_last_checkpoint_savedrloggersafter_save_checkpointr rr#rrr+r,_remove_checkpoint)loggerckpts_to_removeckpt_to_removerrr2rr'r6_cbLs$   zGNeMoModelCheckpoint._get_finalize_save_checkpoint_callback.._cbr')r2rrrrr'rr6rGsz:NeMoModelCheckpoint._get_finalize_save_checkpoint_callbackcsv|jr|s|jd|dS|j|ddt||||}|dur2||}t|||j|dddS)aPerforms checkpoint removal or deferred removal. With async save, `self._remove_checkpoint` is called before the checkpoint is actually finished so we can't remove it. Instead we add it to `self.deferred_ckpts_to_remove` for future removal. rNTrr) r#r+rrr-rrrbr)r2rrr rr4r'r6rks   z&NeMoModelCheckpoint._remove_checkpointcCs||jd|jSNr)replaceFILE_EXTENSIONr2rr'r'r6rbz(NeMoModelCheckpoint._ema_format_filepathrccstfdd|DS)Nc3s|]}|VqdSr7r8)r:rr<r'r6r=sz5NeMoModelCheckpoint._has_ema_ckpts..)any)r2rcr'r<r6r]rz"NeMoModelCheckpoint._has_ema_ckptscCst|d|jSr)rPendswithrrr'r'r6r9rz$NeMoModelCheckpoint._is_ema_filepathcs~tjrt}|jdSddtjdD}|r)tfdd|SddtjdD}tfd d|S) Nz/*.ckptcSsg|]}|r|qSr')r[rr'r'r6rsz?NeMoModelCheckpoint._saved_checkpoint_paths..*c | Sr7rtpr<r'r6r z=NeMoModelCheckpoint._saved_checkpoint_paths..cSsg|]}|qSr'r'r:fr'r'r6rs*.ckptcrr7rrr<r'r6rr )rrrglobrfilterrglob)r2mscdist_checkpointscheckpoint_filesr'r<r6rOs z+NeMoModelCheckpoint._saved_checkpoint_pathscheckpoint_dirc Cs"tstdt|r)t}||dtj}t|}|D]}||qdSt |}dd|dtjD}dd| dD}|D]}t |}||vr_t d|t|qGdd|dD}|D]}t |}||vrt d |t|ql|D]} t| qdS) Nz8_remove_unfinished_checkpoints should run only on rank 0rcSh|] }|r|qSr')is_fileresolver!r'r'r6 s zENeMoModelCheckpoint._remove_unfinished_checkpoints..cSsh|]}|qSr')r-r!r'r'r6r.sr#z Removing unfinished checkpoint: cSr+r')r[r-rr'r'r6r.rz%Removing unfinished dist checkpoint: )rAssertionErrorrrr$rrr rrr&rrr(rpremoverr) r*r'existing_marker_filepathsfs ckpt_filepathcheckpoint_filepathspossible_marker_path all_dirpaths ckpt_dirpathrr'r'r6rs@      z2NeMoModelCheckpoint._remove_unfinished_checkpointsz pl.TrainerpreviouscurrentcCs||krdSt|s dSt|}|jdurt|jnd}|dur8||kr8t|dr6|jdr6ndS|jdurEt|j dt|j}||j vS)aChecks if the previous checkpoint should be deleted. A checkpoint won't be deleted if any of the cases apply: - The previous checkpoint is the same as the current checkpoint (means the old was already overwritten by new) - The previous checkpoint is not in the current checkpoint directory and the filesystem is local - The previous checkpoint is the checkpoint the Trainer resumed from and the filesystem is local and the resumed from checkpoint is not the last checkpoint FTNr@z.dirpath is None.) r rabsolutermrPrnamerrr5r)r2rr8r9 resume_pathrr'r'r6_should_remove_checkpoints   z-NeMoModelCheckpoint._should_remove_checkpoint)FTFrFNFrr7)rnN)F)T)/__name__ __module__ __qualname____doc__rrzrPrr.r1rrrrrrrrrrr_rrr rrrrr staticmethodrrtrrrrrrrbrr]r9propertyrOrr= __classcell__r'r'r4r6r's 0B  2-# "+    ) %",r)-rprTrrcopyrpathlibrtypingrrrrrr r_weakrefr #lightning.fabric.utilities.cloud_ior ,lightning.pytorch.callbacks.model_checkpointr r lightning.pytorch.trainerrlightning.pytorch.utilitiesr!nemo.collections.common.callbacksr nemo.utilsrnemo.utils.app_stater!nemo.utils.callbacks.dist_ckpt_iornemo.utils.get_rankrnemo.utils.model_utilsrrrnemo.utils.msc_utilsrrrr'r'r'r6s*