#include <TBBSession.h>

Inheritance diagram for tensorflow::TBBSession:

Classes
struct	ExecutorsAndKeys

struct	FunctionInfo

struct	PerPartitionExecutorsAndLib

struct	RunState

struct	RunStateArgs

Public Types
typedef std::function< void(Session *)>	CloseCallback

typedef std::vector< std::pair< string, Tensor > >	NamedTensorList

typedef std::unordered_map< StringPiece, Node *, StringPieceHasher >	NameNodeMap

Public Member Functions
::tensorflow::Status	Close () override

::tensorflow::Status	Create (const GraphDef &graph) override

void	ExportCostModels (CostModelManager::CostModelMap *cost_models)

::tensorflow::Status	Extend (const GraphDef &graph) override

::tensorflow::Status	ListDevices (std::vector< DeviceAttributes > *response) override

::tensorflow::Status	LocalDeviceManager (const DeviceMgr **output) override

::tensorflow::Status	Reset (const std::vector< string > &containers)

::tensorflow::Status	Run (const NamedTensorList &inputs, const std::vector< string > &output_names, const std::vector< string > &target_nodes, std::vector< Tensor > *outputs) override

::tensorflow::Status	Run (const ::tensorflow::RunOptions &run_options, const NamedTensorList &inputs, const std::vector< string > &output_names, const std::vector< string > &target_nodes, std::vector< Tensor > outputs, RunMetadata run_metadata) override

	TBBSession (const SessionOptions &options, const DeviceMgr device_mgr, TBBSessionFactory factory)

	~TBBSession () override

Private Member Functions
::tensorflow::Status	CheckNotClosed ()

::tensorflow::Status	CreateDebuggerState (const DebugOptions &debug_options, int64 session_run_index, int64 executor_step_index, const std::vector< string > &input_names, const std::vector< string > &output_names, const std::vector< string > &target_names, std::unique_ptr< DebuggerStateInterface > *debugger_state)

::tensorflow::Status	CreateGraphs (const BuildGraphOptions &options, std::unordered_map< string, std::unique_ptr< Graph >> outputs, std::unique_ptr< FunctionLibraryDefinition > flib_def, RunStateArgs run_state_args, DataTypeVector input_types, DataTypeVector *output_types)

::tensorflow::Status	DecorateAndPublishGraphForDebug (const DebugOptions &debug_options, Graph graph, Device device)

::tensorflow::Status	ExtendLocked (const GraphDef &graph) EXCLUSIVE_LOCKS_REQUIRED(graph_def_lock_)

::tensorflow::Status	GetOrCreateExecutors (gtl::ArraySlice< string > inputs, gtl::ArraySlice< string > outputs, gtl::ArraySlice< string > target_nodes, ExecutorsAndKeys *executors_and_keys, RunStateArgs run_state_args)

bool graph_created_	GUARDED_BY (graph_def_lock_)

GraphDef graph_def_	GUARDED_BY (graph_def_lock_)

std::vector< std::unique_ptr< FunctionInfo > > functions_	GUARDED_BY (executor_lock_)

std::unordered_map< string, std::shared_ptr< ExecutorsAndKeys > > executors_	GUARDED_BY (executor_lock_)

std::unordered_map< string, std::unique_ptr< RunState > > partial_runs_	GUARDED_BY (executor_lock_)

std::unordered_map< string, string > stateful_placements_	GUARDED_BY (graph_def_lock_)

std::unique_ptr< GraphExecutionState > execution_state_	GUARDED_BY (graph_def_lock_)

bool closed_	GUARDED_BY (closed_lock_)

Status	MaybeInitializeExecutionState (const GraphDef &graph, bool *out_already_initialized) EXCLUSIVE_LOCKS_REQUIRED(graph_def_lock_)

::tensorflow::Status	ResourceHandleToInputTensor (const Tensor &resource_tensor, Tensor *retrieved_tensor)

void	SchedClosure (tbb::task_arena &arena, tbb::task_group &g, std::function< void()> c)

	TF_DISALLOW_COPY_AND_ASSIGN (TBBSession)

::tensorflow::Status	WaitForNotification (Notification *n, int64 timeout_in_ms)

void	WaitForNotification (tbb::task_arena &arena, tbb::task_group &group, RunState run_state, CancellationManager cm, int64 timeout_in_ms)

Private Attributes
CancellationManager *	cancellation_manager_

mutex	closed_lock_

CostModelManager	cost_model_manager_

const std::unique_ptr< const DeviceMgr >	device_mgr_

DeviceSet	device_set_

std::vector< Device * >	devices_

std::atomic< int64 >	edge_name_counter_ = {0}

mutex	executor_lock_

TBBSessionFactory *const	factory_

std::unique_ptr< FunctionLibraryDefinition >	flib_def_

mutex	graph_def_lock_

std::atomic< int64 >	handle_name_counter_ = {0}

Status	init_error_

Executor::Args::NodeOutputsCallback	node_outputs_callback_ = nullptr

