#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 92 of file TBBSession.h.

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

Definition at line 93 of file TBBSession.h.

Constructor & Destructor Documentation

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

Definition at line 195 of file TBBSession.cc.

References ztail::d, device_mgr_, device_set_, devices_, dqm::qstatus::ERROR, unpackBuffers-CaloStage2::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 235 of file TBBSession.cc.

References cancellation_manager_, Close(), ztail::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 250 of file TBBSession.h.

References pfDeepBoostedJetPreprocessParams_cfi::input_names, cmsLHEtoEOSManager::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 1044 of file TBBSession.cc.

References cancellation_manager_, closed_lock_, tensorflow::TBBSessionFactory::Deregister(), factory_, and cmsLHEtoEOSManager::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 285 of file TBBSession.cc.

References ExtendLocked(), graph_def_lock_, init_error_, and cmsLHEtoEOSManager::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 325 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 877 of file TBBSession.cc.

References KineDebug3::count(), ztail::d, device_mgr_, device_set_, devices_, edge_name_counter_, flib_def_, tensorflow::TBBSession::RunStateArgs::graph, graph_def_lock_, tensorflow::TBBSession::RunStateArgs::is_partial_run, cmsLHEtoEOSManager::l, eostools::move(), Skims_PA_cff::name, options_, PatBasicFWLiteJetAnalyzer_Selector_cfg::outputs, ZMuMuAnalysisNtupler_cff::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 338 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 121 of file TBBSession.h.

                                                                      {
       cost_model_manager_.ExportCostModels(cost_models);
     }

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

override

Definition at line 297 of file TBBSession.cc.

References CheckNotClosed(), ExtendLocked(), graph_def_lock_, and cmsLHEtoEOSManager::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 303 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 637 of file TBBSession.cc.

References CreateGraphs(), tensorflow::TBBSession::RunStateArgs::debug_options, DecorateAndPublishGraphForDebug(), device_mgr_, device_set_, executor_lock_, dqmdumpme::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, B2GTnPMonitor_cfi::item, crabWrapper::key, cmsLHEtoEOSManager::l, mps_check::lib, eostools::move(), dqmiodumpmetadata::n, names, node_outputs_callback_, AlcaSiPixelAliHarvester0T_cff::options, options_, convertSQLitetoXML_cfg::output, CalibrationSummaryClient_cfi::params, 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 1029 of file TBBSession.cc.

References ztail::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 116 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 257 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 1039 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 612 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 317 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 191 of file TBBSession.cc.

References HltBtagPostValidation_cff::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 1125 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 1102 of file TBBSession.cc.

References tensorflow::TBBSession::RunState::executors_done, cmsLHEtoEOSManager::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();
     }
   }