
    h}                         S SK r S SKrS SKrS SKrS SKJr  S SKrS SKrS SKJ	r	J
r
JrJrJrJrJrJr   " S S5      rS rSS\4S jjrSS\4S	 jjr\S
:X  a   \" 5       r\" \" \R2                  5      5        gg)    N)environ)	NVMLErrornvmlDeviceGetCountnvmlDeviceGetHandleByIndexnvmlDeviceGetMemoryInfonvmlDeviceGetNamenvmlInitnvmlShutdownnvmlSystemGetDriverVersionc                       \ rS rSrSrSS jrS rS\4S jrS\S	\	S\
4S
 jrS\4S jrS\4S jrS\	\
   4S jrS\4S jrS\4S jrS\4S jrSrg)MachineInfo   z'Class encapsulating Machine Info logic.Nc                 F   Xl         UcC  [        R                  " S[        R                  S9  [        R                  " [
        5      U l        OX l        S U l         U R                  5       U l        g ! [         a%    U R                  R                  S5        S U l         g f = f)Nz3%(asctime)s - %(name)s - %(levelname)s: %(message)s)formatlevelz"Exception in getting machine info.)silentloggingbasicConfigINFO	getLogger__name__loggermachine_infoget_machine_info	Exception	exception)selfr   r   s      _/var/www/fran/franai/venv/lib/python3.13/site-packages/onnxruntime/transformers/machine_info.py__init__MachineInfo.__init__   s    >Lll "++H5DK K 	% $ 5 5 7D 	%KK!!"FG $D	%s   A1 1,B B c           
      X   U R                  5       n[        R                  " 5       nUU R                  5       U R                  5       [        R                  " 5       U R                  US/5      U R                  5       U R                  5       U R                  5       U R                  5       S.	nU$ )z!Get machine info in metric formatpython_version)	gpucpumemoryospythonpackagesonnxruntimepytorch