const int64	operation_timeout_in_ms_ = 0

const SessionOptions	options_

string	session_handle_

SessionState	session_state_

bool	sync_on_finish_ = true

Static Private Attributes
static std::atomic_int_fast64_t	step_id_counter_

Friends
class	DebugGateway

Detailed Description

Definition at line 81 of file TBBSession.h.

Member Typedef Documentation

typedef std::function<void(Session*)> tensorflow::TBBSession::CloseCallback

Definition at line 83 of file TBBSession.h.

typedef std::vector<std::pair<string, Tensor> > tensorflow::TBBSession::NamedTensorList

Definition at line 93 of file TBBSession.h.

typedef std::unordered_map<StringPiece, Node*, StringPieceHasher> tensorflow::TBBSession::NameNodeMap

Definition at line 94 of file TBBSession.h.

Constructor & Destructor Documentation

tensorflow::TBBSession::TBBSession	(	const SessionOptions &	options,
		const DeviceMgr *	device_mgr,
		TBBSessionFactory *	factory
	)

Definition at line 197 of file TBBSession.cc.

References edmIntegrityCheck::d, device_mgr_, device_set_, devices_, dqm::qstatus::ERROR, MessageLogger_cfi::INFO, LOG, session_handle_, btagGenBb_cfi::Status, mps_update::status, and sync_on_finish_.

     : options_(options),
       device_mgr_(device_mgr),
       factory_(factory),
       cancellation_manager_(new CancellationManager()),
       operation_timeout_in_ms_(options_.config.operation_timeout_in_ms()) {
   // The default value of sync_on_finish will be flipped soon and this
   // environment variable will be removed as well.
   const Status status =
       ReadBoolFromEnvVar("TF_SYNC_ON_FINISH", true, &sync_on_finish_);
   if (!status.ok()) {
     LOG(ERROR) << status.error_message();
   }
   // NOTE(mrry): We do not need to use a unique string for the session
   // handle, because TBBSession owns its devices. This may change
   // in future versions.
   session_handle_ = "tbb";
   int devices_added = 0;
   if (options.config.log_device_placement()) {
     const string mapping_str = device_mgr_->DeviceMappingString();
     if (mapping_str.empty()) {
       printf("Device mapping: no known devices.\n");
     } else {
       printf("Device mapping:\n%s", mapping_str.c_str());
     }
     LOG(INFO) << "Device mapping:\n" << mapping_str;
   }
   for (auto d : device_mgr_->ListDevices()) {
     devices_.push_back(d);
     device_set_.AddDevice(d);
     d->op_segment()->AddHold(session_handle_);
 
     // The first device added is special: it is the 'client device' (a
     // CPU device) from which we feed and fetch Tensors.
     if (devices_added == 0) {
       device_set_.set_client_device(d);
     }
     ++devices_added;
   }
 }

tensorflow::TBBSession::~TBBSession ( )

override

Definition at line 240 of file TBBSession.cc.

References cancellation_manager_, Close(), edmIntegrityCheck::d, device_mgr_, flib_def_, and session_handle_.

                         {
   if (!closed_) Close().IgnoreError();
   for (auto& it : partial_runs_) {
     it.second.reset(nullptr);
   }
   for (auto& it : executors_) {
     it.second.reset();
   }
   for (auto d : device_mgr_->ListDevices()) {
     d->op_segment()->RemoveHold(session_handle_);
   }
   for (auto d : device_mgr_->ListDevices()) {
     d->ClearResourceMgr();
   }
   functions_.clear();
   delete cancellation_manager_;
 
   execution_state_.reset(nullptr);
   flib_def_.reset(nullptr);
 }

Member Function Documentation

::tensorflow::Status tensorflow::TBBSession::CheckNotClosed ( )

inlineprivate

Definition at line 251 of file TBBSession.h.

References pfDeepBoostedJetPreprocessParams_cfi::input_names, checklumidiff::l, and btagGenBb_cfi::Status.

Referenced by DecorateAndPublishGraphForDebug(), and Extend().

                                     {
     mutex_lock l(closed_lock_);
     if (closed_) return errors::Cancelled("Session has been closed.");
     return ::tensorflow::Status::OK();
   }

tensorflow::Status tensorflow::TBBSession::Close ( )

override

Definition at line 1080 of file TBBSession.cc.

References cancellation_manager_, closed_lock_, tensorflow::TBBSessionFactory::Deregister(), factory_, and checklumidiff::l.

Referenced by ~TBBSession().

                                    {
   cancellation_manager_->StartCancel();
   {
     mutex_lock l(closed_lock_);
     if (closed_) return ::tensorflow::Status::OK();
     closed_ = true;
   }
   if (factory_ != nullptr) factory_->Deregister(this);
   return ::tensorflow::Status::OK();
 }

Status tensorflow::TBBSession::Create ( const GraphDef & graph )

override

Definition at line 292 of file TBBSession.cc.

