o پiW@sdZddlZddlZddlZddlZddlZddlZddlZddl Z ddl Z ddl Z ddl m Z mZmZmZmZmZmZmZee dddddlZddlZddlZddlmZddlmZdd lmZdd lm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0m1Z1m2Z2dd l3m4Z4dd l5m6Z6dd l7m8Z8ddl9m:Z:ddl;mZ>m?Z?ddl@mAZAmBZBddlCmDZDmEZEmFZFmGZGmHZHmIZImJZJmKZKmLZLmMZMmNZNmOZOddlPmQZQddlRmSZSeTeUZVeWeXeIZY d*de?de>dee:dee:e8ffddZZGdddeZ[de?fddZ\ded edeefd!d"Z]de?de>d#efd$d%Z^ d*de?d&ed#ed'edee>dee:e8eee>ff d(d)Z_dS)+zy The entry point of inference server. (SRT = SGLang Runtime) This file implements python APIs for the inference engine. N) AsyncIteratorCallableDictIteratorListOptionalTupleUnion_register_atexitcOsdSN)argskwargsr r Q/home/ubuntu/.local/lib/python3.10/site-packages/sglang/srt/entrypoints/engine.py!sr) EngineBase)$run_data_parallel_controller_process)run_detokenizer_process)CloseSessionReqInput!DestroyWeightsUpdateGroupReqInputEmbeddingReqInputGenerateReqInputGetWeightsByNameReqInputInitWeightsUpdateGroupReqInput"LoadLoRAAdapterFromTensorsReqInputLoadLoRAAdapterReqInputMultimodalDataInputFormatOpenSessionReqInputReleaseMemoryOccupationReqInputResumeMemoryOccupationReqInput RpcReqInput RpcReqOutputUnloadLoRAAdapterReqInputUpdateWeightFromDiskReqInput$UpdateWeightsFromDistributedReqInputUpdateWeightsFromIPCReqInputUpdateWeightsFromTensorReqInput)MultiTokenizerRouter)run_scheduler_process)TemplateManager)TokenizerManager)?parse_remote_instance_transfer_engine_info_from_scheduler_infos)PortArgs ServerArgs)process_tracing_inittrace_set_thread_info) MultiprocessingSerializerassert_pkg_versionconfigure_loggerget_bool_env_varget_zmq_socketis_cudakill_process_tree launch_dummy_health_check_servermaybe_reindex_device_id numa_utilsset_prometheus_multiproc_dir set_ulimit)TorchMemorySaverAdapter) __version__ server_args port_argsTokenizerManagerClassreturncCs8|pt}|||}t}|j||j|j|jd||fS)N)tokenizer_manager model_path chat_templatecompletion_template)r*r)initialize_templatesrCrDrE)r>r?r@rBtemplate_managerr r rinit_tokenizer_manageras rHc4@seZdZUdZeZeed<eeZ e ed<ee Z e ed<ee Ze ed<ddZ                   dd eeeeefd eeeeefd eeeeeeefd eedeedeedeeeeefdeeeeefdeeeeefdeeeeeeefdeeeedeeeeefdedededeeeeefdeeeeefdeeeeefdeedeedeeeeefdeed eed!eeeeff0d"d#Z                   dd eeeeefd eeeeefd eeeeeeefd eedeedeedeeeeefdeeeeefdeeeeefdeeeeeeefdeeeedeeeeefdedededeeeeefdeeeeefdeeeeefdeedeedeeeeefdeed eed!eeeeff0d$d%Z       dd eeeeeeeeefd eedeedeed&eedeeeeeeefdeedeeeeefd!efd'd(Z       dd eeeeeeeeefd eedeedeed&eedeeeeeeefdeedeeeeefd!efd)d*Zd eeeed!efd+d,Zd-d.Zd/d0Z d1d2Z!d3d4Z" dd5ed6eed!efd7d8Z#d6ed!dfd9d:Z$d;d<Z%d=d>Z&d?d@Z'dAdBZ(dCdDZ)dEdFZ* GddHedIedJedKedLedMef dNdOZ+dLefdPdQZ, R S ddTe-edUe-edVe-e-edLedWedXeef dYdZZ.  Sdd[ee/ee0j1fdXeedWefd\d]Z2 dd^edXeefd_d`Z3 SddaeeefdWefdbdcZ4ddeedfefdgdhZ5diedjee/ee0j1fdkefdldmZ6ddiedednefdodpZ7diefdqdrZ8ddseeefdtduZ9ddseeefdvdwZ:dxdyZ; dzefd{d|Z    ddeeeeefdeeeeeeeefdeeededed!eee?f ddZ@    ddeeeeefdeeeeeeeefdeeededed!eee?f ddZAdS)Enginea The entry point to the inference engine. - The engine consists of three components: 1. TokenizerManager: Tokenizes the requests and sends them to the scheduler. 2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager. 3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager. Note: 1. The HTTP server, Engine, and TokenizerManager all run in the main process. 2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library. server_args_classinit_tokenizer_manager_funcrun_scheduler_process_funcrun_detokenizer_process_funcc KsDd|vr |d}nd|vrd|d<|jdi|}||_td|t|jt||j|j |j d\}}}}||_ ||_ |d|_ ||_t||_td}|jjdkrct|tj|jjd|_nd |_|jrt|jd d }|jd kryd }n|jdkrd}t|zt|_Wd St yt!|_t"|jYd Sw)z The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`. Please refer to `ServerArgs` for the documentation. r> log_levelerror server_args=)r>rKrLrMrTNsglang TokenizerprefillzPrefill TokenizerdecodezDecode Tokenizerr )#rJr>loggerinfoatexitregistershutdown_launch_subprocessesrKrLrMrBrGscheduler_infor?r+$remote_instance_transfer_engine_infozmqContext node_rankr4DEALER rpc_ipc_name send_to_rpc enable_tracer.otlp_traces_endpointdisaggregation_moder/asyncioget_running_looploop RuntimeErrornew_event_loopset_event_loop) selfrr>rBrGscheduler_infosr?context thread_labelr r r__init__sV           zEngine.__init__NFpromptsampling_params input_ids image_data audio_data video_datareturn_logproblogprob_start_lentop_logprobs_numtoken_ids_logprob lora_pathcustom_logit_processorreturn_hidden_statesreturn_routed_expertsstreambootstrap_hostbootstrap_portbootstrap_roomdata_parallel_rankexternal_trace_headerridsession_paramspriorityrAcsjjr%|durtdn|dkrtd|jjkr%tdjjtdid|d|d|d |d |d |d |d |d| d| d| d| d| d|d|d|d|d|d|d|d|d|d|}j|d|rfdd}|Sj }|S) The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. Please refer to `GenerateReqInput` for the documentation. N7data_parallel_rank not provided, using default dispatchr'data_parallel_rank must be non-negativez.data_parallel_rank must be less than dp_size: textrtrsrurvrwrxryrzr{r|r}r~rrrrrrrrrrc3s6 z j}|VWn tyYdSwqr )rirun_until_complete __anext__StopAsyncIteration)chunk generatorrmr rgenerator_wrappers  z*Engine.generate..generator_wrapperr ) r>enable_dp_attentionrVdebug ValueErrordp_sizerrBgenerate_requestrirr)rmrrrsrtrurvrwrxryrzr{r|r}r~rrrrrrrrrrobjrretr rrgenerates|(        zEngine.generatecs|jjr)|durtdn|dkrtd||jjkr)td|jjddtd|td!id |d |d |d |d |d|d|d|d| d| d| d| d|d|d| d|d|d|d|d|d|d|d|}|j|d}|d ur|S| IdHS)"rNrrrz(data_parallel_rank must be in range [0, ]zdata_parallel_rank: rrtrsrurvrwrxryrzr{r|r~rrr}rrrrrrrrTr ) r>rrVrrrrrBrr)rmrrrsrtrurvrwrxryrzr{r|r}r~rrrrrrrrrrrrr r rasync_generate)s|)       zEngine.async_generate dimensionsc Cs:t||||||||d} |j| d} |j| } | S) The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. Please refer to `EmbeddingReqInput` for the documentation. rrurvrwrr|rrNrrBrrirr) rmrrrurvrwrr|rrrrrr r rencode}s z Engine.encodec s6t||||||||d} |j| d} | IdHS)z Asynchronous version of encode method. The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. Please refer to `EmbeddingReqInput` for the documentation. rN)rrBrr) rmrrrurvrwrr|rrrrr r r async_encodes zEngine.async_encodecCs.t|dd}|j|d}|j|}|S)rT)ris_cross_encoder_requestNr)rmrrrrrr r rreranks z Engine.rerankcCsttdddS)zShutdown the engineF)include_parentN)r6osgetpidrmr r rrZszEngine.shutdowncCs|Sr r rr r r __enter__szEngine.__enter__cCs |dS)NF)rZ)rmexc_type exc_value tracebackr r r__exit__szEngine.__exit__cCs|j|jSr )rirrB flush_cacherr r rrszEngine.flush_cachecapacity_of_str_len session_idcC"t||d}|j|j|dS)abOpen a session for multi-turn conversation with shared context. Args: capacity_of_str_len: Maximum string length capacity for the session. session_id: Optional session ID. If not provided, a UUID will be generated. Returns: The session ID (either the provided one or a newly generated UUID). )rrN)rrirrB open_session)rmrrrr r rrs zEngine.open_sessioncCs$t|d}|j|j|ddS)zsClose a session and release its resources. Args: session_id: The session ID to close. )rN)rrirrB close_session)rmrrr r rrs zEngine.close_sessioncKs|j|jjdi|dS)Nr )rirrB start_profilermrr r rrszEngine.start_profilecC|j|jdSr )rirrB stop_profilerr r rrszEngine.stop_profilecCrr )rirrB start_expert_distribution_recordrr r rrz'Engine.start_expert_distribution_recordcCrr )rirrBstop_expert_distribution_recordrr r rrrz&Engine.stop_expert_distribution_recordcCrr )rirrBdump_expert_distribution_recordrr r rrrz&Engine.dump_expert_distribution_recordcCs4|j|j}it|jj|j|tdS)N)internal_statesversion) rirrBget_internal_state dataclassesasdictr>r\r=)rmrr r rget_server_infos zEngine.get_server_infoncclmaster_address master_port rank_offset world_size group_namebackendcC*t||||||d}|j|j|dS)z"Initialize parameter update group.)rrrrrrN)rrirrBinit_weights_update_group)rmrrrrrrrr r rr  z Engine.init_weights_update_groupcC t|d}|j|j|dS)zDestroy parameter update group.)rN)rrirrBdestroy_weights_update_group)rmrrr r rr(s  z#Engine.destroy_weights_update_groupweight_update_groupTnamesdtypesshapesr load_formatcCr)z'Update weights from distributed source.)rrrrrrN)r$rirrBupdate_weights_from_distributed)rmrrrrrrrr r rr4rz&Engine.update_weights_from_distributed named_tensorscsL|dkr}n fddt|jjD}t|||d}|j|j|dS)zUpdate weights from distributed source. If there are going to be more updates, set `flush_cache` to be false to avoid duplicated cache cleaning operation.flattened_bucketcsg|]}tqSr )r0 serialize).0_rr r Usz5Engine.update_weights_from_tensor..)serialized_named_tensorsrrN)ranger>tp_sizer&rirrBupdate_weights_from_tensor)rmrrrrrr rrrJs   z!Engine.update_weights_from_tensorrCcCr)a Update the weights from disk inplace without re-launching the engine. This method allows updating the model weights from disk without restarting the engine. It can be used to load a different model or update weights with new training. )rCrN)r#rirrBupdate_weights_from_disk)rmrCrrr r rrbs  zEngine.update_weights_from_disk zmq_handlescCr)z:Update weights from IPC for checkpoint-engine integration.)rrN)r%rirrBupdate_weights_from_ipc)rmrrrr r rrvs zEngine.update_weights_from_ipcdname truncate_sizecCr)zGet weights by parameter name.)rrN)rrirrBget_weights_by_name)rmrrrr r rrs  zEngine.get_weights_by_name lora_nametensors config_dictcCs2tj|dd}t|||d}|j|j|dS)NT) output_str)rrserialized_tensors)r0rrrirrBload_lora_adapter_from_tensors)rmrrrrlora_reqr r rrs z%Engine.load_lora_adapter_from_tensorspinnedcCs$t|||d}|j|j|dS)z8Load a new LoRA adapter without re-launching the engine.)rr|rN)rrirrBload_lora_adapter)rmrr|rrr r rrs zEngine.load_lora_adaptercCr)z6Unload a LoRA adapter without re-launching the engine.)rN)r"rirrBunload_lora_adapter)rmrrr r rrs  zEngine.unload_lora_adaptertagscC t|d}|j|j|dSN)r)rrirrBrelease_memory_occupationrmrrr r rr  z Engine.release_memory_occupationcCrr)rrirrBresume_memory_occupationrr r rrrzEngine.resume_memory_occupationcCs|j|jdS)a< To maintain a high performance server with low latency, we want to reduce the stalls caused by the garbage collector scanning through a large number of objects. It is usually helpful to start the server and warm it up with real requests to initialize many of the long-lived objects that do not need to be garbage collected. After sufficient warmup, we can call this function to freeze the garbage collector so that all objects created before this point are considered out of scope for garbage collection. N)rirrB freeze_gcrr r rrs zEngine.freeze_gcmethodcKsHt||d}|j||jtj}t|tsJ|js"J|j dS)N)r parameters) r rc send_pyobj recv_pyobjr^BLOCKY isinstancer!successmessage)rmrrrrecv_reqr r rcollective_rpcs  zEngine.collective_rpccK|jdi|dS)Nsave_remote_model)rrrr r rrzEngine.save_remote_modelcKr)Nsave_sharded_model)rrrr r rrrzEngine.save_sharded_modelqueryitemslabel_token_ids apply_softmax item_firstc Cs |j|jj|||||ddS)a Score the probability of specified token IDs appearing after the given (query + item) pair. For example: query = "<|user|>Is the following city the capital of France? " items = ["Paris <|assistant|>", "London <|assistant|>", "Berlin <|assistant|>"] label_token_ids = [2332, 1223] # Token IDs for "Yes" and "No" item_first = False This would pass the following prompts to the model: "<|user|>Is the following city the capital of France? Paris <|assistant|>" "<|user|>Is the following city the capital of France? London <|assistant|>" "<|user|>Is the following city the capital of France? Berlin <|assistant|>" The api would then return the probabilities of the model producing "Yes" and "No" as the next token. The output would look like: [[0.9, 0.1], [0.2, 0.8], [0.1, 0.9]] Args: query: The query text or pre-tokenized query token IDs. Must be provided. items: The item text(s) or pre-tokenized item token IDs. Must be provided. label_token_ids: List of token IDs to compute probabilities for. If None, no token probabilities will be computed. apply_softmax: Whether to normalize probabilities using softmax. item_first: If True, prepend items to query. Otherwise append items to query. Returns: List of dictionaries mapping token IDs to their probabilities for each item. Each dictionary in the list corresponds to one item input. Raises: ValueError: If query is not provided, or if items is not provided, or if token IDs are out of vocabulary, or if logprobs are not available for the specified tokens. Nr r r r r request)rirrB score_requestrmr r r r r r r rscores'z Engine.scorecs |jj|||||ddIdHS)zh Asynchronous version of score method. See score() for detailed documentation. Nr)rBrrr r r async_scores  zEngine.async_score)NNNNNNFNNNNNFFFNNNNNNNN)NNNNNNNr )r)rTN)NT)T)r)F)NNNFF)B__name__ __module__ __qualname____doc__r-rJ__annotations__ staticmethodrHrKrr(rLrrMrqrr rstrrintrboolrrrrrrrrZrrrrrrrrrrrrrlistrrtorchTensorrrrrrrrrrrrrrfloatrrr r r rrIvs  E  !"# _  !"# W                       4  rIcCspdtjvs|jrtt|jtjd<dtjvs|js|jr*tt|jp$|jtjd<dtjd<dtjd<tjddd krBdtjd<tjd durOd tjd <tjd dur\dtjd <d tdt ddtjd<|j rst t t ds|jdkrtdddtrtddd|jdurdd}ttj|ntd|jttj|jtjddddS) NNCCL_CUMEM_ENABLENCCL_NVLS_ENABLE8CUDA_DEVICE_MAX_CONNECTIONSAUTOCUDA_MODULE_LOADINGTRTLLM_ENABLE_PDL10CUTE_DSL_LOG_LEVEL30CUTE_DSL_LOG_TO_CONSOLEz sglang-run--ri SGLANG_RUN_ID$SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK flashinferflashinfer_pythonz0.6.3zPlease uninstall the old version and reinstall the latest version by following the instructions at https://docs.flashinfer.ai/installation.html.z sgl-kernelz0.3.21zSPlease reinstall the latest version with `pip install sgl-kernel --force-reinstall`cSstdttdS)NzIReceived sigquit from a child process. It usually means the child failed.)rVrOr6rr)signumframer r rlaunch_phase_sigquit_handlerdsz:_set_envs_and_config..launch_phase_sigquit_handlerzUsing custom SIGQUIT handler: spawnT)force)renvironenable_symm_memrrenable_nccl_nvlsgettimerandomrandintenable_metricsr:r;r3attention_backendr1_is_cudacustom_sigquit_handlersignalSIGQUITrVrOmpset_start_method)r>r4r r r_set_envs_and_config&sX         rFscheduler_pipe_readersscheduler_procsc Csg}tt|D]<}z||}Wn"ty4td|d||td||jw|ddkr?td| |q|S)z@Wait for the model to finish loading and return scheduler infos.zRank z< scheduler is dead. Please check if there are relevant logs.z Exit code: statusreadyz;Initialization failed. Please see the error messages above.) rlenrecvEOFErrorrVrOjoinexitcoderjappend)rGrHrnidatar r r_wait_for_scheduler_readyvs$     rSrLcCsFg}|jdkrtj|jd}g}t|j|jd}t|j|jd}t||j|||j|d}|} |j | } t| |j| | |j| d} |D]} | D]} t j dd\}}|j | || | | |j }|jrq|jnd}|j ||j}| ||j}| |j |j}| |j |j|j |j|j}t|G}t j||||| |||| d|f d}|"t|| |Wdn1swYWdn1swYWdn1swY||||qQqM||fSt j dd\}}|g}t jtt||||dd}|||||fS)Nr)enableF)duplextargetr )r>r? pipe_writerrL)rWr)rr<createenable_memory_savermaxpp_sizennodesrr`rrDPipe base_gpu_id gpu_id_stepr attn_cp_size moe_dp_sizeep_sizer8Processconfigure_subprocessr9startrPrdict)r>r?rLrHmemory_saver_adapterrGpp_size_per_nodennodes_per_pp_rank pp_rank_rangennodes_per_tp_grouptp_size_per_node tp_rank_rangepp_ranktp_rankreaderwritergpu_id attn_dp_size attn_tp_size attn_cp_rank moe_dp_rank moe_ep_rankprocr r r_launch_scheduler_processess            F rzrKrMc Cs*t|t|||durt|}td|t|||d\}}|jdkrbt ||}t ddkr=dd||fSt |j |j|j|D]}|td|jd|jqHdd||fStj|||fd } | |jdkr||||\} } nt||} d} t ||}|d d | _| | ||fS) z Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess. NrP)r>r?rLr"SGLANG_BLOCK_NONZERO_RANK_CHILDRENr)z$Scheduler or DataParallelController z terminated with rVrmax_req_input_len)r2rFcheck_server_argsr,init_newrVrWrzr`rSrgetenvr7hostportr>rNrOpidrOrDrdrftokenizer_worker_numr'r|) r>rKrLrMr?rHrGrnry detoken_procrBrGr r rr[sT          r[r )`rrgrXrloggingmultiprocessingrDrr<rB threadingr;typingrrrrrrrr setattrruvloopr^!sglang.srt.entrypoints.EngineBaser,sglang.srt.managers.data_parallel_controllerr'sglang.srt.managers.detokenizer_managerrsglang.srt.managers.io_structrrrrrrrrrrrrr r!r"r#r$r%r&)sglang.srt.managers.multi_tokenizer_mixinr'sglang.srt.managers.schedulerr($sglang.srt.managers.template_managerr)%sglang.srt.managers.tokenizer_managerr*;sglang.srt.model_loader.remote_instance_weight_loader_utilsr+sglang.srt.server_argsr,r-sglang.srt.tracing.tracer.r/sglang.srt.utilsr0r1r2r3r4r5r6r7r8r9r:r;+sglang.srt.utils.torch_memory_saver_adapterr<sglang.versionr= getLoggerrrVset_event_loop_policyEventLoopPolicyr@rHrIrFrSrzr[r r r rs (   T     8     5P  k