§
    Û×j/>  ã                   óÎ   — d Z ddlZddlZddlZddlZddlZddlZddlmZ ddlmZ ddl	m
Z
mZmZmZ ddlmZmZ dZd	Z G d
„ dej        ¦  «        Zd„ Zd„ Z G d„ d¦  «        ZdS )a‰  CPU-backed expert banks with independently reclaimable GPU expert copies.

An ExpertBank registers checkpoint tensors whose first dimension is the expert
index. MMGP leaves these tensors on pageable CPU memory, outside normal block
swapping/pinning. Its residency controller owns every optional GPU copy. Compute
backends use resident_tensors() and lease(); they never transfer model weights.
é    N)Únn)Útqdmé   )Ú	allocatorÚcuda_graphsÚmemory_pressureÚruntime)Úcold_rangesÚprefetch_rangesi    i   c                   ól   ‡ — e Zd Zˆ fd„Zed„ ¦   «         Zd„ Zd„ Zej	        d„ ¦   «         Z
d„ Zd„ Zˆ xZS )Ú
ExpertBankc                 óp  •— t          ¦   «                              ¦   «          t          |¦  «        | _        t	          |¦  «        | _        d | _        i | _        d| _        d| _	        t          j        | j        t          j        d¬¦  «        | _        | j                             ¦   «         | _        d | _        d S )Nr   Úcpu)ÚdtypeÚdevice)ÚsuperÚ__init__ÚintÚnum_expertsÚtupleÚexpert_tensor_namesÚ
_residencyÚ	_residentÚresidency_versionÚ_leasesÚtorchÚzerosÚuint8Úresident_flagsÚnumpyÚ_flagsÚexpert_priority)Úselfr   Útensor_namesÚ	__class__s      €ú4C:\Users\johnn\OneDrive\Documents\AI\mmgp\experts.pyr   zExpertBank.__init__   s›   ø€ Ý‰Œ×ÒÑÔÐÝ˜{Ñ+Ô+ˆÔÝ#(¨Ñ#6Ô#6ˆÔ ØˆŒØˆŒØ!"ˆÔØˆŒå#œk¨$Ô*:Å%Ä+ÐV[Ð\Ñ\Ô\ˆÔØÔ)×/Ò/Ñ1Ô1ˆŒð  $ˆÔÐÐó    c                 óD   ‡ — t          ˆ fd„‰ j        D ¦   «         ¦  «        S )Nc              3   ó®   •K  — | ]O}t          ‰|¦  «        d                               ¦   «         t          ‰|¦  «                             ¦   «         z  V — ŒPdS )r   N)ÚgetattrÚnumelÚelement_size)Ú.0Únamer#   s     €r&   ú	<genexpr>z+ExpertBank.expert_nbytes.<locals>.<genexpr>,   sZ   øè è € Ð|Ð|Ð[_•7˜4 Ñ&Ô& qÔ)×/Ò/Ñ1Ô1µG¸DÀ$Ñ4GÔ4G×4TÒ4TÑ4VÔ4VÑVÐ|Ð|Ð|Ð|Ð|Ð|r'   )Úsumr   ©r#   s   `r&   Úexpert_nbyteszExpertBank.expert_nbytes*   s*   ø€ åÐ|Ð|Ð|Ð|ÐcgÔc{Ð|Ñ|Ô|Ñ|Ô|Ð|r'   c                 ó   — | j         S )zGExpert id to checkpoint-name -> GPU tensor; callers must not mutate it.)r   r1   s    r&   Úresident_tensorszExpertBank.resident_tensors.   s
   € àŒ~Ðr'   c                 ó   — dS )zCOptional backend scratch; called when the owning model is unloaded.N© r1   s    r&   Úrelease_execution_cachez"ExpertBank.release_execution_cache2   s   € € € r'   c              #   óˆ   K  — | xj         dz  c_         	 | j        V — | xj         dz  c_         dS # | xj         dz  c_         w xY w)zEKeep this bank's copies alive while its compute backend submits work.r   N)r   r   r1   s    r&   ÚleasezExpertBank.lease5   sW   è è € ð 	ˆŒ˜ÑˆŒð	Ø”.Ð Ð Ð àˆLŒL˜AÑˆLŒLˆLˆLøˆDˆLŒL˜AÑˆLŒLˆLˆLˆLˆLs	   ”	/ ¯Ac                 óN   — || j         |<   d| j        |<   | xj        dz  c_        d S )Nr   ©r   r!   r   )r#   ÚexpertÚtensorss      r&   Ú_placezExpertBank._place>   s4   € Ø!(ˆŒ�vÑØˆŒ�FÑØÐÔ !Ñ#ÐÔÐÐr'   c                 óJ   — | j         |= d| j        |<   | xj        dz  c_        d S )Nr   r   r;   )r#   r<   s     r&   Ú_evictzExpertBank._evictC   s2   € ØŒN˜6Ð"ØˆŒ�FÑØÐÔ !Ñ#ÐÔÐÐr'   )Ú__name__Ú
__module__Ú__qualname__r   Úpropertyr2   r4   r7   Ú
contextlibÚcontextmanagerr9   r>   r@   Ú__classcell__)r%   s   @r&   r   r      s´   ø€ € € € € ð$ð $ð $ð $ð $ð ð}ð }ñ „Xð}ðð ð ðRð Rð Rð Ôðð ñ Ôðð$ð $ð $ð
$ð $ð $ð $ð $ð $ð $r'   r   c                 ó>   — d„ |                       ¦   «         D ¦   «         S )Nc                 óD   — i | ]\  }}t          |t          ¦  «        ¯||“ŒS r6   )Ú
isinstancer   )r-   r.   Úmodules      r&   ú
<dictcomp>zbanks.<locals>.<dictcomp>J   s/   € ÐeÐeÐe™\˜T 6ÅjÐQWÕYcÑFdÔFdÐeˆD�&ÐeÐeÐer'   )Únamed_modules)Úmodels    r&   ÚbanksrO   I   s"   € ØeÐe¨U×-@Ò-@Ñ-BÔ-BÐeÑeÔeÐer'   c                 ó|   ‡— t          t          | ¦  «        ¦  «        Šˆfd„|                      ¦   «         D ¦   «         S )Nc                 óR   •‡— h | ]"\  Š}t          ˆfd „‰D ¦   «         ¦  «        ¯ ‰’Œ#S )c              3   óT   •K  — | ]"}‰|k    p‰                      |d z   ¦  «        V — Œ#dS )ú.N)Ú
startswith)r-   Úprefixr.   s     €r&   r/   z+excluded_names.<locals>.<setcomp>.<genexpr>O   sX   øè è € ð  =Eð  =EÐqw¸TÀVº^Ð=lÈtÏÊÐ_eÐhkÑ_kÑOlÔOlð  =Eð  =Eð  =Eð  =Eð  =Eð  =Er'   )Úany)r-   Ú_r.   Úprefixess     @€r&   ú	<setcomp>z!excluded_names.<locals>.<setcomp>O   su   øø€ ð  Fð  Fð  F‘W�T˜1½ð  =Eð  =Eð  =Eð  =Eð  |Dð  =Eñ  =Eô  =Eñ  :Eô  :Eð  FˆDð  Fð  Fð  Fr'   )r   rO   rM   )rN   rX   s    @r&   Úexcluded_namesrZ   M   sT   ø€ Ý•U˜5‘\”\Ñ"Ô"€Hð Fð  Fð  Fð  F × 3Ò 3Ñ 5Ô 5ð  Fñ  Fô  Fð  Fr'   c                   óf   — e Zd Zd„ Zd„ Zd„ Zd„ Zd„ Zd„ Zd„ Z	d„ Z
dd
„Zd„ Zd„ Zd„ Zd„ Zdd„ZdS )Ú	Residencyc                 ór  ‡ ‡‡— t          j        |¦  «        ‰ _        |‰ _        t	          |                     ¦   «         ¦  «        ‰ _        d„ ‰ j        D ¦   «         ‰ _        |‰ _        |‰ _	        d ‰ _
        d‰ _        i ‰ _        d‰ _        d ‰ _        d‰ _        d ‰ _        i ‰ _        d ‰ _        ˆ fd„t'          t)          d„ ‰ j        D ¦   «         ¦  «        ¦  «        D ¦   «         ‰ _        t-          d„ ‰ j        D ¦   «         ¦  «        r@d„ t/          ‰ j        ¦  «        D ¦   «         Šd„ Š‰ j                             ˆˆfd	„¬
