Fix deadlock for in_place/retract in pause_generation

[maocheng23]

Summary

• Fix writer lock deadlock in pause_generation mode during IPC weight updates and post-processing hooks
• When the scheduler is paused, existing readers are blocked on is_pause_cond, so acquiring the writer lock would deadlock. This change skips the writer lock when paused (using nullcontext()) since no concurrent inference can race in that state.
• Applies to both update_weights_from_ipc and post_process_weights paths

Test plan

• Verify IPC weight update completes without deadlock when scheduler is paused
• Verify normal (non-paused) weight updates still acquire the writer lock correctly
• Run existing e2e tests to confirm no regression

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request modifies update_weights_from_ipc and post_process_weights to bypass the model update writer lock when the engine is paused, preventing potential deadlocks. However, the review identifies a critical race condition in both functions: by releasing the is_pause_cond lock before performing the weight updates, the engine could be unpaused by a concurrent call, leading to unsafe concurrent weight updates and inference. It is recommended to hold the is_pause_cond lock during the update when the engine is paused to ensure thread safety.

[gemini-code-assist]

The current implementation introduces a race condition. By releasing is_pause_cond before performing the weight update with nullcontext(), it's possible for continue_generation to be called in the interim. This would set is_pause = False, allowing new inference requests to proceed and the scheduler to unpause. This leads to concurrent weight updates and inference, which can cause data corruption or crashes.

To fix this, the weight update should be performed while holding the is_pause_cond lock if the engine is paused, ensuring that continue_generation cannot unpause the scheduler until the update is complete.

Suggested change

	# Skip the writer lock when paused: readers are blocked on
	# is_pause_cond so no concurrent inference can race, and
	# waiting for the writer lock would deadlock because existing
	# readers are stuck waiting on the paused scheduler.
	async with self.is_pause_cond:
	is_paused = self.is_pause
	lock_context = (
	self.model_update_lock.writer_lock
	if not is_paused
	else nullcontext()
	)
	async with lock_context:
	result = (await self.update_weights_from_ipc_communicator(obj))[0]
	success, message = result.success, result.message
	async with self.is_pause_cond:
	is_paused = self.is_pause
	if is_paused:
	# Skip the writer lock when paused to avoid deadlock with existing readers.
	# We hold is_pause_cond during the update to prevent continue_generation
	# from unpausing the scheduler and allowing new readers to race.
	result = (await self.update_weights_from_ipc_communicator(obj))[0]
	success, message = result.success, result.message
	if not is_paused:
	# This means that weight sync cannot run while requests are in progress.
	async with self.model_update_lock.writer_lock:
	result = (await self.update_weights_from_ipc_communicator(obj))[0]
	success, message = result.success, result.message

[gemini-code-assist]

Similar to the change in update_weights_from_ipc, this implementation is susceptible to a race condition where continue_generation could unpause the engine after is_paused is checked but before the update completes. Holding the is_pause_cond lock during the update when paused prevents this race.

Suggested change

	async with self.is_pause_cond:
	is_paused = self.is_pause
	lock_context = (
	self.model_update_lock.writer_lock if not is_paused else nullcontext()
	)
	async with lock_context:
	results = await self.post_process_weights_communicator(obj)
	return _Communicator.merge_results(results)
	async with self.is_pause_cond:
	is_paused = self.is_pause
	if is_paused:
	results = await self.post_process_weights_communicator(obj)
	if not is_paused:
	async with self.model_update_lock.writer_lock:
	results = await self.post_process_weights_communicator(obj)
	return _Communicator.merge_results(results)