References ExtendLocked(), graph_def_lock_, init_error_, and checklumidiff::l.

                                                {
   TF_RETURN_IF_ERROR(init_error_);
   if (graph.node_size() > 0) {
     mutex_lock l(graph_def_lock_);
     if (graph_created_) {
       return errors::AlreadyExists(
           "A Graph has already been created for this session.");
     }
     return ExtendLocked(graph);
   }
   return Status::OK();
 }

Status tensorflow::TBBSession::CreateDebuggerState	(	const DebugOptions &	debug_options,
		int64	session_run_index,
		int64	executor_step_index,
		const std::vector< string > &	input_names,
		const std::vector< string > &	output_names,
		const std::vector< string > &	target_names,
		std::unique_ptr< DebuggerStateInterface > *	debugger_state
	)

private

Definition at line 335 of file TBBSession.cc.

References pfDeepBoostedJetPreprocessParams_cfi::input_names.

Referenced by DecorateAndPublishGraphForDebug().

                                                            {
   TF_RETURN_IF_ERROR(
       DebuggerStateRegistry::CreateState(debug_options, debugger_state));
   TF_RETURN_IF_ERROR(debugger_state->get()->PublishDebugMetadata(
       debug_options.global_step(), session_run_index, executor_step_index,
       input_names, output_names, target_names));
   return Status::OK();
 }

Status tensorflow::TBBSession::CreateGraphs	(	const BuildGraphOptions &	options,
		std::unordered_map< string, std::unique_ptr< Graph >> *	outputs,
		std::unique_ptr< FunctionLibraryDefinition > *	flib_def,
		RunStateArgs *	run_state_args,
		DataTypeVector *	input_types,
		DataTypeVector *	output_types
	)

private

Definition at line 908 of file TBBSession.cc.

References KineDebug3::count(), edmIntegrityCheck::d, device_mgr_, device_set_, devices_, edge_name_counter_, flib_def_, tensorflow::TBBSession::RunStateArgs::graph, graph_def_lock_, tensorflow::TBBSession::RunStateArgs::is_partial_run, checklumidiff::l, eostools::move(), dataset::name, options_, PatBasicFWLiteJetAnalyzer_Selector_cfg::outputs, tablePrinter::prefix, alignCSCRings::s, btagGenBb_cfi::Status, and std::swap().

