o
    $i00                     @   s  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ZddlmZm	Z	 ddl
mZmZ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 dd	lmZ dd
lmZ eeZdedefddZ dedee fddZ!G dd dZ"G dd deZ#dS )z'Util class to install packages via uv.
    N)create_taskget_running_loop)DictListOptional)try_to_create_directory)dependency_utilsvirtualenv_utils)Protocol	parse_uri)RuntimeEnvPlugincheck_output_cmd)get_directory_size_bytesuv_dictreturnc                 C   s&   t j| dd}t|d }|S )z=Get a deterministic hash value for `uv` related runtime envs.T)	sort_keyszutf-8)jsondumpshashlibsha1encode	hexdigest)r   serialized_uv_spechash_val r   X/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/ray/_private/runtime_env/uv.py_get_uv_hash   s   r   runtime_envc                 C   sn   |  d}|dur3t|trdt|d }|S t|tr(dtt|dd }|S tdt|j dd}|S )zAReturn `"uv://<hashed_dependencies>"`, or None if no GC required.uvNzuv://)r   )packagesz?uv field received by RuntimeEnvAgent must be list or dict, not .)get
isinstancedictr   list	TypeErrortype__name__)r   r   urir   r   r   get_uri   s   



r*   c                
   @   s   e Zd Zefdedddeej fddZdeded	e	dejfd
dZ
dedede	dejdef
ddZdededejddfddZdedee ded	edejf
ddZdd Zdd ZdS )UvProcessor
target_dirr   
RuntimeEnvloggerc                 C   s   zdd l }W n ty   tdtj dw |d| || _tj	| jd| _
|| _|| _| j | _tj | _| j| j  d S )Nr   zPlease install virtualenv `z4 -m pip install virtualenv`to enable uv runtime env.z!Setting up uv for runtime_env: %sexec_cwd)
virtualenvImportErrorRuntimeErrorsys
executabledebug_target_dirospathjoin	_exec_cwd_runtime_env_logger	uv_config
_uv_configenvironcopy_uv_envupdateenv_vars)selfr,   r   r.   r0   r   r   r   __init__2   s"   zUvProcessor.__init__r8   cwdpip_envc           	         sb   t |}t |}dtf fdd}|ddddd| g}|d	| t||||d
I dH  dS )zjBefore package install, make sure the required version `uv` (if specifieds)
        is installed.
        r   c                     s     j dd} | rd|  S dS )z,Get `uv` executable with version to install.
uv_versionNr   )r>   r"   )rH   rD   r   r   _get_uv_exec_to_installU   s   
z8UvProcessor._install_uv.<locals>._get_uv_exec_to_install-mpipinstallz--disable-pip-version-checkz--no-cache-dirzInstalling package uv to %sr.   rF   envN)r	   get_virtualenv_pathget_virtualenv_pythonstrinfor   )	rD   r8   rF   rG   r.   virtualenv_pathpythonrJ   uv_install_cmdr   rI   r   _install_uvL   s   

		zUvProcessor._install_uvrO   r   c                    sJ   t |}|dddg}zt||||dI dH  W dS  ty$   Y dS w )z6Check and return the existence of `uv` in virtual env.rK   r   versionrN   NTF)r	   rQ   r   	Exception)rD   r8   rF   rO   r.   rU   check_existence_cmdr   r   r   _check_uv_existencei   s   
zUvProcessor._check_uv_existencerU   Nc                    s(   |ddddg}t |||dI dH  dS )zCheck virtual env dependency compatibility.
        If any incompatibility detected, exception will be thrown.

        param:
            python: the path for python executable within virtual environment.
        rK   r   rL   check)r.   rF   Nr   )sefrU   rF   r.   cmdr   r   r   	_uv_check}   s   zUvProcessor._uv_checkuv_packagesc                    s   t |}t |}t||}| ||||I dH }	|	r'| jdddur2| ||||I dH  t	 }
|

