Method onTrigger

extensions/tensorflow/TFApplyGraph.cpp:80–189

}

void TFApplyGraph::onTrigger(const std::shared_ptr<core::ProcessContext> &context,
                             const std::shared_ptr<core::ProcessSession> &session) {
  auto flow_file = session->get();

  if (!flow_file) {
    return;
  }

  try {
    // Read graph
    std::string tf_type;
    flow_file->getAttribute("tf.type", tf_type);

    std::shared_ptr<tensorflow::GraphDef> graph_def;
    uint32_t graph_version;

    {
      std::lock_guard<std::mutex> guard(graph_def_mtx_);

      if ("graph" == tf_type) {
        logger_->log_info("Reading new graph def");
        graph_def_ = std::make_shared<tensorflow::GraphDef>();
        GraphReadCallback graph_cb(graph_def_);
        session->read(flow_file, &graph_cb);
        graph_version_++;
        logger_->log_info("Read graph version: %i", graph_version_);
        session->remove(flow_file);
        return;
      }

      graph_version = graph_version_;
      graph_def = graph_def_;
    }

    if (!graph_def) {
      logger_->log_error("Cannot process input because no graph has been defined");
      session->transfer(flow_file, Retry);
      return;
    }

    // Use an existing context, if one is available
    std::shared_ptr<TFContext> ctx;

    if (tf_context_q_.try_dequeue(ctx)) {
      logger_->log_debug("Using available TensorFlow context");

      if (ctx->graph_version != graph_version) {
        logger_->log_info("Allowing session with stale graph to expire");
        ctx = nullptr;
      }
    }

    if (!ctx) {
      logger_->log_info("Creating new TensorFlow context");
      tensorflow::SessionOptions options;
      ctx = std::make_shared<TFContext>();
      ctx->tf_session.reset(tensorflow::NewSession(options));
      ctx->graph_version = graph_version;