Status GrpcServer::Init(
ServiceInitFunction service_func,
const RendezvousMgrCreationFunction& rendezvous_mgr_func,
- const WorkerCreationFunction& worker_func) {
+ const WorkerCreationFunction& worker_func,
+ const StatsPublisherFactory& stats_factory) {
mutex_lock l(mu_);
CHECK_EQ(state_, NEW);
master_env_.env = env_;
master_env_.ops = OpRegistry::Global();
master_env_.worker_cache = worker_cache;
master_env_.master_session_factory =
- [config](
+ [config, stats_factory](
SessionOptions options, const MasterEnv* env,
std::unique_ptr<std::vector<std::unique_ptr<Device>>> remote_devs,
std::unique_ptr<WorkerCacheInterface> worker_cache,
options.config.MergeFrom(config);
return new MasterSession(options, env, std::move(remote_devs),
std::move(worker_cache), std::move(device_set),
- CreateNoOpStatsPublisher);
+ stats_factory);
};
master_env_.worker_cache_factory =
[this](const WorkerCacheFactoryOptions& options,
Status GrpcServer::Init(
ServiceInitFunction service_func,
+ const RendezvousMgrCreationFunction& rendezvous_mgr_func,
+ const WorkerCreationFunction& worker_func) {
+ return Init(std::move(service_func), rendezvous_mgr_func, worker_func,
+ CreateNoOpStatsPublisher);
+}
+
+Status GrpcServer::Init(
+ ServiceInitFunction service_func,
const RendezvousMgrCreationFunction& rendezvous_mgr_func) {
return Init(service_func, rendezvous_mgr_func, nullptr);
}
#include "grpc++/security/credentials.h"
#include "tensorflow/core/common_runtime/process_util.h"
+#include "tensorflow/core/common_runtime/stats_publisher_interface.h"
#include "tensorflow/core/distributed_runtime/master_env.h"
#include "tensorflow/core/distributed_runtime/rpc/async_service_interface.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_channel.h"
protected:
Status Init(ServiceInitFunction service_func,
const RendezvousMgrCreationFunction& rendezvous_mgr_func,
+ const WorkerCreationFunction& worker_func,
+ const StatsPublisherFactory& stats_factory);
+
+ Status Init(ServiceInitFunction service_func,
+ const RendezvousMgrCreationFunction& rendezvous_mgr_func,
const WorkerCreationFunction& worker_func);
Status Init(ServiceInitFunction service_func,