Add health check, make async Engine more robust

[Yard1]

For production usecases, we want to be able to detect Engine failures, especially ones that can happen silently (eg. due to NCCL timeouts). This PR adds a health check method (currently only checking the health of Ray workers) and makes the Async engine more robust by adding a timeout for each iteration as well as better error reporting.

[zhuohan123]

Thanks for the contribution! In general LGTM. Left some small questions.

[zhuohan123]

We raise errors in both try branch and except branch. Then what does the finally here do?

[Yard1]

we want to run the error callback even after we re-raise an exception in except

[njhill]

I think you could just do this in the except block though before re-raising (things will still run in the same order)

[zhuohan123]

Why don't we always clear this flag?

Suggested change

	async def wait_for_new_requests(self, clear: bool):
	if not self.has_new_requests():
	await self.new_requests_event.wait()
	if clear:
	self.new_requests_event.clear()
	async def wait_for_new_requests(self):
	if not self.has_new_requests():
	await self.new_requests_event.wait()
	self.new_requests_event.clear()

[zhuohan123]

Also, what's the reason behind this change? Why do we need to move the clear call from get_new_and_finished_requests to here?

[Yard1]

Yes, we can always clear it.

The reason is to ensure the event is cleared as soon as we have new requests

[njhill]

Thanks @Yard1 this looks great

[njhill]

I think you could just do this in the except block though before re-raising (things will still run in the same order)

[njhill]

Suggestion to only clear before waiting

Suggested change

	if not self.has_new_requests():
	await self.new_requests_event.wait()
	self.new_requests_event.clear()
	if not self.has_new_requests():
	self.new_requests_event.clear()
	if not self.has_new_requests():
	await self.new_requests_event.wait()

[Yard1]

Hmm can you explain why we should do it like that?

[njhill]

Just to avoid flip-flopping the event - it only needs to be cleared when you're actually about to wait on it. But I guess with python/asyncio it doesn't matter anyway.

[Yard1]

Yeah I think it should be fine