support tp>1

[fangyuchu]

Purpose

Test Plan

test with dp2 tp 2, and inject fault on DP rank 1 (on both tp ranks)

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request updates the fault tolerance engine core and GPU worker sentinel to support multiple data parallel (DP) ports instead of a single port, using JSON serialization to store and retrieve the list of ports. However, the reviewer identified critical issues with this implementation: generating a port for every rank in the world size and indexing by the worker's rank is incorrect for multi-TP/PP configurations. Because ranks within the same DP group must share the same port to successfully initialize, the port allocation and indexing should instead be based on the number of independent DP groups (calculated via world size divided by data parallel size).

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Generating world_size ports is redundant and incorrect. In a multi-TP/PP setup, the data parallel group is formed by ranks that share the same TP/PP/PCP indices. Therefore, there are world_size // data_parallel_size independent DP groups, and we only need one port per DP group. Generating world_size ports and indexing them by self.worker.rank will cause different ranks in the same DP group to use different ports, preventing them from connecting and causing initialization to hang.

Suggested change

	if parallel_config.data_parallel_rank == 0:
	worker_port = get_open_port()
	worker_ports = [get_open_port() for _ in range(parallel_config.world_size)]
	engine_port = get_open_port()
	engine.dp_store.set(worker_key, str(worker_port).encode())
	engine.dp_store.set(worker_key, json.dumps(worker_ports).encode())
	engine.dp_store.set(engine_key, str(engine_port).encode())
	else:
	worker_port = int(engine.dp_store.get(worker_key).decode())
	worker_ports = json.loads(engine.dp_store.get(worker_key).decode())
	engine_port = int(engine.dp_store.get(engine_key).decode())
	if parallel_config.data_parallel_rank == 0:
	num_workers_per_dp_rank = parallel_config.world_size // parallel_config.data_parallel_size
	worker_ports = [get_open_port() for _ in range(num_workers_per_dp_rank)]
	engine_port = get_open_port()
	engine.dp_store.set(worker_key, json.dumps(worker_ports).encode())
	engine.dp_store.set(engine_key, str(engine_port).encode())
	else:
	worker_ports = json.loads(engine.dp_store.get(worker_key).decode())
	engine_port = int(engine.dp_store.get(engine_key).decode())

[gemini-code-assist]

Using self.worker.rank % world_size resolves to self.worker.rank because the rank is always less than the world size. This assigns a unique port to every single worker. However, workers belonging to the same DP group must share the same port to initialize their process group. Indexing by self.worker.rank % num_workers_per_dp_rank ensures that all ranks in the same DP group get the same port.

            world_size = self.worker.parallel_config.world_size
            num_workers_per_dp_rank = world_size // self.dp_size
            port = params["new_stateless_dp_group_ports"][self.worker.rank % num_workers_per_dp_rank]