Avoid direct access to the env var TEST_UNDECLARED_OUTPUTS_DIR.

On Windows, Bazel populates this path with `/`s only making proper path
management impossible without sanitizing the path up front. This changes to
accessing the env var through an indirection layer which will fix path problems
on Windows when the codebase is ready to switch over.

PiperOrigin-RevId: 296083765
Change-Id: I26bbaf83ba5e3fafd3ab0a0de08f6cb597b94477
This commit is contained in:
Brian Atkinson 2020-02-19 16:25:30 -08:00 committed by TensorFlower Gardener
parent edaaeaddbd
commit ba2cbe1e55
3 changed files with 18 additions and 19 deletions

View File

@ -27,6 +27,7 @@ limitations under the License.
#include "mlir/IR/Operation.h" // TF:llvm-project
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/path.h"
namespace tensorflow {
@ -97,18 +98,18 @@ struct WritableFileRawStream : public llvm::raw_ostream {
Status CreateFileForDumping(llvm::StringRef name,
std::unique_ptr<llvm::raw_ostream>* os,
std::string* filepath, llvm::StringRef dirname) {
const char* dir = nullptr;
std::string dir;
if (!dirname.empty())
dir = dirname.data();
dir = std::string(dirname);
else
dir = GetDumpDirFromEnvVar();
if (!dir) {
if (dir.empty()) {
return Status(error::Code::INVALID_ARGUMENT,
"(TF_DUMP_GRAPH_PREFIX not specified)");
}
if (std::strncmp(dir, "-", 2) == 0) {
if (dir == "-") {
*os = std::make_unique<LogInfoRawStream>();
*filepath = "LOG(INFO)";
return Status();
@ -151,25 +152,24 @@ std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
return filepath;
}
const char* GetDumpDirFromEnvVar() {
std::string GetDumpDirFromEnvVar() {
const char* prefix_env = getenv("TF_DUMP_GRAPH_PREFIX");
if (!prefix_env) {
LOG(WARNING)
<< "Failed to dump MLIR module because dump location is not "
<< " specified through TF_DUMP_GRAPH_PREFIX environment variable.";
return nullptr;
return "";
}
if (absl::EqualsIgnoreCase(prefix_env, "sponge")) {
const char* tmp_dir = getenv("TEST_UNDECLARED_OUTPUTS_DIR");
if (!tmp_dir) {
LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
"TEST_UNDECLARED_OUTPUT_DIRS is not set";
return nullptr;
}
return tmp_dir;
std::string result = prefix_env;
if (absl::EqualsIgnoreCase(result, "sponge") &&
!io::GetTestUndeclaredOutputsDir(&result)) {
LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
"TEST_UNDECLARED_OUTPUT_DIRS is not set";
return "";
}
return prefix_env;
return result;
}
} // namespace tensorflow

View File

@ -54,7 +54,7 @@ std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
// Default is reading from TF_DUMP_GRAPH_PREFIX, and if the string is 'sponge'
// read from TEST_UNDECLARED_OUTPUTS_DIR. Returns nullptr if the directory
// cannot be determined and generates a warning message.
const char* GetDumpDirFromEnvVar();
std::string GetDumpDirFromEnvVar();
} // namespace tensorflow

View File

@ -35,11 +35,10 @@ namespace tensorflow {
// This require the TF_DUMP_GRAPH_PREFIX to be set to a path that exist (or can
// be created).
static void DumpModule(mlir::ModuleOp module, llvm::StringRef file_prefix) {
const char* prefix_env = GetDumpDirFromEnvVar();
if (!prefix_env) {
std::string prefix = GetDumpDirFromEnvVar();
if (prefix.empty()) {
return;
}
std::string prefix = prefix_env;
auto* env = tensorflow::Env::Default();
auto status = env->RecursivelyCreateDir(prefix);