¦  «         t	          t3          j        ˆ fd„‰ j        D ¦   «         ¦  «        ¦  «        ‰ _        ‰ j        D ]}t          j        ‰ ¦  «        |_        Œt;          j        ‰ d¬¦  «         d S )Nc                 ó   — i | ]
}||j         “ŒS r6   )r2   ©r-   Úbanks     r&   rL   z&Residency.__init__.<locals>.<dictcomp>W   s   € ÐFÐFÐF°4�d˜DÔ.ÐFÐFÐFr'   Fr   c                 óB   •— g | ]}‰j         D ]}||j        k     ¯||f‘ŒŒS r6   )rO   r   )r-   r<   r`   r#   s      €r&   ú
<listcomp>z&Residency.__init__.<locals>.<listcomp>e   sq   ø€ ð  ^ð  ^ð  ^¨6ÐuyÔuð  ^ð  ^Ðmqð  DJð  MQô  M]ò  D]ð  D]˜˜v˜ð  D]ð  D]ð  D]ð  D]r'   c              3   ó$   K  — | ]}|j         V — Œd S ©N)r   r_   s     r&   r/   z%Residency.__init__.<locals>.<genexpr>e   s&   è è € Ð>gÐ>gÐTX¸tÔ?OÐ>gÐ>gÐ>gÐ>gÐ>gÐ>gr'   c              3   ó(   K  — | ]}|j         d uV — Œd S rd   )r"   r_   s     r&   r/   z%Residency.__init__.<locals>.<genexpr>f   s*   è è € ÐGÐG°DˆtÔ#¨4Ð/ÐGÐGÐGÐGÐGÐGr'   c                 ó   — i | ]\  }}||“Œ	S r6   r6   )r-   ÚpositionÚitems      r&   rL   z&Residency.__init__.<locals>.<dictcomp>g   s   € ÐRÐRÐR©¨°$�T˜8ÐRÐRÐRr'   c                 ól   — | d         j         €dn%t          | d         j         | d                  ¦  «        S )Nr   g        r   )r"   Úfloat)rh   s    r&   ú<lambda>z$Residency.__init__.<locals>.<lambda>h   s5   € ¨4°¬7Ô+BÐ+J C CÕPUÐVZÐ[\ÔV]ÔVmÐnrÐstÔnuÔVvÑPwÔPw€ r'   c                 ó,   •—  ‰| ¦  «         ‰|          fS rd   r6   )rh   ÚorderÚprioritys    €€r&   rk   z$Residency.__init__.<locals>.<lambda>i   s   ø€ °°¸$±´°ÀÀtÄÐ/M€ r'   )Úkeyc              3   ó8   •K  — | ]\  }}‰j         |         V — Œd S rd   )Úsizes)r-   r`   rW   r#   s      €r&   r/   z%Residency.__init__.<locals>.<genexpr>j   s.   øè è € Ð3aÐ3aÉÈÈq°D´J¸tÔ4DÐ3aÐ3aÐ3aÐ3aÐ3aÐ3ar'   )rn   )ÚweakrefÚrefÚmanagerÚmodel_idÚlistÚvaluesrO   rq   ÚbudgetÚdynamicÚthreadÚchangingÚloadedÚnbytesÚappliedÚstartedÚpressure_ceilingÚpinnedÚ_pinned_holderÚrangeÚmaxÚrankingrV   Ú	enumerateÚsortÚ	itertoolsÚ
accumulateÚ
cumulativer   r   Úregister)	r#   rt   ru   Úexpert_banksrx   ry   r`   rm   rn   s	   `      @@r&   r   zResidency.__init__S   sÑ  øøø€ Ý”{ 7Ñ+Ô+ˆŒØ ˆŒÝ˜,×-Ò-Ñ/Ô/Ñ0Ô0ˆŒ
ØFÐF¸4¼:ÐFÑFÔFˆŒ
ØˆŒØˆŒØˆŒØˆŒØˆŒØˆŒØˆŒØˆŒØ $ˆÔØˆŒØ"ˆÔð ^ð  ^ð  ^ð  ^µU½3Ð>gÐ>gÐ\`Ô\fÐ>gÑ>gÔ>gÑ;gÔ;gÑ5hÔ5hð  ^ñ  ^ô  ^ˆŒÝÐGÐG¸D¼JÐGÑGÔGÑGÔGð 	OØRÐR½)ÀDÄLÑ:QÔ:QÐRÑRÔRˆEØwÐwˆHØŒL×ÒÐ"MÐ"MÐ"MÐ"MÐ"MÐÑNÔNÐNÝ�yÔ3Ð3aÐ3aÐ3aÐ3aÐTXÔT`Ð3aÑ3aÔ3aÑaÔaÑbÔbˆŒØ”Jð 	0ð 	0ˆDÝ%œk¨$Ñ/Ô/ˆDŒOˆOÝÔ  °Ð2Ñ2Ô2Ð2Ð2Ð2r'   c                 ó–   — |                       ¦   «          i d c| _        | _        t          j        | ¦  «         | j        D ]	}d |_        Œ
d S rd   )Úunloadr�   r‚   r   Ú
unregisterrO   r   )r#   r`   s     r&   ÚclosezResidency.closeo   sU   € Ø�Š‰ŒˆØ+-¨tÐ(ˆŒ�TÔ(ÝÔ" 4Ñ(Ô(Ð(Ø”Jð 	#ð 	#ˆDØ"ˆDŒOˆOð	#ð 	#r'   c                 óœ   — |                       ¦   «         }|j                             | j        ¦  «        }|�|                     ¦   «          d S d S rd   )rt   r   Úgetru   Úclear)r#   rt   Úgraphss      r&   Ú_clear_graphszResidency._clear_graphsv   sF   € Ø—,’,‘.”.ˆØÔ$×(Ò(¨¬Ñ7Ô7ˆØÐØ�LŠL‰NŒNˆNˆNˆNð Ðr'   c                 ó�   ‡ — ‰                       ¦   «         }t          ˆ fd„|j                             ¦   «         D ¦   «         ¦  «        S )Nc              3   ór   •K  — | ]1\  }}|‰j         k    s|                     ‰j         d z   ¦  «        ¯-|V — Œ2dS )ú/N)ru   rT   )r-   r.   Úsizer#   s      €r&   r/   z-Residency._nonexpert_bytes.<locals>.<genexpr>~   sx   øè è € ð  Tð  T™J˜D $ÐUYÐ]aÔ]jÒUjÐUjÐnr×n}Òn}ð  Cô  Lð  ORñ  Rñ  oSô  oSÐUj�4ÐUjÐUjÐUjÐUjð  Tð  Tr'   )rt   r0   Úblocks_of_modules_sizesÚitems)r#   rt   s   ` r&   Ú_nonexpert_byteszResidency._nonexpert_bytes|   s`   ø€ Ø—,’,‘.”.ˆÝð  Tð  Tð  Tð  T¨'Ô*I×*OÒ*OÑ*QÔ*Qð  Tñ  Tô  Tñ  Tô  Tð  	Tr'   c                 ót  ‡— | j         r×|                      ¦   «         }|j                             | j        d¦  «        }t          |dk    r|j        n
| j        g|¢R ¦  «        Št          ˆfd„|j         	                    ¦   «         D ¦   «         ¦  «        }t          j                             t          j                             ¦   «         ¦  «        j        |z
  t          z
  }n3| j        r*t#          d| j        |                      ¦   «         z
  ¦  «        nd}| j        dt)          j        | j        |¦  «        …         S )z­The ranked experts this residency could ever place: within its budget, or (dynamic) within the whole VRAM beside the ordinary
        blocks of the model and its co-tenants.r6   Ú*c              3   ó^   •K  — | ]'\  }}|                      d d¦  «        d         ‰v ¯#|V — Œ(dS )r˜   r   r   N)Úsplit)r-   r.   r™   Ú	model_idss      €r&   r/   z&Residency.placeable.<locals>.<genexpr>‡   sW   øè è € ð  Að  A¡J D¨$Ð]a×]gÒ]gÐhkÐmnÑ]oÔ]oÐpqÔ]rÐvÐ]Ð]˜4Ð]Ð]Ð]Ð]ð  Að  Ar'   r   N)ry   rt   Úcotenants_mapr’   ru   ÚsetÚmodelsr0   rš   r›   r   ÚcudaÚget_device_propertiesÚcurrent_deviceÚtotal_memoryÚMARGINrx   r„   rœ   r…   ÚbisectÚbisect_rightrŠ   )r#   rt   Ú	cotenantsÚordinaryÚtargetr¡   s        @r&   Ú	placeablezResidency.placeable€   s3  ø€ ð Œ<ð 	YØ—l’l‘n”nˆGØÔ-×1Ò1°$´-ÀÑDÔDˆIÝ¨i¸3Ò.>Ð.>˜GœN˜NÀTÄ]ÐD_ÐU^ÐD_ÐD_Ñ`Ô`ˆIÝð  Að  Að  Að  A°'Ô2Q×2WÒ2WÑ2YÔ2Yð  Añ  Aô  Añ  Aô  AˆHÝ”Z×5Ò5µe´j×6OÒ6OÑ6QÔ6QÑRÔRÔ_ÐbjÑjÕmsÑsˆFˆFàFJÄkÐX•S˜˜DœK¨$×*?Ò*?Ñ*AÔ*AÑAÑBÔBÐBÐWXˆFØŒ|ÐI�VÔ0°´À&ÑIÔIÐIÔJÐJr'   c                 ó  ‡‡	— |                       ¦   «         }|sdS t          j        ¦   «         Št          |¦  «        D ]U\  Š	\  }}|j        D ]E}‰                     ‰	› d|› �t          j        t          ||¦  «        |         d¬¦  «        ¦  «         ŒFŒV |‰| j        › d�¦  «         t          |¦  «        D ]I\  Š	}ˆˆ	fd„|d         j        D ¦   «         }d„ | 	                    ¦   «         D ¦   «         }|r
