| 1522 | } |
| 1523 | |
| 1524 | Status DirectSession::PRunSetup(const std::vector<string>& input_names, |
| 1525 | const std::vector<string>& output_names, |
| 1526 | const std::vector<string>& target_nodes, |
| 1527 | string* handle) { |
| 1528 | TF_RETURN_IF_ERROR(CheckNotClosed()); |
| 1529 | TF_RETURN_IF_ERROR(CheckGraphCreated("PRunSetup()")); |
| 1530 | |
| 1531 | // RunOptions is not available in PRunSetup, so use thread pool 0. |
| 1532 | thread::ThreadPool* pool = thread_pools_[0].first; |
| 1533 | |
| 1534 | // Check if we already have an executor for these arguments. |
| 1535 | ExecutorsAndKeys* executors_and_keys; |
| 1536 | // TODO(cais): TFDBG support for partial runs. |
| 1537 | DebugOptions debug_options; |
| 1538 | RunStateArgs run_state_args(debug_options); |
| 1539 | run_state_args.is_partial_run = true; |
| 1540 | TF_RETURN_IF_ERROR(GetOrCreateExecutors(input_names, output_names, |
| 1541 | target_nodes, &executors_and_keys, |
| 1542 | &run_state_args)); |
| 1543 | |
| 1544 | // Create the run state and save it for future PRun calls. |
| 1545 | Executor::Args args; |
| 1546 | if (run_in_caller_thread_) { |
| 1547 | args.executor_policy = ExecutorPolicy::USE_INLINE_EXECUTOR; |
| 1548 | } else if (run_cost_model_executor_) { |
| 1549 | args.executor_policy = ExecutorPolicy::USE_COST_MODEL_EXECUTOR; |
| 1550 | } else { |
| 1551 | args.executor_policy = ExecutorPolicy::USE_NORMAL_EXECUTOR; |
| 1552 | } |
| 1553 | args.step_id = step_id_counter_.fetch_add(1); |
| 1554 | RunState* run_state = |
| 1555 | new RunState(input_names, output_names, args.step_id, &devices_); |
| 1556 | run_state->rendez = new IntraProcessRendezvous(device_mgr_); |
| 1557 | { |
| 1558 | mutex_lock l(executor_lock_); |
| 1559 | if (!partial_runs_ |
| 1560 | .emplace(run_state_args.handle, |
| 1561 | std::unique_ptr<RunState>(run_state)) |
| 1562 | .second) { |
| 1563 | return errors::Internal("The handle '", run_state_args.handle, |
| 1564 | "' created for this partial run is not unique."); |
| 1565 | } |
| 1566 | } |
| 1567 | |
| 1568 | // Start parallel Executors. |
| 1569 | const size_t num_executors = executors_and_keys->items.size(); |
| 1570 | ExecutorBarrier* barrier = new ExecutorBarrier( |
| 1571 | num_executors, run_state->rendez, [run_state](const Status& ret) { |
| 1572 | if (!ret.ok()) { |
| 1573 | mutex_lock l(run_state->mu_); |
| 1574 | run_state->status.Update(ret); |
| 1575 | } |
| 1576 | run_state->executors_done.Notify(); |
| 1577 | }); |
| 1578 | |
| 1579 | args.rendezvous = run_state->rendez; |
| 1580 | args.global_rendezvous = run_state->rendez; |
| 1581 | args.cancellation_manager = cancellation_manager_; |
nothing calls this directly
no test coverage detected