diff options
author | Xiaoqiang Zheng <zhengxq@google.com> | 2016-10-12 15:36:13 -0800 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2016-10-12 16:47:55 -0700 |
commit | 3ac64952aaf944cc6cc498474d3008cdb66bc8c7 (patch) | |
tree | 4496f6899bea3eb6b4234f9d81d5ef4fb5796042 /tensorflow/core/util/env_var.cc | |
parent | 3e581482a97683c90a65969423dd75c41b2f8653 (diff) |
Adding an environmental variable to control the cuda host memory limit.
Change: 135978018
Diffstat (limited to 'tensorflow/core/util/env_var.cc')
-rw-r--r-- | tensorflow/core/util/env_var.cc | 63 |
1 files changed, 63 insertions, 0 deletions
diff --git a/tensorflow/core/util/env_var.cc b/tensorflow/core/util/env_var.cc new file mode 100644 index 0000000000..d4e89b966e --- /dev/null +++ b/tensorflow/core/util/env_var.cc @@ -0,0 +1,63 @@ +/* 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/util/env_var.h" + +#include <stdlib.h> + +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/numbers.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { + +Status ReadBoolFromEnvVar(StringPiece env_var_name, bool default_val, + bool* value) { + *value = default_val; + const char* tf_env_var_val = getenv(env_var_name.ToString().c_str()); + if (tf_env_var_val == nullptr) { + return Status::OK(); + } + string str_value = str_util::Lowercase(tf_env_var_val); + if (str_value == "0" || str_value == "false") { + *value = false; + return Status::OK(); + } else if (str_value == "1" || str_value == "true") { + *value = true; + return Status::OK(); + } + return errors::InvalidArgument(strings::StrCat( + "Failed to parse the env-var ${", env_var_name, "} into bool: ", + tf_env_var_val, ". Use the default value: ", default_val)); +} + +Status ReadInt64FromEnvVar(StringPiece env_var_name, int64 default_val, + int64* value) { + *value = default_val; + const char* tf_env_var_val = getenv(env_var_name.ToString().c_str()); + if (tf_env_var_val == nullptr) { + return Status::OK(); + } + if (strings::safe_strto64(tf_env_var_val, value)) { + return Status::OK(); + } + return errors::InvalidArgument(strings::StrCat( + "Failed to parse the env-var ${", env_var_name, "} into int64: ", + tf_env_var_val, ". Use the default value: ", default_val)); +} + +} // namespace tensorflow |