Referenced by GetOrCreateExecutors().

                                   {
   mutex_lock l(graph_def_lock_);
   std::unique_ptr<ClientGraph> client_graph;
 
   std::unique_ptr<GraphExecutionState> temp_exec_state_holder;
   GraphExecutionState* execution_state = nullptr;
   if (options_.config.graph_options().place_pruned_graph()) {
     // Because we are placing pruned graphs, we need to create a
     // new GraphExecutionState for every new unseen graph,
     // and then place it.
     GraphExecutionStateOptions prune_options;
     prune_options.device_set = &device_set_;
     prune_options.session_options = &options_;
     prune_options.stateful_placements = stateful_placements_;
     TF_RETURN_IF_ERROR(GraphExecutionState::MakeForPrunedGraph(
         execution_state_->original_graph_def().library(), prune_options,
         execution_state_->original_graph_def(), subgraph_options,
         &temp_exec_state_holder, &client_graph));
     execution_state = temp_exec_state_holder.get();
   } else {
     execution_state = execution_state_.get();
     TF_RETURN_IF_ERROR(
         execution_state->BuildGraph(subgraph_options, &client_graph));
   }
 
   if (subgraph_options.feed_endpoints.size() !=
       client_graph->feed_types.size()) {
     return errors::Internal(
         "Graph pruning failed: requested number of feed endpoints = ",
         subgraph_options.feed_endpoints.size(),
         " versus number of pruned feed endpoints = ",
         client_graph->feed_types.size());
   }
   if (subgraph_options.fetch_endpoints.size() !=
       client_graph->fetch_types.size()) {
     return errors::Internal(
         "Graph pruning failed: requested number of fetch endpoints = ",
         subgraph_options.fetch_endpoints.size(),
         " versus number of pruned fetch endpoints = ",
         client_graph->fetch_types.size());
   }
 
   auto current_stateful_placements = execution_state->GetStatefulPlacements();
   // Update our current state based on the execution_state's
   // placements.  If there are any mismatches for a node,
   // we should fail, as this should never happen.
   for (auto placement_pair : current_stateful_placements) {
     const string& node_name = placement_pair.first;
     const string& placement = placement_pair.second;
     auto iter = stateful_placements_.find(node_name);
     if (iter == stateful_placements_.end()) {
       stateful_placements_.insert(std::make_pair(node_name, placement));
     } else if (iter->second != placement) {
       return errors::Internal(
           "Stateful placement mismatch. "
           "Current assignment of ",
           node_name, " to ", iter->second, " does not match ", placement);
     }
   }
 
   stateful_placements_ = execution_state->GetStatefulPlacements();
 
   // Remember the graph in run state if this is a partial run.
   if (run_state_args->is_partial_run) {
     run_state_args->graph.reset(new Graph(flib_def_.get()));
     CopyGraph(*execution_state->full_graph(), run_state_args->graph.get());
   }
 
   // Partition the graph across devices.
   PartitionOptions popts;
   popts.node_to_loc = [](const Node* node) {
     assert(node != nullptr);
     return node->assigned_device_name();
   };
   popts.new_name = [this](const string& prefix) {
     return strings::StrCat(prefix, "/_", edge_name_counter_.fetch_add(1));
   };
   popts.get_incarnation = [](const string& name) {
     // The direct session does not have changing incarnation numbers.
     // Just return '1'.
     return 1;
   };
   popts.flib_def = &client_graph->graph.flib_def();
   popts.control_flow_added = false;
 
   std::unordered_map<string, GraphDef> partitions;
   TF_RETURN_IF_ERROR(Partition(popts, &client_graph->graph, &partitions));
 
   std::vector<string> device_names;
   for (auto device : devices_) {
     // Extract the LocalName from the device.
     device_names.push_back(DeviceNameUtils::LocalName(device->name()));
   }
 
   // Check for valid partitions.
   for (const auto& partition : partitions) {
     const string local_partition_name =
         DeviceNameUtils::LocalName(partition.first);
     if (std::count(device_names.begin(), device_names.end(),
                    local_partition_name) == 0) {
       return errors::InvalidArgument(
           "Creating a partition for ", local_partition_name,
           " which doesn't exist in the list of available devices. Available "
           "devices: ",
           str_util::Join(device_names, ","));
     }
   }
 
   for (const auto& partition : partitions) {
     std::unique_ptr<Graph> device_graph(
         new Graph(client_graph->flib_def.get()));
     GraphConstructorOptions device_opts;
     // There are internal operations (e.g., send/recv) that we now allow.
     device_opts.allow_internal_ops = true;
     device_opts.expect_device_spec = true;
     TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(device_opts, partition.second,
                                               device_graph.get()));
     outputs->emplace(partition.first, std::move(device_graph));
   }
 
   GraphOptimizationPassOptions optimization_options;
   optimization_options.session_options = &options_;
   optimization_options.flib_def = client_graph->flib_def.get();
   optimization_options.partition_graphs = outputs;
   TF_RETURN_IF_ERROR(OptimizationPassRegistry::Global()->RunGrouping(
       OptimizationPassRegistry::POST_PARTITIONING, optimization_options));
 
   Status s;
   for (auto& partition : *outputs) {
     const string& partition_name = partition.first;
     std::unique_ptr<Graph>* graph = &partition.second;
 
     VLOG(2) << "Created " << DebugString(graph->get()) << " for "
             << partition_name;
 
     // Give the device an opportunity to rewrite its subgraph.
     Device* d;
     s = device_mgr_->LookupDevice(partition_name, &d);
     if (!s.ok()) break;
     s = d->MaybeRewriteGraph(graph);
     if (!s.ok()) {
       break;
     }
   }
   *flib_def = std::move(client_graph->flib_def);
   std::swap(*input_types, client_graph->feed_types);
   std::swap(*output_types, client_graph->fetch_types);
   return s;
 }

Status tensorflow::TBBSession::DecorateAndPublishGraphForDebug	(	const DebugOptions &	debug_options,
		Graph *	graph,
		Device *	device
	)

private

Definition at line 349 of file TBBSession.cc.

Referenced by GetOrCreateExecutors().

                                                                      {
   std::unique_ptr<DebugGraphDecoratorInterface> decorator;
   TF_RETURN_IF_ERROR(
       DebugGraphDecoratorRegistry::CreateDecorator(debug_options, &decorator));
 
   TF_RETURN_IF_ERROR(decorator->DecorateGraph(graph, device));
   TF_RETURN_IF_ERROR(decorator->PublishGraph(*graph, device->name()));
   return Status::OK();
 }

void tensorflow::TBBSession::ExportCostModels ( CostModelManager::CostModelMap * cost_models )

inline

Definition at line 123 of file TBBSession.h.

                                                                    {
     cost_model_manager_.ExportCostModels(cost_models);
   }

Status tensorflow::TBBSession::Extend ( const GraphDef & graph )

override

Definition at line 305 of file TBBSession.cc.

References CheckNotClosed(), ExtendLocked(), graph_def_lock_, and checklumidiff::l.

                                                {
   TF_RETURN_IF_ERROR(CheckNotClosed());
   mutex_lock l(graph_def_lock_);
   return ExtendLocked(graph);
 }

Status tensorflow::TBBSession::ExtendLocked ( const GraphDef & graph )

private

Definition at line 311 of file TBBSession.cc.

References flib_def_, and MaybeInitializeExecutionState().

