[Remote functions] Only set the default runner *after* resolving the remote FLR.
Previously, if the `runner` was not specified for a function execution, we would immediately set it to the default runner of the *local* FLR, even if the function was to be executed remotely. This change postpones the resolution of the default runner until after the function invocation has been routed to the FLR that will actually execute it. As a result, we avoid the pathological case where a GPU device using a private threadpool (TF_GPU_THREAD_MODE=gpu_private) ends up running all of the ops for the CPU-side input pipeline on the private threadpool. PiperOrigin-RevId: 195735734
Loading
Please sign in to comment