|| j
        |<   ŒJ‰| _        dS )	aC  Pinned copies of the experts this residency could place, in ranking order and as far as the reserved RAM left by the models
        allows (pin_module(module, name) pins a module's parameters and stops there): their uploads become plain transfers. The banks
        keep their pageable sources, which the CPU backend reads.NrW   F)Úrequires_gradz expertsc           	      óF   •— i | ]}|t          ‰‰› d |› �¦  «        j        “ŒS )rW   )r*   Údata)r-   r.   ÚholderÚindexs     €€r&   rL   z!Residency.pin.<locals>.<dictcomp>š   s6   ø€ ÐmÐmÐmÈ�t�W V°Ð->Ð->¸Ð->Ð->Ñ?Ô?ÔDÐmÐmÐmr'   r   c                 óB   — i | ]\  }}|                      ¦   «         ¯||“ŒS r6   )Ú	is_pinned)r-   r.   Útensors      r&   rL   z!Residency.pin.<locals>.<dictcomp>›   s/   € Ð^Ð^Ð^©¨¨fÈ6×K[ÒK[ÑK]ÔK]Ð^�t˜VÐ^Ð^Ð^r'   )r¯   r   ÚModuler†   r   Úregister_parameterÚ	Parameterr*   ru   r›   r�   r‚   )