Referenced by Create(), and Extend().

                                                      {
   bool already_initialized;
   // If this is the first call, we can initialize the execution state
   // with `graph` and do not need to call `Extend()`.
   TF_RETURN_IF_ERROR(
       MaybeInitializeExecutionState(graph, &already_initialized));
   if (already_initialized) {
     TF_RETURN_IF_ERROR(flib_def_->AddLibrary(graph.library()));
     std::unique_ptr<GraphExecutionState> state;
     TF_RETURN_IF_ERROR(execution_state_->Extend(graph, &state));
     execution_state_.swap(state);
   }
   return Status::OK();
 }

Status tensorflow::TBBSession::GetOrCreateExecutors	(	gtl::ArraySlice< string >	inputs,
		gtl::ArraySlice< string >	outputs,
		gtl::ArraySlice< string >	target_nodes,
		ExecutorsAndKeys **	executors_and_keys,
		RunStateArgs *	run_state_args
	)

private

Definition at line 667 of file TBBSession.cc.

References CreateGraphs(), tensorflow::TBBSession::RunStateArgs::debug_options, DecorateAndPublishGraphForDebug(), device_mgr_, device_set_, executor_lock_, plotBeamSpotDB::first, tensorflow::TBBSession::RunStateArgs::graph, graph_def_lock_, cuy::graphs, tensorflow::TBBSession::RunStateArgs::handle, handle_name_counter_, mps_fire::i, triggerObjects_cff::id, input, tensorflow::TBBSession::RunStateArgs::is_partial_run, crabWrapper::key, checklumidiff::l, mps_check::lib, eostools::move(), gen::n, names, node_outputs_callback_, AlcaSiPixelAliHarvester0T_cff::options, options_, convertSQLitetoXML_cfg::output, and session_handle_.

