Link Management

Function Description

The HCCL collective communication APIs used during NN model execution employ a bilateral communication mode, requiring both sides to initiate link establishment simultaneously. In the Prefill-Decode disaggregation solution, this process is simplified to unilateral link establishment initiated by the client. Either the Prefill or Decode side can be defined as the client or server. In the following example, the Decode side is defined as the client, and link establishment follows a procedure where the Decode side initiates link establishment requests to the Prefill side.

Two core APIs are provided: link_clusters and unlink_clusters, both called by the client. Link establishment operations are performed on a point-to-point basis.

link_clusters: Establishes links between nodes.
unlink_clusters: Disconnects links between nodes.

Scenario

Link establishment is a prerequisite for KV cache transmission between Prefill and Decode nodes; therefore, links must be established before KV cache transmission is enabled.
In cluster reliability scenarios, if a node on the Prefill or Decode side experiences an anomaly, the faulty node can be taken offline by disconnecting its links, provided that overall cluster availability remains unaffected.
Dynamic adjustment of the PD cluster ratio is achieved through link establishment and disconnection: Nodes are added or removed dynamically based on workload, where node addition requires link establishment and node removal requires link disconnection.

Example

The following pseudocode demonstrates the link establishment process between a Prefill node and a Decode node:

Launch the Prefill- and Decode-side scripts, where the LLM-DataDist initialization API is called. The listening host IP address and port number must be set on the Prefill side.

# Script on the Prefill side
from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo

 # llm datadist initialization
 llm_datadist = LLMDataDist(LLMRole.Prompt, cluster_id=0)
 llm_config = LLMConfig()
 llm_config.listen_ip_info = "10.10.1.1:26000" # local_host_ip + port
 llm_config.device_id = 0
 llm_options = llm_config.generate_options()
 llm_datadist.init(llm_options)

# Script on the Decode side
from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo

 # LLM-DataDist initialization
 llm_datadist = LLMDataDist(LLMRole.DECODER, cluster_id=0)
 llm_config = LLMConfig()
 llm_config.device_id =0 
 llm_options = llm_config.generate_options()
 llm_datadist.init(llm_options)

In the Decode-side script, call link_clusters to initiate link establishment. When the service exits, call unlink_clusters on the Decode side to disconnect the link.

 # Generate cluster information for link establishment.
 cluster = LLMClusterInfo()
 cluster.remote_cluster_id = 1  # remote_cluster_id must match the LLM-DataDist instance created on the Prefill side.
 cluster.append_local_ip_info("192.168.2.1", 26000) # local_ip_info is the IP address of the host with which a link needs to be established locally.
 cluster.append_remote_ip_info("192.168.1.1", 26000) # remote_ip_info is the IP address of the host that wants to establish a link with the peer end.

 # Call link_clusters to establish a link.
 # ret is the returned value of the API, and rets indicates the link establishment result of each cluster.
 ret, rets = llm_datadist.link_clusters([cluster], timeout=5000)
 # Check the link establishment result.
 if ret != LLMStatusCode.LLM_SUCCESS:
     raise Exception("link failed.")
 for cluster_i in range(len(rets)):
     link_ret = rets[cluster_i]
     if link_ret != LLMStatusCode.LLM_SUCCESS:
         print(f"{cluster_i} link failed.")

When the service ends, the link is disconnected on the Decoder side. Both Prefill and Decoder call the finalize method of the llm_datadist API to release resources.

 # Script on the Prefill side
 # Call llm_datadist to allocate the KV cache.
 # Perform service inference.
 # ...

 # Exit the service.
 llm_datadist.finalize()

 # Script on the Decode side
 # pull_cache and model inference
 # ...

 # Call unlink_clusters to disconnect the link upon service exit.
 ret, rets = llm_datadist.unlink_clusters([cluster], timeout=5000)
 if ret != LLMStatusCode.LLM_SUCCESS:
     raiseRuntimeError(f'[unlink_cluster] failed, ret={ret}')
 llm_datadist.finalize()

To add a node or bring an offline node back online, the same procedure must be followed. To remove a node, under normal circumstances, the Decode side should actively call unlink_clusters. If the Decode side cannot call unlink_clusters, the Prefill side must call unlink_clusters instead.

The link_clusters API is called when nodes are brought online, and the unlink_clusters API is called when nodes are brought offline to flexibly scale in or out distributed clusters.

Exception Handling

If an exception occurs on the Decode side and prevents it from calling unlink_clusters, the Prefill side must call unlink_clusters to clean up resources; otherwise, link establishment cannot be performed again.

The unlink_clusters API provides a forcible disconnection capability, which is useful in case of link faults where normal disconnection may take a long time. When the forcible disconnection option is enabled (by setting force=True), the call must be initiated on both sides, and only the local link is cleaned up.

# Forcibly disconnect the link.
 ret, rets = llm_datadist.unlink_clusters([cluster], timeout=5000, force=True)
 if ret != LLMStatusCode.LLM_SUCCESS:
     raise RuntimeError(f'[unlink_clusters] failed, ret={ret}')

Parent topic: Function Description