r#   Ú
pin_moduleÚexpertsr`   r<   r.   rh   r=   r´   rµ   s
           @@r&   ÚpinzResidency.pin�   sN  øø€ ð —.’.Ñ"Ô"ˆØð 	ØˆFÝ”‘”ˆÝ%.¨wÑ%7Ô%7ð 	}ð 	}Ñ!ˆE‘>�D˜&ØÔ0ð }ð }�Ø×)Ò)¨UÐ*;Ð*;°TÐ*;Ð*;½R¼\Í'ÐRVÐX\ÑJ]ÔJ]Ð^dÔJeÐuzÐ={Ñ={Ô={Ñ|Ô|Ð|Ð|ð}àˆ
�6˜dœmÐ5Ð5Ð5Ñ6Ô6Ð6Ý$ WÑ-Ô-ð 	,ð 	,‰KˆE�4ØmÐmÐmÐmÐmÐQUÐVWÔQXÔQlÐmÑmÔmˆGØ^Ð^¸¿º¹¼Ð^Ñ^Ô^ˆGØð ,Ø$+�”˜DÑ!øØ$ˆÔÐÐr'   c                 ó^  ‡‡— |                       ¦   «         }|j                             | j        d¦  «        }|dk    r|j        n
| j        g|¢R }|D ][Š‰|j        vrŒ‰|j        vr dS |j        ‰         Št          ˆˆfd„|j         	                    ¦   «         D ¦   «         ¦  «        r dS Œ\dS )Nr6   rž   Fc              3   óŽ   •K  — | ]?\  }}|                      ‰d z   ¦  «        ¯|¯|                     d d¦  «        d         ‰vV — Œ@dS )r˜   r   N)rT   r    )r-   r.   r™   ru   Ú	preloadeds      €€r&   r/   z1Residency._nonexperts_resident.<locals>.<genexpr>ª   sž   øè è € ð  að  a¹j¸dÀDÐw{÷  xGò  xGð  HPð  SVñ  HVñ  xWô  xWð  að  \`ð  a�4—:’:˜c 1Ñ%Ô% aÔ(°	Ð9ð  að  að  að  að  að  ar'   T)
rt   r¢   r’   ru   r¤   Úactive_models_idsÚpreloaded_blocks_per_modelrV   rš   r›   )r#   rt   r¬   r¡   ru   rÁ   s       @@r&   Ú_nonexperts_residentzResidency._nonexperts_resident    sô   øø€ Ø—,’,‘.”.ˆØÔ)×-Ò-¨d¬m¸RÑ@Ô@ˆ	Ø&/°3Ò&6Ð&6�G”N�N¸T¼]Ð<WÈYÐ<WÐ<Wˆ	Ø!ð 	ð 	ˆHØ˜wœ~Ð-Ð-ØØ˜wÔ8Ð8Ð8Ø�u�uØÔ:¸8ÔDˆIÝð  að  að  að  að  aÈGÔLk×LqÒLqÑLsÔLsð  añ  aô  añ  aô  að Ø�u�uðàˆtr'   c                 ó^   — |                       ¦   «         sdS |                      d¬¦  «         dS )a  Placement requested outside graph recording (runtime.place_experts): grows the GPU copies to what the budget or the free VRAM
        allows without waiting for a measuring cycle, never shrinks them (pressure does). False while the ordinary blocks are not all
        resident yet.FT)Ú	grow_only)rÄ   Úpreparer1   s    r&   ÚplacezResidency.place®   s6   € ð ×(Ò(Ñ*Ô*ð 	Ø�5Ø�Š˜tˆÑ$Ô$Ð$Øˆtr'   Fc                 ó~  ‡ ‡	— t          j        ¦   «         ‰ _        ‰ j        rdS ‰                      ¦   «         s‰                      d¬¦  «         dS ‰ j        r‰ j        s|s	d‰ _        dS ‰ j        r*t          d‰ j        ‰  
                    ¦   «         z
  ¦  «        nd}‰ j        r¼t          j        dk    r9t          j        ¦   «         }t          j        ¦   «         |d         z   |d         z
  }n$t          j                             ¦   «         d         }t          d‰ j        |z   t&          z
  t)          d	t          j                             ¦   «         z  ¦  «        z
  ¦  «        }‰ j        �t/          |‰ j        ¦  «        }t1          j        ‰ j        |¦  «        }|‰ j        k    s|r|t9          ‰ j        ¦  «        k    rdS ‰ j        r6‰ j        r/t=          |r‰ j        |d
z
           nd‰ j        z
  ¦  «        dk     rdS ‰ j        d|…         }t9          |¦  «        t9          ‰ j        ¦  «        k    r$tA          ˆ fd„|D ¦   «         ¦  «        r	|‰ _        dS d‰ _        	 ‰  !                    ¦   «          t          j         "                    ¦   «          tG          |¦  «        Š	ˆ	fd„tI          ‰ j        ¦  «        D ¦   «         D ]} ‰ j%        |Ž  Œˆ fd„|D ¦   «         }|r‰  &                    |¦  «         tN           (                    |¦  «        ‰ _        |‰ _        ‰  )                    ¦   «         j*        d
k    rStW          j,        d‰ j-        › dt9          ‰ j        ¦  «        › dt9          ‰ j        ¦  «        › d‰ j        dz  d›d�	¦  «         d‰ _        ‰  .                    ¦   «          dS # d‰ _        ‰  .                    ¦   «          w xY w)zCAt a model/cycle boundary, after the ordinary preload has priority.NF)ÚresetTr   ÚvmmÚreservedÚ	allocatedgš™™™™™©?r   i   c              3   ó*   •K  — | ]}|‰j         v V — Œd S rd   ©r|   ©r-   rh   r#   s     €r&   r/   z$Residency.prepare.<locals>.<genexpr>Ô   s*   øè è € Ð2ZÐ2ZÈ4°4¸4¼;Ð3FÐ2ZÐ2ZÐ2ZÐ2ZÐ2ZÐ2Zr'   c                 ó   •— g | ]}|‰v¯|‘Œ	S r6   r6   )r-   rh   Úkeeps     €r&   rb   z%Residency.prepare.<locals>.<listcomp>Ü   s#   ø€ ÐTÐTÐT $À4ÈtÐCSÐCS˜ÐCSÐCSÐCSr'   c                 ó&   •— g | ]}|‰j         v¯|‘ŒS r6   rÏ   rÐ   s     €r&   rb   z%Residency.prepare.<locals>.<listcomp>Þ   s%   ø€ ÐLÐLÐL $°DÀÄÐ4KÐ4K˜Ð4KÐ4KÐ4Kr'   zExpert Residency - z: r˜   z Experts on GPU (i   z.0fz! MiB); CPU Experts Are Not Pinned)/Ú	threadingÚ	get_identrz   r{   rÄ   rŽ   ry   r   rx   r„   rœ   r   ÚactiveÚstatsÚroomr   r¥   Úmem_get_infor}   r©   r   Úmemory_allocatedr€   Úminrª   r«   rŠ   r~   Úlenr|   Úabsr…   Úallr•   Úsynchronizer£   ÚreversedÚ_dropÚ_uploadÚdictÚfromkeysrt   ÚverboseLevelr   Úwriteru   Ú_placement_changed)
r#   rÆ   r®   ÚstateÚfreeÚcountÚwantedrh   Ú	additionsrÒ   s
   `        @r&   rÇ   zResidency.prepare·   sï  øø€ åÔ)Ñ+Ô+ˆŒØŒ=ð 	ØˆFØ×(Ò(Ñ*Ô*ð 	Ø�KŠK˜eˆKÑ$Ô$Ð$ØˆFØŒ<ð 	 ¤ð 	°Yð 	àˆDŒLØˆFØBFÄ+ÐT•�Q˜œ d×&;Ò&;Ñ&=Ô&=Ñ=Ñ>Ô>Ð>ÐSTˆØŒ<ð 	eÝÔ 5Ò(Ð(Ý!œÑ)Ô)�Ý ”~Ñ'Ô'¨%°
Ô*;Ñ;¸eÀKÔ>PÑP��å”z×.Ò.Ñ0Ô0°Ô3�Ý˜˜DœK¨$Ñ.µÑ7½#¸dÅUÄZ×E`ÒE`ÑEbÔEbÑ>bÑ:cÔ:cÑcÑdÔdˆFØÔ Ð,Ý˜ Ô!6Ñ7Ô7ˆFåÔ# D¤O°VÑ<Ô<ˆØ�D”LÒ Ð  IÐ °%½3¸t¼{Ñ;KÔ;KÒ2KÐ2KØˆFØŒ<ð 	˜DœKð 	­CÈuÐ1[°´ÀÈÁÔ1KÐ1KÐZ[Ð_cÔ_jÑ0jÑ,kÔ,kÐowÒ,xÐ,xØˆFØ”˜f˜u˜fÔ%ˆÝˆv‰;Œ;�#˜dœkÑ*Ô*Ò*Ð*­sÐ2ZÐ2ZÐ2ZÐ2ZÐSYÐ2ZÑ2ZÔ2ZÑ/ZÔ/ZÐ*Ø ˆDŒLØˆFØˆŒð	&Ø×ÒÑ Ô Ð ÝŒJ×"Ò"Ñ$Ô$Ð$Ý�v‘;”;ˆDØTÐTÐTÐT­(°4´;Ñ*?Ô*?ÐTÑTÔTð "ð "�Ø�”
˜DÐ!Ð!Ð!ØLÐLÐLÐL¨&ÐLÑLÔLˆIØð (Ø—’˜YÑ'Ô'Ð'ÝŸ-š-¨Ñ/Ô/ˆDŒKØ ˆDŒLØ�|Š|‰~Œ~Ô*¨aÒ/Ð/Ý”
ð  x°´ð  xð  xÅ#ÀdÄkÑBRÔBRð  xð  xÕUXÐY]ÔYeÑUfÔUfð  xð  xÐy}ô  zEð  IPñ  zQð  xð  xð  xð  xñ  yô  yð  yà!ˆDŒMØ×#Ò#Ñ%Ô%Ð%Ð%Ð%øð "ˆDŒMØ×#Ò#Ñ%Ô%Ð%Ð%øøøs   ÉD)N ÎN<c           
      ót  ‡ ‡‡— t          t          dt          ‰ j                             ¦   «         ¦  «        z  ¦  «        Šˆfd„t	          d¦  «        D ¦   «         }t
          j                             ¦   «         t
          j                             ¦   «         g}d\  }}ˆˆ fd„}‰r%t          t           |d¦  «        ¦  «        ¦  «         	 t          t          ‰d‰ j        › �dd‰                      ¦   «         j        d	k     ¬