Referenced by DecorateAndPublishGraphForDebug().

                                   {
   int64 handle_name_counter_value = -1;
   if (LogMemory::IsEnabled() || run_state_args->is_partial_run) {
     handle_name_counter_value = handle_name_counter_.fetch_add(1);
   }
 
   string debug_tensor_watches_summary;
   if (!run_state_args->debug_options.debug_tensor_watch_opts().empty()) {
     debug_tensor_watches_summary = SummarizeDebugTensorWatches(
         run_state_args->debug_options.debug_tensor_watch_opts());
   }
 
   // Fast lookup path, no sorting.
   const string key = strings::StrCat(
       str_util::Join(inputs, ","), "->", str_util::Join(outputs, ","), "/",
       str_util::Join(target_nodes, ","), "/", run_state_args->is_partial_run,
       "/", debug_tensor_watches_summary);
   // Set the handle, if it's needed to log memory or for partial run.
   if (handle_name_counter_value >= 0) {
     run_state_args->handle =
         strings::StrCat(key, ";", handle_name_counter_value);
   }
 
   // See if we already have the executors for this run.
   {
     mutex_lock l(executor_lock_);  // could use reader lock
     auto it = executors_.find(key);
     if (it != executors_.end()) {
       *executors_and_keys = it->second.get();
       return Status::OK();
     }
   }
 
   // Slow lookup path, the unsorted key missed the cache.
   // Sort the inputs and outputs, and look up with the sorted key in case an
   // earlier call used a different order of inputs and outputs.
   //
   // We could consider some other signature instead of sorting that
   // preserves the same property to avoid the sort in the future.
   std::vector<string> inputs_sorted(inputs.begin(), inputs.end());
   std::sort(inputs_sorted.begin(), inputs_sorted.end());
   std::vector<string> outputs_sorted(outputs.begin(), outputs.end());
   std::sort(outputs_sorted.begin(), outputs_sorted.end());
   std::vector<string> tn_sorted(target_nodes.begin(), target_nodes.end());
   std::sort(tn_sorted.begin(), tn_sorted.end());
 
   const string sorted_key = strings::StrCat(
       str_util::Join(inputs_sorted, ","), "->",
       str_util::Join(outputs_sorted, ","), "/", str_util::Join(tn_sorted, ","),
       "/", run_state_args->is_partial_run, "/", debug_tensor_watches_summary);
   // Set the handle, if its needed to log memory or for partial run.
   if (handle_name_counter_value >= 0) {
     run_state_args->handle =
         strings::StrCat(sorted_key, ";", handle_name_counter_value);
   }
 
   // See if we already have the executors for this run.
   {
     mutex_lock l(executor_lock_);
     auto it = executors_.find(sorted_key);
     if (it != executors_.end()) {
       *executors_and_keys = it->second.get();
       // Insert this under the original key.
       executors_.emplace(key, it->second);
       return Status::OK();
     }
   }
 
   // Nothing found, so create the executors and store in the cache.
   BuildGraphOptions options;
   options.feed_endpoints = inputs_sorted;
   options.fetch_endpoints = outputs_sorted;
   options.target_nodes = tn_sorted;
   options.use_function_convention = !run_state_args->is_partial_run;
   if (!run_state_args->debug_options.debug_tensor_watch_opts().empty()) {
     options.debug_options = run_state_args->debug_options;
   }
 
   std::unique_ptr<FunctionInfo> func_info(new FunctionInfo);
   std::shared_ptr<ExecutorsAndKeys> ek(new ExecutorsAndKeys);
 
   // The executor_lock_ is intentionally released while executor is
   // being created.
   std::unordered_map<string, std::unique_ptr<Graph>> graphs;
   TF_RETURN_IF_ERROR(CreateGraphs(options, &graphs, &func_info->flib_def,
                                   run_state_args, &ek->input_types,
                                   &ek->output_types));
 
   if (run_state_args->is_partial_run) {
     ek->graph = std::move(run_state_args->graph);
     std::unordered_set<StringPiece, StringPieceHasher> names;
     for (const string& input : inputs) {
       TensorId id(ParseTensorName(input));
       names.emplace(id.first);
     }
     for (const string& output : outputs) {
       TensorId id(ParseTensorName(output));
       names.emplace(id.first);
     }
     for (Node* n : ek->graph->nodes()) {
       if (names.count(n->name()) > 0) {
         ek->name_to_node.insert({n->name(), n});
       }
     }
   }
   ek->items.reserve(graphs.size());
   const auto& optimizer_opts =
       options_.config.graph_options().optimizer_options();
 
   int graph_def_version;
   {
     mutex_lock l(graph_def_lock_);
     graph_def_version =
         execution_state_->original_graph_def().versions().producer();
   }
   func_info->proc_flr.reset(new ProcessFunctionLibraryRuntime(
       device_mgr_.get(), options_.env, graph_def_version,
       func_info->flib_def.get(), optimizer_opts));
 
   GraphOptimizer optimizer(optimizer_opts);
   for (auto iter = graphs.begin(); iter != graphs.end(); ++iter) {
     const string& partition_name = iter->first;
     std::unique_ptr<Graph>& partition_graph = iter->second;
 
     Device* device;
     TF_RETURN_IF_ERROR(device_mgr_->LookupDevice(partition_name, &device));
 
     ek->items.resize(ek->items.size() + 1);
     auto* item = &(ek->items.back());
     auto lib = func_info->proc_flr->GetFLR(partition_name);
     if (lib == nullptr) {
       return errors::Internal("Could not find device: ", partition_name);
     }
     item->flib = lib;
 
     LocalExecutorParams params;
     params.device = device;
     params.function_library = lib;
     auto opseg = device->op_segment();
     params.create_kernel = [this, lib, opseg](const NodeDef& ndef,
                                               OpKernel** kernel) {
       // We do not share the kernel via the OpSegment if the node is
       // stateless, or a function.
       // NOTE(mrry): We must not share function kernels (implemented
       // using `CallOp`) between subgraphs, because `CallOp::handle_`
       // is tied to a particular subgraph. Even if the function itself
       // is stateful, the `CallOp` that invokes it is not.
       if (!lib->IsStateful(ndef.op()) ||
           lib->GetFunctionLibraryDefinition()->Find(ndef.op()) != nullptr) {
         return lib->CreateKernel(ndef, kernel);
       }
       auto create_fn = [lib, &ndef](OpKernel** kernel) {
         return lib->CreateKernel(ndef, kernel);
       };
       // Kernels created for subgraph nodes need to be cached.  On
       // cache miss, create_fn() is invoked to create a kernel based
       // on the function library here + global op registry.
       return opseg->FindOrCreate(session_handle_, ndef.name(), kernel,
                                  create_fn);
     };
     params.delete_kernel = [lib](OpKernel* kernel) {
       // If the node is stateful, opseg owns it. Otherwise, delete it.
       if (kernel && !lib->IsStateful(kernel->type_string())) {
         delete kernel;
       }
     };
     params.node_outputs_cb = node_outputs_callback_;
 
     optimizer.Optimize(lib, options_.env, device, &iter->second,
                        /*shape_map=*/nullptr);
 
     // EXPERIMENTAL: tfdbg inserts debug nodes in the graph.
     if (!options.debug_options.debug_tensor_watch_opts().empty()) {
       TF_RETURN_IF_ERROR(DecorateAndPublishGraphForDebug(
           options.debug_options, partition_graph.get(), params.device));
     }
 
     TF_RETURN_IF_ERROR(EnsureMemoryTypes(DeviceType(device->device_type()),
                                          device->name(),
                                          partition_graph.get()));
     // NewLocalExecutor takes ownership of partition_graph.
     item->graph = partition_graph.get();
     item->executor = nullptr;
     item->device = device;
     Executor* executor;
     TF_RETURN_IF_ERROR(
         NewLocalExecutor(params, partition_graph.release(), &executor));
     item->executor.reset(executor);
   }
 
   // Cache the mapping from input/output names to graph elements to
   // avoid recomputing it every time.
   if (!run_state_args->is_partial_run) {
     // For regular `Run()`, we use the function calling convention, and so
     // maintain a mapping from input/output names to
     // argument/return-value ordinal index.
     for (size_t i = 0; i < inputs_sorted.size(); ++i) {
       const string& input = inputs_sorted[i];
       ek->input_name_to_index[input] = i;
     }
     for (size_t i = 0; i < outputs_sorted.size(); ++i) {
       const string& output = outputs_sorted[i];
       ek->output_name_to_index[output] = i;
     }
   } else {
     // For `PRun()`, we use the rendezvous calling convention, and so
     // maintain a mapping from input/output names to rendezvous keys.
     //
     // We always use the first device as the device name portion of the
     // key, even if we're feeding another graph.
     for (size_t i = 0; i < inputs_sorted.size(); ++i) {
       const string& input = inputs_sorted[i];
       ek->input_name_to_rendezvous_key[input] = GetRendezvousKey(
           input, device_set_.client_device()->attributes(), FrameAndIter(0, 0));
     }
     for (size_t i = 0; i < outputs_sorted.size(); ++i) {
       const string& output = outputs_sorted[i];
       ek->output_name_to_rendezvous_key[output] =
           GetRendezvousKey(output, device_set_.client_device()->attributes(),
                            FrameAndIter(0, 0));
     }
   }
 
   // Reacquire the lock, try to insert into the map.
   mutex_lock l(executor_lock_);
   functions_.push_back(std::move(func_info));
 
   // Another thread may have created the entry before us, in which case we will
   // reuse the already created one.
   auto insert_result = executors_.emplace(sorted_key, ek);
   // Insert the value under the original key, so the fast path lookup will work
   // if the user uses the same order of inputs, outputs, and targets again.
   executors_.emplace(key, insert_result.first->second);
   *executors_and_keys = insert_result.first->second.get();
 
   return Status::OK();
 }

