Method TryUpdateTaskIndex

oneflow/core/graph/task_id_generator.cpp:30–50

}

void TaskIdGenerator::TryUpdateTaskIndex(const HashMap<int64_t, uint32_t>& task_index_state) {
  for (auto& pair : stream_id2task_index_counter_) {
    const int64_t i64_stream_id = EncodeStreamIdToInt64(pair.first);
    uint32_t initial_task_index = 0;
    if (task_index_state.count(i64_stream_id) != 0) {
      initial_task_index = task_index_state.at(i64_stream_id);
    }
    pair.second = std::max(pair.second, initial_task_index);
  }

  // try update the task_index_init_state
  for (const auto& pair : task_index_state) {
    const auto& key = pair.first;
    const auto& val = pair.second;
    if (task_index_init_state_.count(key) != 0) {
      task_index_init_state_[key] = std::max(task_index_init_state_.at(key), val);
    } else {
      task_index_init_state_[key] = val;
    }
  }
}

TaskId TaskIdGenerator::Generate(const StreamId& stream_id) {
  std::unique_lock<std::mutex> lock(mutex_);