¦  «        ¦  «        D �]:\  }\  }}	t          j        ¦   «          |d	z   t#          ‰¦  «        k     r(t          t           ||d	z   ¦  «        ¦  «        ¦  «         |‰ j        |         z   dt#          |j        ¦  «        z  z   ‰k    r;||                              ¦   «          d	|z
  d}}||                              ¦   «          i }
‰ j                             ||	fi ¦  «        }|j        D �]}||v r;t          j        ||         d¬¦  «                             ||         d¬¦  «        |
|<   ŒBt3          ||¦  «        |	         }|                     ¦   «         |                     ¦   «         z  }||         |||z   …                              |j        ¦  «                             |j        ¦  «        }|                     |¦  «         t          j        |d¬¦  «                             |d¬¦  «        |
|<   || dz   dz  z  }�Œ|                     |	|
¦  «         d‰ j         ||	f<   ‰ xj!        ‰ j        |         z  c_!        �Œ<	 ||                              ¦   «          ||                              ¦   «          dS # ||                              ¦   «          ||                              ¦   «          w xY w)aÁ  GPU copies of new experts, from their pinned copies (pin()) or else through two pinned buffers: copying the next experts from
        their pageable (mmap) sources overlaps the transfer of the previous ones. These copies are only made when the placement changes.
        Sources that a drive must read are read ahead in large requests, the next expert while the current one is copied (the parallel CPU
        copy would fault them in 64 KiB reads).é   c                 óT   •— g | ]$}t          j        ‰t           j        d d¬¦  «        ‘Œ%S )r   T)r   r   Ú
