sg $UddlZddlZddlmZmZmZmZmZddlm Z ddl m Z ddl m cmZddlmZgZeeed<ej(eaedZe dZd eeeffd Zd eeeffd Z dd ed ed eeeefgeeefffdZy)N)AnyCallableDictListTypeVar) ParamSpec)uuid4)DCP_LOGGER_NAME__all___T_Preturnci}|jdd}|jdd}|jdd}|jdd}|s|xs|x}r t|dd}| t|n||d<ttj|d<|r|j j |d<|r|j j |d<|r|j j |d<|S)z0 Extracts log data from dcp method args storage_writerNstorage_readerplanner checkpoint_iduuid)getgetattrstrr int __class____name__)argskwargsmsg_dictrrrr serializers V/var/www/html/venv/lib/python3.12/site-packages/torch/distributed/checkpoint/logger.py_msg_dict_from_dcp_method_argsr sHZZ 0$7NZZ 0$7NjjD)GJJ5M N,LnLjL OTB ,7M] _ 57;;'HV%3%=%=%F%F!"%3%=%=%F%F!"%//88 Ocft|i|}|jtj|fi||S)N)r update c10d_logger _get_msg_dict) func_namerrrs rr%r%7s2-t>v>H OOK--iD8DE Or!log_exceptionswrapper_kwargsc >dtttfffd }|S)zKThis method decorator logs the start, end, and exception of wrapped events.funcctjdtjdtjdt ffd }|S)Nrrrctjg|ii |}d|d<tj}||d<|d<tj | |i|}d|d<tj}tj|d<||z |d<tj ||S#t $r?}r7d|d<||d<tj|d<tj|d}~wwxYw) Nstarteventtimer' exceptionerrorend times_spent)r%rr/time_ns _dcp_loggerdebug BaseExceptionr1) rrrt0resultr1t1r*r'r(s rwrapperz6_dcp_method_logger..decorator..wrapperDs$  $(D>(DV(DH !(HW B!HV )7H% &   h ' t.v.!&HW B#||~HV &(2gH] #   h 'M! !(3HW%+0'HW%'+||~HV$%%h/  sB// C78:C22C7) functoolswrapsr rrr )r*r;r'r(s` r decoratorz%_dcp_method_logger..decoratorCsC   277 bii B   @r!)rr r )r'r(r>s`` r_dcp_method_loggerr?>s! "R("H r!)F)r<r/typingrrrrrtyping_extensionsrrr torch.distributed.c10d_logger distributedr$-torch.distributed.checkpoint.logging_handlersr r r__annotations___get_or_create_loggerr5r r r r%boolr?r!rrIs 55'33Ic0k//@  T]t_tCH~Dc3h!))47) xB (2r6"223)r!