[NextPluggableDevice]fix NextPluggableDevice allocator(use pjrt allocator) and Sync() #59898
Conversation
jzhoulon
changed the title
google-ml-butler
bot
added
kokoro:force-run
|
@jyingl3 I saw the PR is stuck with some internal checks Failed, if there is something need to be fixed in my PR, please let me know, thanks. |
|
Thanks Zhoulong! Could you update tensorflow/tensorflow/core/common_runtime/next_pluggable_device/BUILD to include tensorflow/core/tfrt/common:async_value_tensor? Sorry not noticing it earlier. Thank you! |
@jyingl3 done, please have a look. thanks |
google-ml-butler
bot
added
kokoro:force-run
|
Hi Zhoulong, we are actively testing some TPU implementation that uses NPD to meet deadline and worry this change may cause breakage. Hope you don't mind if we withhold merging this change temporarily. |
sure, we will wait. Thanks |
gbaned
removed
awaiting review
|
hi, @jyingl3 , is there any update on this pr? thanks |
| device_ordinal_(options.device_ordinal), | ||
| compilation_device_type_(options.compilation_device_name) { | ||
| allocator_ = std::make_unique<NextPluggableDeviceAllocator>(device_ordinal_); | ||
| if (absl::GetFlag(FLAGS_next_pluggable_device_use_pjrt)) { |
There was a problem hiding this comment.
Do you mind adding a new flag FLAGS_next_pluggable_device_use_pjrt_allocator to guard this behavior?
There was a problem hiding this comment.
@jyingl3 done, please have a review. thanks very much.
There was a problem hiding this comment.
Approved the change. Thank you so much!
|
Thank you Zhoulong and sorry for the delay. Could you add a new flag so that the allocator can be switched? Left a comment about that. Thank you!! |
google-ml-butler
bot
added
kokoro:force-run
6 participants
use AsyncValueAllocator if context is using PJRT runtime. And fix Sync() to call done callback to avoid hang issue in direct session.