pin_memory)r   Úemptyr   )r-   rW   r™   s     €r&   rb   z%Residency._upload.<locals>.<listcomp>ï   s0   ø€ ÐiÐiÐiÐ[\•5”;˜t­5¬;¸uÐQUÐVÑVÔVÐiÐiÐir'   )r   r   c                 óH   •‡‡— ‰|          \  ŠŠˆˆˆfd„‰j         D ¦   «         S )Nc                 óô   •— g | ]t}|‰j                              ‰‰fi ¦  «        v¯!t          ‰|¦  «        ‰         }|                     ¦   «         |                     ¦   «         |                     ¦   «         z  f‘ŒuS r6   )r�   r’   r*   Údata_ptrr+   r,   )r-   r.   r¸   r`   r<   r#   s      €€€r&   rb   z6Residency._upload.<locals>.sources.<locals>.<listcomp>õ   s„   ø€ ð uð uð uÐTXØ 4¤;§?¢?°D¸&°>À2Ñ#FÔ#FÐFÐFÕV]Ð^bÐdhÑViÔViÐjpÔVqÈ6ð —_’_Ñ&Ô&¨¯ª©¬¸×9LÒ9LÑ9NÔ9NÑ(NÐOØFÐFÐFr'   )r   )rµ   r`   r<   rì   r#   s    @@€€r&   Úsourcesz"Residency._upload.<locals>.sourcesó   sU   øøø€ Ø$ UÔ+‰LˆD�&ðuð uð uð uð uð uÐ\`Ô\tð uñ uô uð ur'   r   zPreloading Experts - r<   gUUUUUUÕ?r   )ÚdescÚunitÚminintervalÚdisableé@   r¥   )r   T)Únon_blockingN)"r„   ÚUPLOAD_STAGINGrq   rw   rƒ   r   r¥   ÚEventr   r
   r†   r   ru   rt   rå   r	   Ú
checkpointrÜ   r   Úrecordrß   r�   r’   Ú
empty_likeÚcopy_r*   r+   r,   Úviewr   Úshaper>   r|   r}   )r#   rì   ÚstagingÚ	transfersÚslotÚoffsetrõ   rµ   r`   r<   r=   Úkeptr.   Úsourcer}   r�   r™   s   ``              @r&   râ   zResidency._uploadé   sñ  øøø€ õ
 •> 1¥s¨4¬:×+<Ò+<Ñ+>Ô+>Ñ'?Ô'?Ñ#?Ñ@Ô@ˆØiÐiÐiÐiÕ`eÐfgÑ`hÔ`hÐiÑiÔiˆÝ”Z×%Ò%Ñ'Ô'­¬×)9Ò)9Ñ);Ô);Ð<ˆ	Ø‰ˆˆfð	uð 	uð 	uð 	uð 	uð 	uð
 ð 	5Ý�K¨¨°©
