1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/tpu/tpu_initializer_helper.h"
17
18 #if defined(LIBTPU_ON_GCE)
19 #include <fcntl.h>
20 #include <stdlib.h>
21 #include <unistd.h>
22 #endif // LIBTPU_ON_GCE
23
24 #include "absl/strings/str_split.h"
25 #include "absl/synchronization/mutex.h"
26 #include "tensorflow/core/platform/logging.h"
27
28 namespace tensorflow {
29 namespace tpu {
30
TryAcquireTpuLock()31 bool TryAcquireTpuLock() {
32 #if defined(LIBTPU_ON_GCE)
33 static absl::Mutex* mu = new absl::Mutex();
34 absl::MutexLock l(mu);
35
36 static bool attempted_file_open = false;
37 static bool should_load_library = false;
38
39 if (!attempted_file_open) {
40 should_load_library = true;
41
42 // if the TPU_HOST_BOUNDS env var is set, that means we are loading each
43 // chip in a different process and thus multiple libtpu loads are OK.
44 if (getenv("TPU_HOST_BOUNDS") == nullptr) {
45 int fd = open("/tmp/libtpu_lockfile", O_CREAT | O_RDWR, 0644);
46
47 // This lock is held until the process exits intentionally. The underlying
48 // TPU device will be held on until it quits.
49 if (lockf(fd, F_TLOCK, 0) != 0) {
50 LOG(WARNING) << "libtpu.so already in used by another process. Not "
51 "attempting to load libtpu.so in this process.";
52 should_load_library = false;
53 } else {
54 should_load_library = true;
55 }
56 } else {
57 LOG(INFO) << "TPU_HOST_BOUNDS is set, allowing multiple libtpu.so loads.";
58 should_load_library = true;
59 }
60 }
61
62 return should_load_library;
63 #else // LIBTPU_ON_GCE
64 return false;
65 #endif
66 }
67
68 std::pair<std::vector<std::string>, std::vector<const char*>>
GetLibTpuInitArguments()69 GetLibTpuInitArguments() {
70 // We make copies of the arguments returned by getenv because the memory
71 // returned may be altered or invalidated by further calls to getenv.
72 std::vector<std::string> args;
73 std::vector<const char*> arg_ptrs;
74
75 // Retrieve arguments from environment if applicable.
76 char* env = getenv("LIBTPU_INIT_ARGS");
77 if (env != nullptr) {
78 // TODO(frankchn): Handles quotes properly if necessary.
79 args = absl::StrSplit(env, ' ');
80 }
81
82 arg_ptrs.reserve(args.size());
83 for (int i = 0; i < args.size(); ++i) {
84 arg_ptrs.push_back(args[i].data());
85 }
86
87 return {args, arg_ptrs};
88 }
89
90 } // namespace tpu
91 } // namespace tensorflow
92