bool graph_created_ tensorflow::TBBSession::GUARDED_BY ( graph_def_lock_ )

private

GraphDef graph_def_ tensorflow::TBBSession::GUARDED_BY ( graph_def_lock_ )

private

std::vector<std::unique_ptr<FunctionInfo> > functions_ tensorflow::TBBSession::GUARDED_BY ( executor_lock_ )

private

std::unordered_map<string, std::shared_ptr<ExecutorsAndKeys> > executors_ tensorflow::TBBSession::GUARDED_BY ( executor_lock_ )

private

std::unordered_map<string, std::unique_ptr<RunState> > partial_runs_ tensorflow::TBBSession::GUARDED_BY ( executor_lock_ )

private

std::unordered_map<string, string> stateful_placements_ tensorflow::TBBSession::GUARDED_BY ( graph_def_lock_ )

private

std::unique_ptr<GraphExecutionState> execution_state_ tensorflow::TBBSession::GUARDED_BY ( graph_def_lock_ )

private

bool closed_ tensorflow::TBBSession::GUARDED_BY ( closed_lock_ )

private

tensorflow::Status tensorflow::TBBSession::ListDevices ( std::vector< DeviceAttributes > * response )

override

Definition at line 1063 of file TBBSession.cc.

References edmIntegrityCheck::d, and devices_.

                                            {
   response->clear();
   response->reserve(devices_.size());
   for (Device* d : devices_) {
     const DeviceAttributes& attrs = d->attributes();
     response->emplace_back(attrs);
   }
   return ::tensorflow::Status::OK();
 }

::tensorflow::Status tensorflow::TBBSession::LocalDeviceManager ( const DeviceMgr ** output )

inlineoverride

Definition at line 118 of file TBBSession.h.

                                                                          {
     *output = device_mgr_.get();
     return ::tensorflow::Status::OK();
   }

Status tensorflow::TBBSession::MaybeInitializeExecutionState	(	const GraphDef &	graph,
		bool *	out_already_initialized
	)

private

Definition at line 261 of file TBBSession.cc.

References device_set_, flib_def_, AlcaSiPixelAliHarvester0T_cff::options, options_, and groupFilesInBlocks::temp.

Referenced by ExtendLocked().

                                                           {
   // If already initialized, do nothing.
   if (flib_def_ && execution_state_) {
     *out_already_initialized = true;
     return Status::OK();
   }
   // Set up the per-session execution state.
   // NOTE(mrry): The function library created here will be used for
   // all subsequent extensions of the graph.
   flib_def_.reset(
       new FunctionLibraryDefinition(OpRegistry::Global(), graph.library()));
   GraphExecutionStateOptions options;
   options.device_set = &device_set_;
   options.session_options = &options_;
   // TODO(mrry,suharshs): We explicitly copy `graph` so that
   // `MakeForBaseGraph()` can take ownership of its
   // contents. Previously this happened implicitly in calls to the
   // `GraphExecutionState`. Other sessions call
   // `MakeForBaseGraph` in such a way that we can destructively read
   // the passed-in `GraphDef`. In principle we could do the same here,
   // with a wider refactoring; we might revise the direct session so
   // that it copies the graph fewer times.
   GraphDef temp(graph);
   TF_RETURN_IF_ERROR(
       GraphExecutionState::MakeForBaseGraph(&temp, options, &execution_state_));
   graph_created_ = true;
   *out_already_initialized = false;
   return Status::OK();
 }

