xref: /external/tensorflow/tensorflow/core/debug/debug_io_utils.cc

/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/debug/debug_io_utils.h"

#include <stddef.h>
#include <string.h>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <utility>
#include <vector>

#ifndef PLATFORM_WINDOWS
#include "grpcpp/create_channel.h"
#else
// winsock2.h is used in grpc, so Ws2_32.lib is needed
#pragma comment(lib, "Ws2_32.lib")
#endif  // #ifndef PLATFORM_WINDOWS

#include "tensorflow/core/debug/debug_callback_registry.h"
#include "tensorflow/core/debug/debugger_event_metadata.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/summary.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/lib/core/bits.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/util/event.pb.h"

#define GRPC_OSS_WINDOWS_UNIMPLEMENTED_ERROR \
  return errors::Unimplemented(              \
      kGrpcURLScheme, " debug URL scheme is not implemented on Windows yet.")

namespace tensorflow {

namespace {

// Creates an Event proto representing a chunk of a Tensor. This method only
// populates the field of the Event proto that represent the envelope
// information (e.g., timestamp, device_name, num_chunks, chunk_index, dtype,
// shape). It does not set the value.tensor field, which should be set by the
// caller separately.
Event PrepareChunkEventProto(const DebugNodeKey& debug_node_key,
                             const uint64 wall_time_us, const size_t num_chunks,
                             const size_t chunk_index,
                             const DataType& tensor_dtype,
                             const TensorShapeProto& tensor_shape) {
  Event event;
  event.set_wall_time(static_cast<double>(wall_time_us));
  Summary::Value* value = event.mutable_summary()->add_value();

  // Create the debug node_name in the Summary proto.
  // For example, if tensor_name = "foo/node_a:0", and the debug_op is
  // "DebugIdentity", the debug node_name in the Summary proto will be
  // "foo/node_a:0:DebugIdentity".
  value->set_node_name(debug_node_key.debug_node_name);

  // Tag by the node name. This allows TensorBoard to quickly fetch data
  // per op.
  value->set_tag(debug_node_key.node_name);

  // Store data within debugger metadata to be stored for each event.
  third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
  metadata.set_device(debug_node_key.device_name);
  metadata.set_output_slot(debug_node_key.output_slot);
  metadata.set_num_chunks(num_chunks);
  metadata.set_chunk_index(chunk_index);

  // Encode the data in JSON.
  string json_output;
  tensorflow::protobuf::util::JsonPrintOptions json_options;
  json_options.always_print_primitive_fields = true;
  auto status = tensorflow::protobuf::util::MessageToJsonString(
      metadata, &json_output, json_options);
  if (status.ok()) {
    // Store summary metadata. Set the plugin to use this data as "debugger".
    SummaryMetadata::PluginData* plugin_data =
        value->mutable_metadata()->mutable_plugin_data();
    plugin_data->set_plugin_name(DebugIO::kDebuggerPluginName);
    plugin_data->set_content(json_output);
  } else {
    LOG(WARNING) << "Failed to convert DebuggerEventMetadata proto to JSON. "
                 << "The debug_node_name is " << debug_node_key.debug_node_name
                 << ".";
  }

  value->mutable_tensor()->set_dtype(tensor_dtype);
  *value->mutable_tensor()->mutable_tensor_shape() = tensor_shape;

  return event;
}

// Translates the length of a string to number of bytes when the string is
// encoded as bytes in protobuf. Note that this makes a conservative estimate
// (i.e., an estimate that is usually too large, but never too small under the
// gRPC message size limit) of the Varint-encoded length, to workaround the lack
// of a portable length function.
const size_t StringValMaxBytesInProto(const string& str) {
#if defined(PLATFORM_GOOGLE)
  return str.size() + DebugGrpcIO::kGrpcMaxVarintLengthSize;
#else
  return str.size();
#endif
}

// Breaks a string Tensor (represented as a TensorProto) as a vector of Event
// protos.
Status WrapStringTensorAsEvents(const DebugNodeKey& debug_node_key,
                                const uint64 wall_time_us,
                                const size_t chunk_size_limit,
                                TensorProto* tensor_proto,
                                std::vector<Event>* events) {
  const protobuf::RepeatedPtrField<string>& strs = tensor_proto->string_val();
  const size_t num_strs = strs.size();
  const size_t chunk_size_ub = chunk_size_limit > 0
                                   ? chunk_size_limit
                                   : std::numeric_limits<size_t>::max();

  // E.g., if cutoffs is {j, k, l}, the chunks will have index ranges:
  //   [0:a), [a:b), [c:<end>].
  std::vector<size_t> cutoffs;
  size_t chunk_size = 0;
  for (size_t i = 0; i < num_strs; ++i) {
    // Take into account the extra bytes in proto buffer.
    if (StringValMaxBytesInProto(strs[i]) > chunk_size_ub) {
      return errors::FailedPrecondition(
          "string value at index ", i, " from debug node ",
          debug_node_key.debug_node_name,
          " does not fit gRPC message size limit (", chunk_size_ub, ")");
    }
    if (chunk_size + StringValMaxBytesInProto(strs[i]) > chunk_size_ub) {
      cutoffs.push_back(i);
      chunk_size = 0;
    }
    chunk_size += StringValMaxBytesInProto(strs[i]);
  }
  cutoffs.push_back(num_strs);
  const size_t num_chunks = cutoffs.size();

  for (size_t i = 0; i < num_chunks; ++i) {
    Event event = PrepareChunkEventProto(debug_node_key, wall_time_us,
                                         num_chunks, i, tensor_proto->dtype(),
                                         tensor_proto->tensor_shape());
    Summary::Value* value = event.mutable_summary()->mutable_value(0);

    if (cutoffs.size() == 1) {
      value->mutable_tensor()->mutable_string_val()->Swap(
          tensor_proto->mutable_string_val());
    } else {
      const size_t begin = (i == 0) ? 0 : cutoffs[i - 1];
      const size_t end = cutoffs[i];
      for (size_t j = begin; j < end; ++j) {
        value->mutable_tensor()->add_string_val(strs[j]);
      }
    }

    events->push_back(std::move(event));
  }

  return Status::OK();
}

// Encapsulates the tensor value inside a vector of Event protos. Large tensors
// are broken up to multiple protos to fit the chunk_size_limit. In each Event
// proto the field summary.tensor carries the content of the tensor.
// If chunk_size_limit <= 0, the tensor will not be broken into chunks, i.e., a
// length-1 vector will be returned, regardless of the size of the tensor.
Status WrapTensorAsEvents(const DebugNodeKey& debug_node_key,
                          const Tensor& tensor, const uint64 wall_time_us,
                          const size_t chunk_size_limit,
                          std::vector<Event>* events) {
  TensorProto tensor_proto;
  if (tensor.dtype() == DT_STRING) {
    // Treat DT_STRING specially, so that tensor_util.MakeNdarray in Python can
    // convert the TensorProto to string-type numpy array. MakeNdarray does not
    // work with strings encoded by AsProtoTensorContent() in tensor_content.
    tensor.AsProtoField(&tensor_proto);

    TF_RETURN_IF_ERROR(WrapStringTensorAsEvents(
        debug_node_key, wall_time_us, chunk_size_limit, &tensor_proto, events));
  } else {
    tensor.AsProtoTensorContent(&tensor_proto);

    const size_t total_length = tensor_proto.tensor_content().size();
    const size_t chunk_size_ub =
        chunk_size_limit > 0 ? chunk_size_limit : total_length;
    const size_t num_chunks =
        (total_length == 0)
            ? 1
            : (total_length + chunk_size_ub - 1) / chunk_size_ub;
    for (size_t i = 0; i < num_chunks; ++i) {
      const size_t pos = i * chunk_size_ub;
      const size_t len =
          (i == num_chunks - 1) ? (total_length - pos) : chunk_size_ub;
      Event event = PrepareChunkEventProto(debug_node_key, wall_time_us,
                                           num_chunks, i, tensor_proto.dtype(),
                                           tensor_proto.tensor_shape());
      event.mutable_summary()
          ->mutable_value(0)
          ->mutable_tensor()
          ->set_tensor_content(tensor_proto.tensor_content().substr(pos, len));
      events->push_back(std::move(event));
    }
  }

  return Status::OK();
}

// Appends an underscore and a timestamp to a file path. If the path already
// exists on the file system, append a hyphen and a 1-up index. Consecutive
// values of the index will be tried until the first unused one is found.
// TOCTOU race condition is not of concern here due to the fact that tfdbg
// sets parallel_iterations attribute of all while_loops to 1 to prevent
// the same node from between executed multiple times concurrently.
string AppendTimestampToFilePath(const string& in, const uint64 timestamp) {
  string out = strings::StrCat(in, "_", timestamp);

  uint64 i = 1;
  while (Env::Default()->FileExists(out).ok()) {
    out = strings::StrCat(in, "_", timestamp, "-", i);
    ++i;
  }
  return out;
}

#ifndef PLATFORM_WINDOWS
// Publishes encoded GraphDef through a gRPC debugger stream, in chunks,
// conforming to the gRPC message size limit.
Status PublishEncodedGraphDefInChunks(const string& encoded_graph_def,
                                      const string& device_name,
                                      const int64 wall_time,
                                      const string& debug_url) {
  const uint64 hash = ::tensorflow::Hash64(encoded_graph_def);
  const size_t total_length = encoded_graph_def.size();
  const size_t num_chunks =
      static_cast<size_t>(std::ceil(static_cast<float>(total_length) /
                                    DebugGrpcIO::kGrpcMessageSizeLimitBytes));
  for (size_t i = 0; i < num_chunks; ++i) {
    const size_t pos = i * DebugGrpcIO::kGrpcMessageSizeLimitBytes;
    const size_t len = (i == num_chunks - 1)
                           ? (total_length - pos)
                           : DebugGrpcIO::kGrpcMessageSizeLimitBytes;
    Event event;
    event.set_wall_time(static_cast<double>(wall_time));
    // Prefix the chunk with
    //   <hash64>,<device_name>,<wall_time>|<index>|<num_chunks>|.
    // TODO(cais): Use DebuggerEventMetadata to store device_name, num_chunks
    // and chunk_index, instead.
    event.set_graph_def(strings::StrCat(hash, ",", device_name, ",", wall_time,
                                        "|", i, "|", num_chunks, "|",
                                        encoded_graph_def.substr(pos, len)));
    const Status s = DebugGrpcIO::SendEventProtoThroughGrpcStream(
        event, debug_url, num_chunks - 1 == i);
    if (!s.ok()) {
      return errors::FailedPrecondition(
          "Failed to send chunk ", i, " of ", num_chunks,
          " of encoded GraphDef of size ", encoded_graph_def.size(), " bytes, ",
          "due to: ", s.error_message());
    }
  }
  return Status::OK();
}
#endif  // #ifndef PLATFORM_WINDOWS

}  // namespace

const char* const DebugIO::kDebuggerPluginName = "debugger";

const char* const DebugIO::kCoreMetadataTag = "core_metadata_";

const char* const DebugIO::kGraphTag = "graph_";

const char* const DebugIO::kHashTag = "hash";

Status ReadEventFromFile(const string& dump_file_path, Event* event) {
  Env* env(Env::Default());

  string content;
  uint64 file_size = 0;

  Status s = env->GetFileSize(dump_file_path, &file_size);
  if (!s.ok()) {
    return s;
  }

  content.resize(file_size);

  std::unique_ptr<RandomAccessFile> file;
  s = env->NewRandomAccessFile(dump_file_path, &file);
  if (!s.ok()) {
    return s;
  }

  StringPiece result;
  s = file->Read(0, file_size, &result, &(content)[0]);
  if (!s.ok()) {
    return s;
  }

  event->ParseFromString(content);
  return Status::OK();
}

const char* const DebugIO::kFileURLScheme = "file://";
const char* const DebugIO::kGrpcURLScheme = "grpc://";
const char* const DebugIO::kMemoryURLScheme = "memcbk://";

// Publishes debug metadata to a set of debug URLs.
Status DebugIO::PublishDebugMetadata(
    const int64 global_step, const int64 session_run_index,
    const int64 executor_step_index, const std::vector<string>& input_names,
    const std::vector<string>& output_names,
    const std::vector<string>& target_nodes,
    const std::unordered_set<string>& debug_urls) {
  std::ostringstream oss;

  // Construct a JSON string to carry the metadata.
  oss << "{";
  oss << "\"global_step\":" << global_step << ",";
  oss << "\"session_run_index\":" << session_run_index << ",";
  oss << "\"executor_step_index\":" << executor_step_index << ",";
  oss << "\"input_names\":[";
  for (size_t i = 0; i < input_names.size(); ++i) {
    oss << "\"" << input_names[i] << "\"";
    if (i < input_names.size() - 1) {
      oss << ",";
    }
  }
  oss << "],";
  oss << "\"output_names\":[";
  for (size_t i = 0; i < output_names.size(); ++i) {
    oss << "\"" << output_names[i] << "\"";
    if (i < output_names.size() - 1) {
      oss << ",";
    }
  }
  oss << "],";
  oss << "\"target_nodes\":[";
  for (size_t i = 0; i < target_nodes.size(); ++i) {
    oss << "\"" << target_nodes[i] << "\"";
    if (i < target_nodes.size() - 1) {
      oss << ",";
    }
  }
  oss << "]";
  oss << "}";

  const string json_metadata = oss.str();
  Event event;
  event.set_wall_time(static_cast<double>(Env::Default()->NowMicros()));
  LogMessage* log_message = event.mutable_log_message();
  log_message->set_message(json_metadata);

  Status status;
  for (const string& url : debug_urls) {
    if (str_util::Lowercase(url).find(kGrpcURLScheme) == 0) {
#ifndef PLATFORM_WINDOWS
      Event grpc_event;

      // Determine the path (if any) in the grpc:// URL, and add it as a field
      // of the JSON string.
      const string address = url.substr(strlen(DebugIO::kFileURLScheme));
      const string path = address.find("/") == string::npos
                              ? ""
                              : address.substr(address.find("/"));
      grpc_event.set_wall_time(event.wall_time());
      LogMessage* log_message_grpc = grpc_event.mutable_log_message();
      log_message_grpc->set_message(
          strings::StrCat(json_metadata.substr(0, json_metadata.size() - 1),
                          ",\"grpc_path\":\"", path, "\"}"));

      status.Update(
          DebugGrpcIO::SendEventProtoThroughGrpcStream(grpc_event, url, true));
#else
      GRPC_OSS_WINDOWS_UNIMPLEMENTED_ERROR;
#endif
    } else if (str_util::Lowercase(url).find(kFileURLScheme) == 0) {
      const string dump_root_dir = url.substr(strlen(kFileURLScheme));
      const string core_metadata_path = AppendTimestampToFilePath(
          io::JoinPath(
              dump_root_dir,
              strings::StrCat(DebugNodeKey::kMetadataFilePrefix,
                              DebugIO::kCoreMetadataTag, "sessionrun",
                              strings::Printf("%.14lld", session_run_index))),
          Env::Default()->NowMicros());
      status.Update(DebugFileIO::DumpEventProtoToFile(
          event, string(io::Dirname(core_metadata_path)),
          string(io::Basename(core_metadata_path))));
    }
  }

  return status;
}

Status DebugIO::PublishDebugTensor(const DebugNodeKey& debug_node_key,
                                   const Tensor& tensor,
                                   const uint64 wall_time_us,
                                   const gtl::ArraySlice<string>& debug_urls,
                                   const bool gated_grpc) {
  int32 num_failed_urls = 0;
  std::vector<Status> fail_statuses;
  for (const string& url : debug_urls) {
    if (str_util::Lowercase(url).find(kFileURLScheme) == 0) {
      const string dump_root_dir = url.substr(strlen(kFileURLScheme));

      const int64 tensorBytes =
          tensor.IsInitialized() ? tensor.TotalBytes() : 0;
      if (!DebugFileIO::requestDiskByteUsage(tensorBytes)) {
        return errors::ResourceExhausted(
            "TensorFlow Debugger has exhausted file-system byte-size "
            "allowance (",
            DebugFileIO::globalDiskBytesLimit, "), therefore it cannot ",
            "dump an additional ", tensorBytes, " byte(s) of tensor data ",
            "for the debug tensor ", debug_node_key.node_name, ":",
            debug_node_key.output_slot, ". You may use the environment ",
            "variable TFDBG_DISK_BYTES_LIMIT to set a higher limit.");
      }

      Status s = DebugFileIO::DumpTensorToDir(
          debug_node_key, tensor, wall_time_us, dump_root_dir, nullptr);
      if (!s.ok()) {
        num_failed_urls++;
        fail_statuses.push_back(s);
      }
    } else if (str_util::Lowercase(url).find(kGrpcURLScheme) == 0) {
#ifndef PLATFORM_WINDOWS
      Status s = DebugGrpcIO::SendTensorThroughGrpcStream(
          debug_node_key, tensor, wall_time_us, url, gated_grpc);

      if (!s.ok()) {
        num_failed_urls++;
        fail_statuses.push_back(s);
      }
#else
      GRPC_OSS_WINDOWS_UNIMPLEMENTED_ERROR;
#endif
    } else if (str_util::Lowercase(url).find(kMemoryURLScheme) == 0) {
      const string dump_root_dir = url.substr(strlen(kMemoryURLScheme));
      auto* callback_registry = DebugCallbackRegistry::singleton();
      auto* callback = callback_registry->GetCallback(dump_root_dir);
      CHECK(callback) << "No callback registered for: " << dump_root_dir;
      (*callback)(debug_node_key, tensor);
    } else {
      return Status(error::UNAVAILABLE,
                    strings::StrCat("Invalid debug target URL: ", url));
    }
  }

  if (num_failed_urls == 0) {
    return Status::OK();
  } else {
    string error_message = strings::StrCat(
        "Publishing to ", num_failed_urls, " of ", debug_urls.size(),
        " debug target URLs failed, due to the following errors:");
    for (Status& status : fail_statuses) {
      error_message =
          strings::StrCat(error_message, " ", status.error_message(), ";");
    }

    return Status(error::INTERNAL, error_message);
  }
}

Status DebugIO::PublishDebugTensor(const DebugNodeKey& debug_node_key,
                                   const Tensor& tensor,
                                   const uint64 wall_time_us,
                                   const gtl::ArraySlice<string>& debug_urls) {
  return PublishDebugTensor(debug_node_key, tensor, wall_time_us, debug_urls,
                            false);
}

Status DebugIO::PublishGraph(const Graph& graph, const string& device_name,
                             const std::unordered_set<string>& debug_urls) {
  GraphDef graph_def;
  graph.ToGraphDef(&graph_def);

  string buf;
  graph_def.SerializeToString(&buf);

  const int64 now_micros = Env::Default()->NowMicros();
  Event event;
  event.set_wall_time(static_cast<double>(now_micros));
  event.set_graph_def(buf);

  Status status = Status::OK();
  for (const string& debug_url : debug_urls) {
    if (debug_url.find(kFileURLScheme) == 0) {
      const string dump_root_dir =
          io::JoinPath(debug_url.substr(strlen(kFileURLScheme)),
                       DebugNodeKey::DeviceNameToDevicePath(device_name));
      const uint64 graph_hash = ::tensorflow::Hash64(buf);
      const string file_name =
          strings::StrCat(DebugNodeKey::kMetadataFilePrefix, DebugIO::kGraphTag,
                          DebugIO::kHashTag, graph_hash, "_", now_micros);

      status.Update(
          DebugFileIO::DumpEventProtoToFile(event, dump_root_dir, file_name));
    } else if (debug_url.find(kGrpcURLScheme) == 0) {
#ifndef PLATFORM_WINDOWS
      status.Update(PublishEncodedGraphDefInChunks(buf, device_name, now_micros,
                                                   debug_url));
#else
      GRPC_OSS_WINDOWS_UNIMPLEMENTED_ERROR;
#endif
    }
  }

  return status;
}

bool DebugIO::IsCopyNodeGateOpen(
    const std::vector<DebugWatchAndURLSpec>& specs) {
#ifndef PLATFORM_WINDOWS
  for (const DebugWatchAndURLSpec& spec : specs) {
    if (!spec.gated_grpc || spec.url.compare(0, strlen(DebugIO::kGrpcURLScheme),
                                             DebugIO::kGrpcURLScheme)) {
      return true;
    } else {
      if (DebugGrpcIO::IsReadGateOpen(spec.url, spec.watch_key)) {
        return true;
      }
    }
  }
  return false;
#else
  return true;
#endif
}

bool DebugIO::IsDebugNodeGateOpen(const string& watch_key,
                                  const std::vector<string>& debug_urls) {
#ifndef PLATFORM_WINDOWS
  for (const string& debug_url : debug_urls) {
    if (debug_url.compare(0, strlen(DebugIO::kGrpcURLScheme),
                          DebugIO::kGrpcURLScheme)) {
      return true;
    } else {
      if (DebugGrpcIO::IsReadGateOpen(debug_url, watch_key)) {
        return true;
      }
    }
  }
  return false;
#else
  return true;
#endif
}

bool DebugIO::IsDebugURLGateOpen(const string& watch_key,
                                 const string& debug_url) {
#ifndef PLATFORM_WINDOWS
  if (debug_url.find(kGrpcURLScheme) != 0) {
    return true;
  } else {
    return DebugGrpcIO::IsReadGateOpen(debug_url, watch_key);
  }
#else
  return true;
#endif
}

Status DebugIO::CloseDebugURL(const string& debug_url) {
  if (debug_url.find(DebugIO::kGrpcURLScheme) == 0) {
#ifndef PLATFORM_WINDOWS
    return DebugGrpcIO::CloseGrpcStream(debug_url);
#else
    GRPC_OSS_WINDOWS_UNIMPLEMENTED_ERROR;
#endif
  } else {
    // No-op for non-gRPC URLs.
    return Status::OK();
  }
}

Status DebugFileIO::DumpTensorToDir(const DebugNodeKey& debug_node_key,
                                    const Tensor& tensor,
                                    const uint64 wall_time_us,
                                    const string& dump_root_dir,
                                    string* dump_file_path) {
  const string file_path =
      GetDumpFilePath(dump_root_dir, debug_node_key, wall_time_us);

  if (dump_file_path != nullptr) {
    *dump_file_path = file_path;
  }

  return DumpTensorToEventFile(debug_node_key, tensor, wall_time_us, file_path);
}

string DebugFileIO::GetDumpFilePath(const string& dump_root_dir,
                                    const DebugNodeKey& debug_node_key,
                                    const uint64 wall_time_us) {
  return AppendTimestampToFilePath(
      io::JoinPath(dump_root_dir, debug_node_key.device_path,
                   strings::StrCat(debug_node_key.node_name, "_",
                                   debug_node_key.output_slot, "_",
                                   debug_node_key.debug_op)),
      wall_time_us);
}

Status DebugFileIO::DumpEventProtoToFile(const Event& event_proto,
                                         const string& dir_name,
                                         const string& file_name) {
  Env* env(Env::Default());

  Status s = RecursiveCreateDir(env, dir_name);
  if (!s.ok()) {
    return Status(error::FAILED_PRECONDITION,
                  strings::StrCat("Failed to create directory  ", dir_name,
                                  ", due to: ", s.error_message()));
  }

  const string file_path = io::JoinPath(dir_name, file_name);

  string event_str;
  event_proto.SerializeToString(&event_str);

  std::unique_ptr<WritableFile> f = nullptr;
  TF_CHECK_OK(env->NewWritableFile(file_path, &f));
  f->Append(event_str).IgnoreError();
  TF_CHECK_OK(f->Close());

  return Status::OK();
}

Status DebugFileIO::DumpTensorToEventFile(const DebugNodeKey& debug_node_key,
                                          const Tensor& tensor,
                                          const uint64 wall_time_us,
                                          const string& file_path) {
  std::vector<Event> events;
  TF_RETURN_IF_ERROR(
      WrapTensorAsEvents(debug_node_key, tensor, wall_time_us, 0, &events));
  return DumpEventProtoToFile(events[0], string(io::Dirname(file_path)),
                              string(io::Basename(file_path)));
}

Status DebugFileIO::RecursiveCreateDir(Env* env, const string& dir) {
  if (env->FileExists(dir).ok() && env->IsDirectory(dir).ok()) {
    // The path already exists as a directory. Return OK right away.
    return Status::OK();
  }

  string parent_dir(io::Dirname(dir));
  if (!env->FileExists(parent_dir).ok()) {
    // The parent path does not exist yet, create it first.
    Status s = RecursiveCreateDir(env, parent_dir);  // Recursive call
    if (!s.ok()) {
      return Status(
          error::FAILED_PRECONDITION,
          strings::StrCat("Failed to create directory  ", parent_dir));
    }
  } else if (env->FileExists(parent_dir).ok() &&
             !env->IsDirectory(parent_dir).ok()) {
    // The path exists, but it is a file.
    return Status(error::FAILED_PRECONDITION,
                  strings::StrCat("Failed to create directory  ", parent_dir,
                                  " because the path exists as a file "));
  }

  env->CreateDir(dir).IgnoreError();
  // Guard against potential race in creating directories by doing a check
  // after the CreateDir call.
  if (env->FileExists(dir).ok() && env->IsDirectory(dir).ok()) {
    return Status::OK();
  } else {
    return Status(error::ABORTED,
                  strings::StrCat("Failed to create directory  ", parent_dir));
  }
}

// Default total disk usage limit: 100 GBytes
const uint64 DebugFileIO::defaultGlobalDiskBytesLimit = 107374182400L;
uint64 DebugFileIO::globalDiskBytesLimit = 0;
uint64 DebugFileIO::diskBytesUsed = 0;

mutex DebugFileIO::bytes_mu(LINKER_INITIALIZED);

bool DebugFileIO::requestDiskByteUsage(uint64 bytes) {
  mutex_lock l(bytes_mu);
  if (globalDiskBytesLimit == 0) {
    const char* env_tfdbg_disk_bytes_limit = getenv("TFDBG_DISK_BYTES_LIMIT");
    if (env_tfdbg_disk_bytes_limit == nullptr ||
        strlen(env_tfdbg_disk_bytes_limit) == 0) {
      globalDiskBytesLimit = defaultGlobalDiskBytesLimit;
    } else {
      strings::safe_strtou64(string(env_tfdbg_disk_bytes_limit),
                             &globalDiskBytesLimit);
    }
  }

  if (bytes == 0) {
    return true;
  }
  if (diskBytesUsed + bytes < globalDiskBytesLimit) {
    diskBytesUsed += bytes;
    return true;
  } else {
    return false;
  }
}

void DebugFileIO::resetDiskByteUsage() {
  mutex_lock l(bytes_mu);
  diskBytesUsed = 0;
}

#ifndef PLATFORM_WINDOWS
DebugGrpcChannel::DebugGrpcChannel(const string& server_stream_addr)
    : server_stream_addr_(server_stream_addr),
      url_(strings::StrCat(DebugIO::kGrpcURLScheme, server_stream_addr)) {}

Status DebugGrpcChannel::Connect(const int64 timeout_micros) {
  ::grpc::ChannelArguments args;
  args.SetInt(GRPC_ARG_MAX_MESSAGE_LENGTH, std::numeric_limits<int32>::max());
  // Avoid problems where default reconnect backoff is too long (e.g., 20 s).
  args.SetInt(GRPC_ARG_MAX_RECONNECT_BACKOFF_MS, 1000);
  channel_ = ::grpc::CreateCustomChannel(
      server_stream_addr_, ::grpc::InsecureChannelCredentials(), args);
  if (!channel_->WaitForConnected(
          gpr_time_add(gpr_now(GPR_CLOCK_REALTIME),
                       gpr_time_from_micros(timeout_micros, GPR_TIMESPAN)))) {
    return errors::FailedPrecondition(
        "Failed to connect to gRPC channel at ", server_stream_addr_,
        " within a timeout of ", timeout_micros / 1e6, " s.");
  }
  stub_ = EventListener::NewStub(channel_);
  reader_writer_ = stub_->SendEvents(&ctx_);

  return Status::OK();
}

bool DebugGrpcChannel::WriteEvent(const Event& event) {
  mutex_lock l(mu_);
  return reader_writer_->Write(event);
}

bool DebugGrpcChannel::ReadEventReply(EventReply* event_reply) {
  mutex_lock l(mu_);
  return reader_writer_->Read(event_reply);
}

void DebugGrpcChannel::ReceiveAndProcessEventReplies(const size_t max_replies) {
  EventReply event_reply;
  size_t num_replies = 0;
  while ((max_replies == 0 || ++num_replies <= max_replies) &&
         ReadEventReply(&event_reply)) {
    for (const EventReply::DebugOpStateChange& debug_op_state_change :
         event_reply.debug_op_state_changes()) {
      string watch_key = strings::StrCat(debug_op_state_change.node_name(), ":",
                                         debug_op_state_change.output_slot(),
                                         ":", debug_op_state_change.debug_op());
      DebugGrpcIO::SetDebugNodeKeyGrpcState(url_, watch_key,
                                            debug_op_state_change.state());
    }
  }
}

Status DebugGrpcChannel::ReceiveServerRepliesAndClose() {
  reader_writer_->WritesDone();
  // Read all EventReply messages (if any) from the server.
  ReceiveAndProcessEventReplies(0);

  if (reader_writer_->Finish().ok()) {
    return Status::OK();
  } else {
    return Status(error::FAILED_PRECONDITION,
                  "Failed to close debug GRPC stream.");
  }
}

mutex DebugGrpcIO::streams_mu(LINKER_INITIALIZED);

int64 DebugGrpcIO::channel_connection_timeout_micros = 900 * 1000 * 1000;
// TODO(cais): Make this configurable?

const size_t DebugGrpcIO::kGrpcMessageSizeLimitBytes = 4000 * 1024;

const size_t DebugGrpcIO::kGrpcMaxVarintLengthSize = 6;

std::unordered_map<string, std::unique_ptr<DebugGrpcChannel>>*
DebugGrpcIO::GetStreamChannels() {
  static std::unordered_map<string, std::unique_ptr<DebugGrpcChannel>>*
      stream_channels =
          new std::unordered_map<string, std::unique_ptr<DebugGrpcChannel>>();
  return stream_channels;
}

Status DebugGrpcIO::SendTensorThroughGrpcStream(
    const DebugNodeKey& debug_node_key, const Tensor& tensor,
    const uint64 wall_time_us, const string& grpc_stream_url,
    const bool gated) {
  if (gated &&
      !IsReadGateOpen(grpc_stream_url, debug_node_key.debug_node_name)) {
    return Status::OK();
  } else {
    std::vector<Event> events;
    TF_RETURN_IF_ERROR(WrapTensorAsEvents(debug_node_key, tensor, wall_time_us,
                                          kGrpcMessageSizeLimitBytes, &events));
    for (const Event& event : events) {
      TF_RETURN_IF_ERROR(
          SendEventProtoThroughGrpcStream(event, grpc_stream_url));
    }
    if (IsWriteGateOpen(grpc_stream_url, debug_node_key.debug_node_name)) {
      DebugGrpcChannel* debug_grpc_channel = nullptr;
      TF_RETURN_IF_ERROR(
          GetOrCreateDebugGrpcChannel(grpc_stream_url, &debug_grpc_channel));
      debug_grpc_channel->ReceiveAndProcessEventReplies(1);
      // TODO(cais): Support new tensor value carried in the EventReply for
      // overriding the value of the tensor being published.
    }
    return Status::OK();
  }
}

Status DebugGrpcIO::ReceiveEventReplyProtoThroughGrpcStream(
    EventReply* event_reply, const string& grpc_stream_url) {
  DebugGrpcChannel* debug_grpc_channel = nullptr;
  TF_RETURN_IF_ERROR(
      GetOrCreateDebugGrpcChannel(grpc_stream_url, &debug_grpc_channel));
  if (debug_grpc_channel->ReadEventReply(event_reply)) {
    return Status::OK();
  } else {
    return errors::Cancelled(strings::StrCat(
        "Reading EventReply from stream URL ", grpc_stream_url, " failed."));
  }
}

Status DebugGrpcIO::GetOrCreateDebugGrpcChannel(
    const string& grpc_stream_url, DebugGrpcChannel** debug_grpc_channel) {
  const string addr_with_path =
      grpc_stream_url.find(DebugIO::kGrpcURLScheme) == 0
          ? grpc_stream_url.substr(strlen(DebugIO::kGrpcURLScheme))
          : grpc_stream_url;
  const string server_stream_addr =
      addr_with_path.substr(0, addr_with_path.find('/'));
  {
    mutex_lock l(streams_mu);
    std::unordered_map<string, std::unique_ptr<DebugGrpcChannel>>*
        stream_channels = GetStreamChannels();
    if (stream_channels->find(grpc_stream_url) == stream_channels->end()) {
      std::unique_ptr<DebugGrpcChannel> channel(
          new DebugGrpcChannel(server_stream_addr));
      TF_RETURN_IF_ERROR(channel->Connect(channel_connection_timeout_micros));
      stream_channels->insert(
          std::make_pair(grpc_stream_url, std::move(channel)));
    }
    *debug_grpc_channel = (*stream_channels)[grpc_stream_url].get();
  }
  return Status::OK();
}

Status DebugGrpcIO::SendEventProtoThroughGrpcStream(
    const Event& event_proto, const string& grpc_stream_url,
    const bool receive_reply) {
  DebugGrpcChannel* debug_grpc_channel;
  TF_RETURN_IF_ERROR(
      GetOrCreateDebugGrpcChannel(grpc_stream_url, &debug_grpc_channel));

  bool write_ok = debug_grpc_channel->WriteEvent(event_proto);
  if (!write_ok) {
    return errors::Cancelled(strings::StrCat("Write event to stream URL ",
                                             grpc_stream_url, " failed."));
  }

  if (receive_reply) {
    debug_grpc_channel->ReceiveAndProcessEventReplies(1);
  }

  return Status::OK();
}

bool DebugGrpcIO::IsReadGateOpen(const string& grpc_debug_url,
                                 const string& watch_key) {
  const DebugNodeName2State* enabled_node_to_state =
      GetEnabledDebugOpStatesAtUrl(grpc_debug_url);
  return enabled_node_to_state->find(watch_key) != enabled_node_to_state->end();
}

bool DebugGrpcIO::IsWriteGateOpen(const string& grpc_debug_url,
                                  const string& watch_key) {
  const DebugNodeName2State* enabled_node_to_state =
      GetEnabledDebugOpStatesAtUrl(grpc_debug_url);
  auto it = enabled_node_to_state->find(watch_key);
  if (it == enabled_node_to_state->end()) {
    return false;
  } else {
    return it->second == EventReply::DebugOpStateChange::READ_WRITE;
  }
}

Status DebugGrpcIO::CloseGrpcStream(const string& grpc_stream_url) {
  mutex_lock l(streams_mu);

  std::unordered_map<string, std::unique_ptr<DebugGrpcChannel>>*
      stream_channels = GetStreamChannels();
  if (stream_channels->find(grpc_stream_url) != stream_channels->end()) {
    // Stream of the specified address exists. Close it and remove it from
    // record.
    Status s =
        (*stream_channels)[grpc_stream_url]->ReceiveServerRepliesAndClose();
    (*stream_channels).erase(grpc_stream_url);
    return s;
  } else {
    // Stream of the specified address does not exist. No action.
    return Status::OK();
  }
}

std::unordered_map<string, DebugGrpcIO::DebugNodeName2State>*
DebugGrpcIO::GetEnabledDebugOpStates() {
  static std::unordered_map<string, DebugNodeName2State>*
      enabled_debug_op_states =
          new std::unordered_map<string, DebugNodeName2State>();
  return enabled_debug_op_states;
}

DebugGrpcIO::DebugNodeName2State* DebugGrpcIO::GetEnabledDebugOpStatesAtUrl(
    const string& grpc_debug_url) {
  static mutex* debug_ops_state_mu = new mutex();
  std::unordered_map<string, DebugNodeName2State>* states =
      GetEnabledDebugOpStates();

  mutex_lock l(*debug_ops_state_mu);
  if (states->find(grpc_debug_url) == states->end()) {
    DebugNodeName2State url_enabled_debug_op_states;
    (*states)[grpc_debug_url] = url_enabled_debug_op_states;
  }
  return &(*states)[grpc_debug_url];
}

void DebugGrpcIO::SetDebugNodeKeyGrpcState(
    const string& grpc_debug_url, const string& watch_key,
    const EventReply::DebugOpStateChange::State new_state) {
  DebugNodeName2State* states = GetEnabledDebugOpStatesAtUrl(grpc_debug_url);
  if (new_state == EventReply::DebugOpStateChange::DISABLED) {
    if (states->find(watch_key) == states->end()) {
      LOG(ERROR) << "Attempt to disable a watch key that is not currently "
                 << "enabled at " << grpc_debug_url << ": " << watch_key;
    } else {
      states->erase(watch_key);
    }
  } else if (new_state != EventReply::DebugOpStateChange::STATE_UNSPECIFIED) {
    (*states)[watch_key] = new_state;
  }
}

void DebugGrpcIO::ClearEnabledWatchKeys() {
  GetEnabledDebugOpStates()->clear();
}

#endif  // #ifndef PLATFORM_WINDOWS

}  // namespace tensorflow