¬
Ñ3Ô3Ñ4Ô4Ð4ð	*Ý)2µ4¸	ÐHoÐ`dÔ`mÐHoÐHoÐv~ð  MRð  \`÷  \hò  \hñ  \jô  \jô  \wð  z{ò  \{ð  4|ñ  4|ô  4|ñ  *}ô  *}ð 0ñ 0Ñ%�‘~˜˜fÝÔ"Ñ$Ô$Ð$Ø˜1‘9�s 9™~œ~Ò-Ð-Ý#¥K°°¸À¹	Ñ0BÔ0BÑ$CÔ$CÑDÔDÐDØ˜DœJ tÔ,Ñ,¨rµC¸Ô8PÑ4QÔ4QÑ/QÑQÐTXÒXÐXØ˜d”O×*Ò*Ñ,Ô,Ð,Ø#$ t¡8¨Q˜&�DØ˜d”O×/Ò/Ñ1Ô1Ð1Ø�Ø”{—’¨¨f ~°rÑ:Ô:�Ø Ô4ð 	4ñ 	4�DØ˜t�|�|Ý(-Ô(8¸¸d¼ÈFÐ(SÑ(SÔ(S×(YÒ(YÐZ^Ð_cÔZdÐswÐ(YÑ(xÔ(x˜ ™Ø Ý$ T¨4Ñ0Ô0°Ô8�FØ#Ÿ\š\™^œ^¨f×.AÒ.AÑ.CÔ.CÑC�FØ$ Tœ]¨6°&¸6±/Ð+AÔB×GÒGÈÌÑUÔU×ZÒZÐ[aÔ[gÑhÔh�FØ—L’L Ñ(Ô(Ð(Ý$)Ô$4°VÀFÐ$KÑ$KÔ$K×$QÒ$QÐRXÐgkÐ$QÑ$lÔ$l�G˜D‘MØ  ¨2¡Ð.°Ñ3Ñ3�F‘FØ—’˜F GÑ,Ô,Ð,Ø.2�”˜T 6˜NÑ+Ø�”˜tœz¨$Ô/Ñ/�”‘ð-0ð0 �dŒO×"Ò"Ñ$Ô$Ð$Ø�dŒO×'Ò'Ñ)Ô)Ð)Ð)Ð)øð �dŒO×"Ò"Ñ$Ô$Ð$Ø�dŒO×'Ò'Ñ)Ô)Ð)Ð)øøøs   Ã
J N Î6N7c                 óˆ   — |                      |¦  «         | j        ||f= | xj        | j        |         z  c_        d | _        d S rd   )r@   r|   r}   rq   r~   )r#   r`   r<   s      r&   rá   zResidency._drop  sD   € Ø�Š�FÑÔÐØŒK˜˜v˜Ð'ØˆŒ�t”z $Ô'Ñ'ˆŒØˆŒˆˆr'   c                 óh   — |                       ¦   «         j        }|�|                     ¦   «          d S d S rd   )rt   Úauto_preloadÚon_expert_change)r#   r  s     r&   rç   zResidency._placement_changed  s8   € Ø—|’|‘~”~Ô2ˆØÐ#Ø×)Ò)Ñ+Ô+Ð+Ð+Ð+ð $Ð#r'   c                 ó\  — |dk    sT| j         sM| j        t          j        ¦   «         k    s1t	          j        ¦   «         st          j                             ¦   «         rdS d„ t          | j
        ¦  «        D ¦   «         }|sdS d| _         	 |                      ¦   «          t          j                             ¦   «          d}|D ];\  }}|| j        |         z  }|                      ||¦  «         ||t          z   k    r nŒ<| j        | _        |                      ¦   «          |d| _         S # d| _         w xY w)Nr   c                 ó*   — g | ]\  }}|j         °||f‘ŒS r6   )r   )r-   r`   r<   s      r&   rb   z)Residency.on_pressure.<locals>.<listcomp>$  s)   € Ð`Ð`Ð`¡l d¨FÐSWÔS_Ð`�D˜&�>Ð`Ð`Ð`r'   TF)r{   rz   rÔ   rÕ   r   Ú	recordingr   r¥   Úis_current_stream_capturingrà   r|   r•   rß   rq   rá   r©   r}   r€   rç   )r#   ÚneedÚvictimsÚfreedr`   r<   s         r&   Úon_pressurezResidency.on_pressure!  sL  € Ø�1Š9ˆ9˜œˆ9¨¬½	Ô8KÑ8MÔ8MÒ)MÐ)MÕQ\ÔQfÑQhÔQhÐ)MÕlqÔlv÷  mSò  mSñ  mUô  mUÐ)MØ�1Ø`Ð`µh¸t¼{Ñ6KÔ6KÐ`Ñ`Ô`ˆØð 	Ø�1ØˆŒð	"Ø×ÒÑ Ô Ð ÝŒJ×"Ò"Ñ$Ô$Ð$ØˆEØ 'ð ð ‘��fØ˜œ DÔ)Ñ)�Ø—