tensorflow)
get_gpu_info_by_nvmlcpuinfoget_cpu_infoget_memory_infoplatform_try_getget_related_packagesget_onnxruntime_infoget_pytorch_infoget_tensorflow_info)r   gpu_infocpu_infor   s       r   r   MachineInfo.get_machine_info2   s    ,,.'') $$&**,##%mmH/?.@A113446,,.224

     returnc                 ^    [         R                  " 5       nUR                  UR                  S.$ )zGet memory info)total	available)psutilvirtual_memoryr<   r=   )r   mems     r   r/   MachineInfo.get_memory_infoD   s#    ##%??r9   r7   namesc           	          U HX  nX1;   d  M
  X   n[        U[        [        45      (       a-  SR                  U Vs/ s H  n[	        U5      PM     sn5      s  $ Us  $    gs  snf )N, )
isinstancelisttuplejoinstr)r   r7   rB   namevalueis         r   r1   MachineInfo._try_getI   s[    D edE]3388U$;USVU$;<<   %<s   A 
c           	      L   [         R                  " 5       nU R                  USS/5      [        R                  " SS9[        R                  " SS9U R                  US/5      U R                  US/5      U R                  US/5      [
        R                  " 5       R                  S	.$ )
zGet CPU infobrand	brand_rawF)logicalT	hz_actuall2_cache_sizeflags)rP   coreslogical_coreshzl2_cacherU   	processor)r-   r.   r1   r>   	cpu_countr0   unamerZ   )r   r7   s     r   r.   MachineInfo.get_cpu_infoR   s    '') ]]8g{-CD%%e4#--d;--;-8h0AB]]8gY7!)33
 	
r9   c                    / nSn [        5         [        5       n[        5       n[        U5       HX  n[	        U5      n[        U5      n0 nUR                  US'   UR                  US'   [        U5      US'   UR                  U5        MZ     [        5         X!S.n	S[        ;   a  [        S   U	S'   U	$ ! [         a7  nU R                  (       d  U R                  R                  SU5         SnAgSnAff = f)	zGet GPU info using nvmlNmemory_totalmemory_availablerK   z-Error fetching GPU information using nvml: %s)driver_versiondevicesCUDA_VISIBLE_DEVICEScuda_visible)r	   r   r   ranger   r   r<   freer   appendr
   r   r   r   errorr   )
r   gpu_info_listra   deviceCountrM   handleinfor6   rh   results
             r   r,    MachineInfo.get_gpu_info_by_nvml`   s    	J79N,.K;'3A6.v6+/::(/3yy+,#4V#< $$X. ( N %3M!W,%,-C%DF>"  	;;!!"QSXY	s   BB1 1
C2;-C--C2c                     SS K nUR                  n/ SQnU Vs0 s H+  oDR                  U;   d  M  UR                  UR                  _M-     nnU$ s  snf )Nr   )zonnxruntime-gpur)   onnxtransformersprotobufsympytorchr+   flatbuffersnumpyzonnxconverter-common)pkg_resourcesworking_setkeyversion)r   rw   installed_packagesrelated_packagesrM   related_packages_lists         r   r2    MachineInfo.get_related_packages|   s[    *66
 <N k;MaQVQVZjQj!1		!1;M k$$ !ls
   AAc                 P    SS K nUR                  SUR                  5       ;   S.$ ! [         a6  nU R                  (       d  U R
                  R                  U5         S nAg S nAf[         a7  nU R                  (       d  U R
                  R                  US5         S nAg S nAff = f)Nr   CUDAExecutionProvider)rz   support_gpuF)r)   __version__get_available_providersImportErrorr   r   r   r   )r   r)   rh   r   s       r   r3    MachineInfo.get_onnxruntime_info   s    	 '226+:]:]:__   	;;%%e, 	;;%%i7	s    #& 
B%,A!!B%.-B  B%c                     SS K nUR                  UR                  R                  5       UR                  R                  S.$ ! [
         a6  nU R                  (       d  U R                  R                  U5         S nAg S nAf[         a7  nU R                  (       d  U R                  R                  US5         S nAg S nAff = f)Nr   )rz   r   cudaF)
rt   r   r   is_availablerz   r   r   r   r   r   )r   rt   rh   r   s       r   r4   MachineInfo.get_pytorch_info   s    	 !,,$zz668** 
  	;;%%e, 	;;%%i7	s#   ?A 
C,A==C
-B<<Cc                     SS K nUR                  R                  UR                  R                  UR                  R                  5       S.$ ! [         a6  nU R                  (       d  U R                  R                  U5         S nAg S nAf[         a6  nU R                  (       d  U R                  R                  U5         S nAg S nAff = f)Nr   )rz   git_versionr   )r+   rz   VERSIONGIT_VERSIONtestis_built_with_cudar   r   r   r   ModuleNotFoundError)r   tfrh   s      r   r5   MachineInfo.get_tensorflow_info   s    	# ::--!zz55!ww99; 
  	;;%%e," 	;;%%e,	s$   A	A 
C
,BC
,CC
)r   r   r   )FN)r   
__module____qualname____firstlineno____doc__r   r   dictr/   rG   rJ   r1   r.   r,   r2   r3   r4   r5   __static_attributes__ r9   r   r   r      s    1%&$@ @
 d s 
d 
d 8%d3i %(d "$ $T r9   r   c                      [         R                  " 5       n U R                  SSSSS9  U R                  SS9  U R	                  5       nU$ )Nz--silentF
store_truezDo not print error message)requiredactionhelp)r   )argparseArgumentParseradd_argumentset_defaults
parse_args)parserargss     r   parse_argumentsr      sU    $$&F
)	   u%DKr9   r:   c                 V    [        U 5      n[        R                  " UR                  SS9$ )N   indent)r   jsondumpsr   )r   machines     r   r   r      s#    &!G::g**155r9   c                     [        U 5      nUR                  nU(       a,  UR                  5        VVs0 s H  u  p4US;   d  M  X4_M     nnn[        R                  " USS9$ s  snnf )N)r#   r$   r%   r   r   )r   r   itemsr   r   )r   r   rl   ry   rL   s        r   get_device_infor      sY    &!GD-1ZZ\]\zsSD\=\

\]::d1%% ^s   A A __main__)T)r   r   r   r0   r&   r   r-   r>   py3nvml.py3nvmlr   r   r   r   r   r	   r
   r   r   r   rJ   r   r   r   r   printr   r   r9   r   <module>r      s~          	 	 	g gT6S 6
&C & zD	
4;;
'( r9   