diff --git a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc index 11dec3947a..3f4718b0ac 100644 --- a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc +++ b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.cc @@ -90,14 +90,14 @@ void DumpTensorToStdOut(const Tensor& tensor) { std::cout << std::endl; } -PathString MakeTensorFileName(const std::string& tensor_name) { +PathString MakeTensorFileName(const std::string& tensor_name, const NodeDumpOptions& dump_options) { auto make_valid_name = [](std::string name) { std::replace_if( name.begin(), name.end(), [](char c) { return !std::isalnum(c); }, '_'); return name; }; - return path_utils::MakePathString(make_valid_name(tensor_name), ".tensorproto"); + return path_utils::MakePathString(make_valid_name(tensor_name), dump_options.file_suffix, ".tensorproto"); } void DumpTensorToFile(const Tensor& tensor, const std::string& tensor_name, const Path& file_path) { @@ -125,7 +125,7 @@ void DumpCpuTensor( break; } case NodeDumpOptions::DataDestination::TensorProtoFiles: { - const Path tensor_file = dump_options.output_dir / Path::Parse(MakeTensorFileName(tensor_name)); + const Path tensor_file = dump_options.output_dir / Path::Parse(MakeTensorFileName(tensor_name, dump_options)); DumpTensorToFile(tensor, tensor_name, tensor_file); break; } @@ -203,6 +203,15 @@ const NodeDumpOptions& NodeDumpOptionsFromEnvironmentVariables() { opts.data_destination = NodeDumpOptions::DataDestination::TensorProtoFiles; } + if (get_bool_env_var(env_vars::kAppendRankToFileName)) { + std::string rank = Env::Default().GetEnvironmentVar("OMPI_COMM_WORLD_RANK"); + if (rank.empty()) { + opts.file_suffix = "_default_rank_0"; + } else { + opts.file_suffix = "_rank_" + rank; + } + } + opts.output_dir = Path::Parse(ToPathString(Env::Default().GetEnvironmentVar(env_vars::kOutputDir))); // check for confirmation for dumping data to files for all nodes diff --git a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.h b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.h index a0649fc66c..bed1a9cb3c 100644 --- a/onnxruntime/core/framework/debug_node_inputs_outputs_utils.h +++ b/onnxruntime/core/framework/debug_node_inputs_outputs_utils.h @@ -30,6 +30,8 @@ constexpr const char* kNameFilter = "ORT_DEBUG_NODE_IO_NAME_FILTER"; constexpr const char* kOpTypeFilter = "ORT_DEBUG_NODE_IO_OP_TYPE_FILTER"; // set to non-zero to dump data to files instead of stdout constexpr const char* kDumpDataToFiles = "ORT_DEBUG_NODE_IO_DUMP_DATA_TO_FILES"; +// set to non-zero to append OpenMPI world rank to filename +constexpr const char* kAppendRankToFileName = "ORT_DEBUG_NODE_IO_APPEND_RANK_TO_FILE_NAME"; // specify the output directory for any data files produced constexpr const char* kOutputDir = "ORT_DEBUG_NODE_IO_OUTPUT_DIR"; // set to non-zero to confirm that dumping data files for all nodes is acceptable @@ -75,6 +77,7 @@ struct NodeDumpOptions { TensorProtoFiles, } data_destination{DataDestination::StdOut}; + std::string file_suffix; // the output directory for dumped data files Path output_dir; };