dtj||I dH  |ddddd|g}| jdd	g}|rX||7 }|d
| t||||dI dH  | jddr|| |||I dH  dS dS )z*Install required python packages via `uv`.NrH   rK   r   rL   rM   z-ruv_pip_install_optionsz
--no-cachez$Installing python requirements to %srN   uv_checkF)r	   rP   rQ   r   get_requirements_filer[   r>   r"   rW   r   run_in_executorgen_requirements_txtrS   r   r_   )rD   r8   r`   rF   rG   r.   rT   rU   requirements_file	uv_existslooprV   uv_opt_listr   r   r   _install_uv_packages   s6   
	


	
z UvProcessor._install_uv_packagesc              	      s   | j }| j}| jd }tj| jdd zEt|| j|I d H  t|}t	
|| j|4 I d H  | ||| j| j|I d H  W d   I d H  W d S 1 I d H sRw   Y  W d S  tys   |d| tj|dd |d  w )Nr    T)exist_okz Delete incomplete virtualenv: %s)ignore_errorszFailed to install uv packages.)r6   r<   r>   r7   makedirsr:   r	   create_or_get_virtualenvrQ   r   	check_rayrj   rA   rY   rS   shutilrmtree	exception)rD   r8   r.   r`   rU   r   r   r   _run   s2   


2	
zUvProcessor._runc                 C   s   |    S N)rs   	__await__rI   r   r   r   ru      s   zUvProcessor.__await__)r(   
__module____qualname__default_loggerrR   r   loggingLoggerrE   r$   rW   boolr[   r_   r   r   rj   rs   ru   r   r   r   r   r+   1   sV    



4r+   c                   @   s   e Zd ZdZdefddZdedefddZd	d
dee fddZe	fdede
ej defddZe	fded	d
ddde
ej def
ddZe	fdee d	d
dddejfddZdS )UvPluginr   resources_dirc                 C   s0   t j|d| _i | _i | _i | _t| j d S )Nr   )r7   r8   r9   _uv_resource_dir_creating_task_create_locks_created_hash_bytesr   )rD   r}   r   r   r   rE      s
   zUvPlugin.__init__r   r   c                 C   s   t j| j|S )zGenerate a path from the hash of a uv spec.

        Example output:
            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
                /uv/ray-9a7972c3a75f55e976e620484f58410c920db091
        )r7   r8   r9   r~   )rD   r   r   r   r   _get_path_from_hash   s   zUvPlugin._get_path_from_hashr   r-   c                 C   s   |  }|r	|gS g S )zCReturn the uv URI from the RuntimeEnv if it exists, else return [].)uv_uri)rD   r   r   r   r   r   get_uris   s   zUvPlugin.get_urisr)   r.   c           	   
   C   s   | d| t|\}}|tjkrtd| d| | j|d}|dur*|  | j|= | 	|}t
|}| j|= zt| W |S  tyc } z|d| dt|  W Y d}~dS d}~ww )z2Delete URI and return the number of bytes deleted.zGot request to delete uv URI %szBUvPlugin can only delete URIs with protocol uv. Received protocol z, URI NzError when deleting uv env z: r   )rS   r   r
   UV
ValueErrorr   popcancelr   r   r   r   rp   rq   OSErrorwarningrR   )	rD   r)   r.   protocolr   taskuv_env_pathlocal_dir_sizeer   r   r   
delete_uri   s2   

zUvPlugin.delete_uricontextRuntimeEnvContextc           	   	      s     sdS t|\}  fdd}|jvr&t j|< j| 4 I d H ?  jv rDj  W  d   I d H  S t|  j < }|	 fdd |I d H }|j < |W  d   I d H  S 1 I d H stw   Y  d S )Nr   c                     s.   t  I d H  t } | d tI d H S rt   )r+   r   rd   r   )rh   )r.   r   r,   r   r   _create_for_hash'  s   

z)UvPlugin.create.<locals>._create_for_hashc                    s   j  d S rt   )r   r   )_)r   rD   r   r   <lambda>;  s    z!UvPlugin.create.<locals>.<lambda>)
has_uvr   r   r   asyncioLockr   r   r   add_done_callback)	rD   r)   r   r   r.   r   r   r   uv_dir_bytesr   )r   r.   r   rD   r,   r   create  s$   




0zUvPlugin.createurisc           
      C   st   |  sd S |d }t|\}}| |}t|}	tj|	s+td| d| d|	|_	| j
t|7  _
d S )Nr   zLocal directory z	 for URI zm does not exist on the cluster. Something may have gone wrong while installing the runtime_env `uv` packages.)r   r   r   r	   rQ   r7   r8   existsr   py_executablecommand_prefixget_virtualenv_activate_command)
rD   r   r   r   r.   r)   r   r   r,   virtualenv_pythonr   r   r   modify_context@  s   


zUvPlugin.modify_contextN)r(   rv   rw   namerR   rE   r   r   r   rx   r   ry   rz   intr   r   r   r   r   r   r   r|      sF    
	
#
+r|   )$__doc__r   r   r   ry   r7   rp   r3   r   r   typingr   r   r   ray._common.utilsr   ray._private.runtime_envr   r	   "ray._private.runtime_env.packagingr
   r   ray._private.runtime_env.pluginr   ray._private.runtime_env.utilsr   ray._private.utilsr   	getLoggerr(   rx   rR   r   r*   r+   r|   r   r   r   r   <module>   s,    
 /