tensorflow::Status tensorflow::TBBSession::Reset ( const std::vector< string > & containers )

Definition at line 1074 of file TBBSession.cc.

References device_mgr_.

                                          {
   device_mgr_->ClearContainers(containers);
   return ::tensorflow::Status::OK();
 }

Status tensorflow::TBBSession::ResourceHandleToInputTensor	(	const Tensor &	resource_tensor,
		Tensor *	retrieved_tensor
	)

private

Definition at line 641 of file TBBSession.cc.

References session_state_.

Referenced by DecorateAndPublishGraphForDebug().

                                                                             {
   if (resource_tensor.dtype() != DT_RESOURCE) {
     return errors::InvalidArgument(strings::StrCat(
         "ResourceHandleToInputTensor() received non-DT_RESOURCE Tensor: ",
         resource_tensor.dtype()));
   }
 
   const ResourceHandle& resource_handle =
       resource_tensor.scalar<ResourceHandle>()();
 
   if (resource_handle.container() ==
       SessionState::kTensorHandleResourceTypeName) {
     return session_state_.GetTensor(resource_handle.name(), retrieved_tensor);
   } else {
     return errors::InvalidArgument(strings::StrCat(
         "Invalid resource type hash code: ", resource_handle.hash_code(),
         "(name: ", resource_handle.name(),
         " type: ", resource_handle.maybe_type_name(),
         "). Perhaps a resource tensor was being provided as a feed? That is "
         "not currently allowed. Please file an issue at "
         "https://github.com/tensorflow/tensorflow/issues/new, ideally with a "
         "short code snippet that leads to this error message."));
   }
 }

Status tensorflow::TBBSession::Run	(	const NamedTensorList &	inputs,
		const std::vector< string > &	output_names,
		const std::vector< string > &	target_nodes,
		std::vector< Tensor > *	outputs
	)

override

Definition at line 326 of file TBBSession.cc.

Referenced by DecorateAndPublishGraphForDebug().

                                                       {
   RunMetadata run_metadata;
   return Run(RunOptions(), inputs, output_names, target_nodes, outputs,
              &run_metadata);
 }

::tensorflow::Status tensorflow::TBBSession::Run	(	const ::tensorflow::RunOptions &	run_options,
		const NamedTensorList &	inputs,
		const std::vector< string > &	output_names,
		const std::vector< string > &	target_nodes,
		std::vector< Tensor > *	outputs,
		RunMetadata *	run_metadata
	)

override

void tensorflow::TBBSession::SchedClosure	(	tbb::task_arena &	arena,
		tbb::task_group &	g,
		std::function< void()>	c
	)

private

Definition at line 193 of file TBBSession.cc.

References EnergyCorrector::c.

Referenced by DecorateAndPublishGraphForDebug().

                                                                                            {
   arena.execute( [&g, &c] () {g.run( c ); } );
 }

tensorflow::TBBSession::TF_DISALLOW_COPY_AND_ASSIGN ( TBBSession )

private

tensorflow::Status tensorflow::TBBSession::WaitForNotification	(	Notification *	n,
		int64	timeout_in_ms
	)

private

Definition at line 1156 of file TBBSession.cc.

References btagGenBb_cfi::Status.

Referenced by DecorateAndPublishGraphForDebug(), and WaitForNotification().

                                                      {
   if (timeout_in_ms > 0) {
     const int64 timeout_in_us = timeout_in_ms * 1000;
     const bool notified =
         WaitForNotificationWithTimeout(notification, timeout_in_us);
     if (!notified) {
       return Status(error::DEADLINE_EXCEEDED,
                     "Timed out waiting for notification");
     }
   } else {
     notification->WaitForNotification();
   }
   return Status::OK();
 }

void tensorflow::TBBSession::WaitForNotification	(	tbb::task_arena &	arena,
		tbb::task_group &	group,
		RunState *	run_state,
		CancellationManager *	cm,
		int64	timeout_in_ms
	)

private

Definition at line 1135 of file TBBSession.cc.

References tensorflow::TBBSession::RunState::executors_done, checklumidiff::l, tensorflow::TBBSession::RunState::mu_, btagGenBb_cfi::Status, mps_update::status, and WaitForNotification().

                                                                        {
   // Doing the wait in the arena adds this thread to the arena
   // and therefore tasks associated to the group can run on this thread
   arena.execute([&taskGroup]() { taskGroup.wait();} );
 
   const Status status =
       WaitForNotification(&run_state->executors_done, timeout_in_ms);
   if (!status.ok()) {
     {
       mutex_lock l(run_state->mu_);
       run_state->status.Update(status);
     }
     cm->StartCancel();
     // We must wait for the executors to complete, because they have borrowed
     // references to `cm` and other per-step state. After this notification, it
     // is safe to clean up the step.
     run_state->executors_done.WaitForNotification();
   }
 }