o `۷i;E@s`dZddlZddlZddlZddlZddlZddlZddlZddlZddl Z ddl m Z ddl Z ddl Z ddlZddlmZddlZddlmZddZdd Zd d Zd d ZddZddZddZddZddZddZddZddZ dBdd Z e!d!kre\Z"Z#Z$Z%Z&Z'e e"Z"ee"e(d"e"e(d#e#e(d$e$e(d%e%e)e"*Z+e+d&d'e,e-ee+d&<e(d(e(d)e'e'ree+e'e(d(e(d*e&ee&Z.e(d+e.e.ree+e.e+/d,i/d-Z0d.e+d,d/<e0d0krenje0d1kree+d,1d2dnYe0d3kr]eej2gd4ej3d5d6j45d76Z7e(d8e7ej2gd9ej3d5d6j45d76Z8e(d:e8d;e8e+d,d<<ne0d=krke(d(e(d>n e(d(e(d?e9dej:d@dA%Z;e;d7e;?e e;j@Z"e e+e"e#e$e%WddS1swYdSdS)Ca This script automates the process of launching and verifying a Ray cluster using a given cluster configuration file. It also handles cluster cleanup before and after the verification process. The script requires one command-line argument: the path to the cluster configuration file. Usage: python launch_and_verify_cluster.py [--no-config-cache] [--retries NUM_RETRIES] [--num-expected-nodes NUM_NODES] [--docker-override DOCKER_OVERRIDE] [--wheel-override WHEEL_OVERRIDE] N)Path)storage)RAYcCstjdd}|jdddd|jdtdd d |jd td d d |jdgdddd|jdtddd |jdtdd|}|jdkrN|jdkrNJd|j|j |j |j |j|jfS)z Check command line arguments and return the cluster configuration file path, the number of retries, the number of expected nodes, and the value of the --no-config-cache flag. zLaunch and verify a Ray cluster) description--no-config-cache store_truez3Pass the --no-config-cache flag to Ray CLI commands)actionhelpz --retriesz;Number of retries for verifying Ray is running (default: 3))typedefaultr z--num-expected-nodesz9Number of nodes for verifying Ray is running (default: 1)z--docker-override)disablelatestnightlycommitrzAOverride the docker image used for the head node and worker nodes)choicesr r z--wheel-overridez:Override the wheel used for the head node and worker nodescluster_configz&Path to the cluster configuration file)r r z%Cannot override both docker and wheel) argparseArgumentParser add_argumentintstr parse_argsdocker_overridewheel_overriderretriesno_config_cachenum_expected_nodes)parserargsr"^/home/ubuntu/vllm_env/lib/python3.10/site-packages/ray/autoscaler/launch_and_verify_cluster.pycheck_arguments!sX r$cCsj|dkrdS|dkr dS|dkr3tdtjr&dtjdtjd d d Std tjtd d S)a  Get the docker image to use for the head node and worker nodes. Args: docker_override: The value of the --docker-override flag. Returns: The docker image to use for the head node and worker nodes, or None if not applicable. rzrayproject/ray:latest-py310rzrayproject/ray:nightly-py310rz^[0-9]+.[0-9]+.[0-9]+$zrayproject/ray:.Nz-py310zGError: docker image is only available for release version, but we get: r )rematchray __version__ __commit__printsysexit)rr"r"r#get_docker_imageXs  r/cCs6|r t|tjstd|tddSdS)z Check if the provided file path is valid and readable. Args: file_path: The path of the file to check. Raises: SystemExit: If the file is not readable or does not exist. z/Error: Cannot read cluster configuration file: r N)is_fileosaccessR_OKr,r-r.) file_pathr"r"r# check_filess r5cCs*|dg}|d|d||d<dS)Nsetup_commandsz9pip3 uninstall -y ray && pip3 install -U "ray[default] @ ")getappend) config_yaml wheel_urlr6r"r"r#override_wheels_urls   r<cCsT|di}||d<d|d<|ddusJd|ddus$Jd||d<dS) Ndockerimage ray_containercontainer_name head_imagezCannot override head_image worker_imagezCannot override worker_image)r8)r: docker_image docker_configr"r"r#override_docker_images  rEc Cstdtdtjddd}|jdd}|d}t|d }t|d }|jd d}|d}t}tj |d }t |d  } | |Wdn1sQwYt dddd|d|d|g dS)zRGet Azure service principal credentials from AWS Secrets Manager and authenticate.&======================================z5Getting Azure credentials from AWS Secrets Manager...secretsmanager us-west-2 region_namez#azure-service-principal-oss-release)SecretId SecretString client_id tenant_idz#azure-service-principal-certificatezazure_cert.pemwNazloginz--service-principalz --usernamez --certificatez--tenant)r,boto3clientget_secret_valuejsonloadstempfilemkdtempr1pathjoinopenwrite subprocess check_call) secrets_clientsecret_responsesecretrMrN cert_responsecerttmp_dir cert_pathfr"r"r#azure_authenticates<  rgcCs|tdtdtjddd}d}d}tjd|}tjtj|s/ttj|| |||t |d d S) z=Download the ssh key from the S3 bucket to the local machine.rFzDownloading ssh key...s3rHrIzaws-cluster-launcher-testzray-autoscaler_59_us-west-2.pem~/.ssh/N) r,rRrSr1rY expanduserexistsdirnamemakedirs download_filechmod) s3_client bucket_namekey_namelocal_key_pathr"r"r#download_ssh_key_awssrucCstdtdt}d}d}||}||}tjd|}tjtj |s6t tj || |t |ddS)zGDownload the ssh key from the google cloud bucket to the local machine.rFzDownloading ssh key from GCP...z*gcp-cluster-launcher-release-test-ssh-keyszAray-autoscaler_gcp_us-west1_anyscale-bridge-cd812d38_ubuntu_0.pemrirjN) r,rClient get_bucketget_blobr1rYrkrlrmrndownload_to_filenamerp)rSrrrsbucketblobrtr"r"r#download_ssh_key_gcps   r|c Cstdtdd}d}t|D]i}z#tj}|ddtjdddd t|gd d |d t |WdStj yy}z7td |d d|dt|j dt td|j dtd|j d|}WYd}~qd}~ww|)z Clean up the cluster using the given cluster configuration file. Args: cluster_config: The path of the cluster configuration file. rFzCleaning up cluster...Nr PYTHONPATHr)down-v-yTcheckcapture_outputenvzray down fails[r /z]: utf-8stdout: stderr: )r,ranger1environcopypopr]runrcleanup_security_groupsCalledProcessErroroutputdecode traceback print_excstdoutstderr)r:r last_error num_triesirer"r"r#cleanup_clusters4    rc Csd}|dkrYz |j|dWdStjjyR}z5|jdddkr:d|}td|d |t||d 7}ntd |WYd}~dSWYd}~nd}~ww|dksdSdS) Nr )GroupIdErrorCodeDependencyViolationzWaiting zGs for the instance to be terminated before deleting the security group r zError deleting security group: )delete_security_groupbotocore exceptions ClientErrorresponser,timesleep) ec2_clientidretryr sleep_timer"r"r#cleanup_security_groups&    rc Cstdtdtjd}tjd}tj|r&tj|r&tddStdtj|}tj|s=tj|dd ztjd d d d dd|ddddg dddtdWdStj y}z&tdtd|j dtd|j dt dWYd}~dSd}~ww)zX Ensure that the SSH keys for Azure tests exist, and create them if they don't. rFz Ensuring Azure SSH keys exist...z#~/.ssh/ray-autoscaler-tests-ssh-keyz'~/.ssh/ray-autoscaler-tests-ssh-key.pubzAzure SSH keys already exist.Nz.Azure SSH keys not found. Creating new keys...T)exist_okz ssh-keygenz-trsaz-b4096z-fz-Nrz-Czray-autoscaler-azure)rrz$Successfully created Azure SSH keys.zError creating SSH keys:rrrr )r,r1rYrkrlrmrnr]rrrrrr-r.)private_key_pathpublic_key_pathssh_dirrr"r"r#ensure_ssh_keys_azure/sH     rc Cs|did}|dkrdSz*tjddd}|jdtgdd |d gdgd }|d D] }t||d q-WdStyS}ztd|WYd}~dSd}~ww)Nproviderr awsec2rHrIztag-key)NameValuesztag:ray-cluster-name cluster_name)FiltersSecurityGroupsrz#Error cleaning up security groups: )r8rRrSdescribe_security_groupsrr Exceptionr,)config provider_typerrsecurity_grouprr"r"r#r[s* rr c Cs|did}|dkrttdtdgd}|r#|d|t|tj}|dd z t j |d d |d Wn*t j yj}zt|j td |j d td|jd |d }~wwtdtdd} d} | |krz"dddt|dd|dg}|r|dt j |d |dd } Wn"t j y| d7} td| d|dtdYnw| |ks{| stdtd|dt||tdtd tdtdtd!t||tdtd"d S)#ak Run the necessary Ray commands to start a cluster, verify Ray is running, and clean up the cluster. Args: cluster_config: The path of the cluster configuration file. retries: The number of retries for the verification step. no_config_cache: Whether to pass the --no-config-cache flag to the ray CLI commands. rr azurerFzStarting new cluster...)r)uprrrr}NTrrrrzVerifying Ray is running...Frr)execrz2python -c 'import ray; ray.init("localhost:6379");z assert len(ray.nodes()) >= ')rrr z#Verification failed. Retry attempt z of z...<z!Error: Verification failed after z0 attempts. Cleaning up cluster before exiting...zExiting script.zRay verification successful.z'Finished executing script successfully.)r8rgr,r9rr1rrrr]rrrrrrrrrr-r.) r:rrrrrcmdrrsuccesscountr"r"r#run_ray_commandstsv           r__main__z"Using cluster configuration file: z4Number of retries for 'verify ray is running' step: zUsing --no-config-cache flag: z6Number of expected nodes for 'verify ray is running': r-rFzOverriding ray wheel...: zOverriding docker image...: zUsing docker image: rr Fcache_stopped_nodesrravailability_zonegcp)gcloudr get-valueaccountT)rrrzActive account email:)rrrprojectzInjecting GCP project 'z$' into cluster configuration file... project_idvspherezVSPHERE provider detected.z-Provider type not recognized. Exiting script.z.yaml)suffix)r )A__doc__rrUr1r'r]r-rWrrpathlibrrRryaml google.cloudrr)"ray.autoscaler._private.aws.configrr$r/r5r<rErgrur|rrrrr__name__rrrrrrr, safe_load read_textr:rrdocker_override_imager8rrrPIPErrstrip account_emailrr.NamedTemporaryFiletempr\dumpencodeflushnamer"r"r"r#s    7 --, T             $