’
˜4 Ñ(Ô(Ð(Ø˜D¥6™MÒ)Ð)Ø�Eð *à$(¤KˆDÔ!Ø×#Ò#Ñ%Ô%Ð%Øà!ˆDŒMˆMø˜EˆDŒMÐ!Ð!Ð!Ð!s   ÂBD" Ä"	D+Tc                 ó\  — | j         rt|                      ¦   «          t          j                             ¦   «          t          t          | j         ¦  «        ¦  «        D ]} | j        |Ž  Œ|                      ¦   «          |r,d| _	        d | _
        | j        D ]}|                     ¦   «          Œd S d S )NF)r|   r•   r   r¥   rß   rv   rà   rá   rç   r   r€   rO   r7   )r#   rÊ   rh   r`   s       r&   rŽ   zResidency.unload7  sÈ   € ØŒ;ð 	&Ø×ÒÑ Ô Ð ÝŒJ×"Ò"Ñ$Ô$Ð$Ý�X d¤kÑ2Ô2Ñ3Ô3ð "ð "�Ø�”
˜DÐ!Ð!Ð!Ø×#Ò#Ñ%Ô%Ð%Øð 	/Ø ˆDŒLØ$(ˆDÔ!Øœ
ð /ð /�Ø×,Ò,Ñ.Ô.Ð.Ð.ð		/ð 	/ð/ð /r'   N)F)T)rA   rB   rC   r   r�   r•   rœ   r¯   r¾   rÄ   rÈ   rÇ   râ   rá   rç   r  rŽ   r6   r'   r&   r\   r\   R   sø   € € € € € ð3ð 3ð 3ð8#ð #ð #ðð ð ðTð Tð TðKð Kð Kð%ð %ð %ð&ð ð ðð ð ð0&ð 0&ð 0&ð 0&ðd+*ð +*ð +*ðZð ð ð,ð ,ð ,ð
"ð "ð "ð,/ð /ð /ð /ð /ð /r'   r\   )Ú__doc__rª   rE   rˆ   rÔ   rr   r   r   r   Ú r   r   r   r	   Ú
read_aheadr
   r   r©   rü   r¹   r   rO   rZ   r\   r6   r'   r&   ú<module>r     sG  ððð ð €€€Ø Ð Ð Ð Ø Ð Ð Ð Ø Ð Ð Ð Ø €€€à €€€Ø Ð Ð Ð Ð Ð Ø Ð Ð Ð Ð Ð à >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ø 4Ð 4Ð 4Ð 4Ð 4Ð 4Ð 4Ð 4à	€Ø€ð,$ð ,$ð ,$ð ,$ð ,$�”ñ ,$ô ,$ð ,$ð^fð fð fðFð Fð Fð
p/ð p/ð p/ð p/ð p/ñ p/ô p/ð p/ð p/ð p/r'   