- Adds support to optionally adds replica as part of job name

PiperOrigin-RevId: 351419660
Change-Id: I59d6f801f40fd598e51f80284b4db0064f86e5c5

tensorflow/c/eager/c_api.cc

@@ -63,7 +63,7 @@ limitations under the License.
 // PLATFORM_GOOGLE, IS_MOBILE_PLATFORM, etc.
 #if defined(PLATFORM_GOOGLE) && !defined(LIBTPU_ON_GCE)
 #include "tensorflow/core/tfrt/eager/c_api_tfrt.h"
-#include "tensorflow/core/tfrt/eager/c_api_tfrt_distributed.h"
+#include "tensorflow/core/tfrt/eager/c_api_tfrt_distributed_impl.h"
 #endif  // PLATFORM_GOOGLE && !LIBTPU_ON_GCE
 
 #if !defined(IS_MOBILE_PLATFORM)
@@ -120,7 +120,7 @@ TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) {
         opts->async);
 #if !defined(IS_MOBILE_PLATFORM)
     tfrt_context->SetDistributedManager(
-        std::make_unique<tfrt::tf::DistributedManagerContextInterface>(
+        tfrt::tf::CreateDistributedManagerContext(
             tfrt_context->GetCoreRuntime()->GetHostContext()));
 #endif  // !IS_MOBILE_PLATFORM
     return tensorflow::wrap(tfrt_context);

tensorflow/core/distributed_runtime/remote_device.cc

@@ -149,4 +149,9 @@ void NewRemoteDevices(Env* env, WorkerCacheInterface* worker_cache,
                      /*fail_fast=*/false, cb);
 }
 
+std::unique_ptr<Device> NewRemoteDevice(Env* env,
+                                        DeviceAttributes device_attribute) {
+  return std::make_unique<RemoteDevice>(env, device_attribute);
+}
+
 }  // namespace tensorflow

tensorflow/core/distributed_runtime/remote_device.h

@@ -60,6 +60,9 @@ typedef std::function<void(const Status&, std::vector<Device*>*)>
 void NewRemoteDevices(Env* env, WorkerCacheInterface* worker_cache,
                       const string& worker_name, NewRemoteDevicesDone done);
 
+// Create Remote Device based on the given attributes.
+std::unique_ptr<Device> NewRemoteDevice(Env* env,
+                                        DeviceAttributes device_attribute);
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_REMOTE_DEVICE_H_