1 /* Copyright 2015 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/framework/resource_mgr.h" 17 18 #include "tensorflow/core/framework/device_attributes.pb.h" 19 #include "tensorflow/core/framework/node_def.pb.h" 20 #include "tensorflow/core/framework/node_def_util.h" 21 #include "tensorflow/core/lib/core/errors.h" 22 #include "tensorflow/core/lib/core/status_test_util.h" 23 #include "tensorflow/core/lib/strings/strcat.h" 24 #include "tensorflow/core/platform/test.h" 25 26 namespace tensorflow { 27 28 class Resource : public ResourceBase { 29 public: 30 explicit Resource(const string& label) : label_(label) {} 31 ~Resource() override {} 32 33 string DebugString() override { return strings::StrCat("R/", label_); } 34 35 private: 36 string label_; 37 }; 38 39 class Other : public ResourceBase { 40 public: 41 explicit Other(const string& label) : label_(label) {} 42 ~Other() override {} 43 44 string DebugString() override { return strings::StrCat("O/", label_); } 45 46 private: 47 string label_; 48 }; 49 50 template <typename T> 51 string Find(const ResourceMgr& rm, const string& container, 52 const string& name) { 53 T* r; 54 TF_CHECK_OK(rm.Lookup(container, name, &r)); 55 const string ret = r->DebugString(); 56 r->Unref(); 57 return ret; 58 } 59 60 template <typename T> 61 string LookupOrCreate(ResourceMgr* rm, const string& container, 62 const string& name, const string& label) { 63 T* r; 64 TF_CHECK_OK(rm->LookupOrCreate<T>(container, name, &r, [&label](T** ret) { 65 *ret = new T(label); 66 return Status::OK(); 67 })); 68 const string ret = r->DebugString(); 69 r->Unref(); 70 return ret; 71 } 72 73 static void HasError(const Status& s, const string& substr) { 74 EXPECT_TRUE(StringPiece(s.ToString()).contains(substr)) 75 << s << ", expected substring " << substr; 76 } 77 78 template <typename T> 79 Status FindErr(const ResourceMgr& rm, const string& container, 80 const string& name) { 81 T* r; 82 Status s = rm.Lookup(container, name, &r); 83 CHECK(!s.ok()); 84 return s; 85 } 86 87 TEST(ResourceMgrTest, Basic) { 88 ResourceMgr rm; 89 TF_CHECK_OK(rm.Create("foo", "bar", new Resource("cat"))); 90 TF_CHECK_OK(rm.Create("foo", "baz", new Resource("dog"))); 91 TF_CHECK_OK(rm.Create("foo", "bar", new Other("tiger"))); 92 93 // Expected to fail. 94 HasError(rm.Create("foo", "bar", new Resource("kitty")), 95 "Already exists: Resource foo/bar"); 96 97 // Expected to be found. 98 EXPECT_EQ("R/cat", Find<Resource>(rm, "foo", "bar")); 99 EXPECT_EQ("R/dog", Find<Resource>(rm, "foo", "baz")); 100 EXPECT_EQ("O/tiger", Find<Other>(rm, "foo", "bar")); 101 102 // Expected to be not found. 103 HasError(FindErr<Resource>(rm, "bar", "foo"), "Not found: Container bar"); 104 HasError(FindErr<Resource>(rm, "foo", "xxx"), "Not found: Resource foo/xxx"); 105 HasError(FindErr<Other>(rm, "foo", "baz"), "Not found: Resource foo/baz"); 106 107 // Delete foo/bar/Resource. 108 TF_CHECK_OK(rm.Delete<Resource>("foo", "bar")); 109 HasError(FindErr<Resource>(rm, "foo", "bar"), "Not found: Resource foo/bar"); 110 111 TF_CHECK_OK(rm.Create("foo", "bar", new Resource("kitty"))); 112 EXPECT_EQ("R/kitty", Find<Resource>(rm, "foo", "bar")); 113 114 // Drop the whole container foo. 115 TF_CHECK_OK(rm.Cleanup("foo")); 116 HasError(FindErr<Resource>(rm, "foo", "bar"), "Not found: Container foo"); 117 118 // Dropping it a second time is OK. 119 TF_CHECK_OK(rm.Cleanup("foo")); 120 HasError(FindErr<Resource>(rm, "foo", "bar"), "Not found: Container foo"); 121 122 // Dropping a non-existent container is also ok. 123 TF_CHECK_OK(rm.Cleanup("bar")); 124 } 125 126 TEST(ResourceMgr, CreateOrLookup) { 127 ResourceMgr rm; 128 EXPECT_EQ("R/cat", LookupOrCreate<Resource>(&rm, "foo", "bar", "cat")); 129 EXPECT_EQ("R/cat", LookupOrCreate<Resource>(&rm, "foo", "bar", "dog")); 130 EXPECT_EQ("R/cat", Find<Resource>(rm, "foo", "bar")); 131 132 EXPECT_EQ("O/tiger", LookupOrCreate<Other>(&rm, "foo", "bar", "tiger")); 133 EXPECT_EQ("O/tiger", LookupOrCreate<Other>(&rm, "foo", "bar", "lion")); 134 TF_CHECK_OK(rm.Delete<Other>("foo", "bar")); 135 HasError(FindErr<Other>(rm, "foo", "bar"), "Not found: Resource foo/bar"); 136 } 137 138 Status ComputePolicy(const string& attr_container, 139 const string& attr_shared_name, 140 bool use_node_name_as_default, string* result) { 141 ContainerInfo cinfo; 142 ResourceMgr rmgr; 143 NodeDef ndef; 144 ndef.set_name("foo"); 145 if (attr_container != "none") { 146 AddNodeAttr("container", attr_container, &ndef); 147 } 148 if (attr_shared_name != "none") { 149 AddNodeAttr("shared_name", attr_shared_name, &ndef); 150 } 151 TF_RETURN_IF_ERROR(cinfo.Init(&rmgr, ndef, use_node_name_as_default)); 152 *result = cinfo.DebugString(); 153 return Status::OK(); 154 } 155 156 string Policy(const string& attr_container, const string& attr_shared_name, 157 bool use_node_name_as_default) { 158 string ret; 159 TF_CHECK_OK(ComputePolicy(attr_container, attr_shared_name, 160 use_node_name_as_default, &ret)); 161 return ret; 162 } 163 164 TEST(ContainerInfo, Basic) { 165 // Correct cases. 166 EXPECT_EQ(Policy("", "", false), "[localhost,_0_foo,private]"); 167 EXPECT_EQ(Policy("", "", true), "[localhost,foo,public]"); 168 EXPECT_EQ(Policy("", "bar", false), "[localhost,bar,public]"); 169 EXPECT_EQ(Policy("", "bar", true), "[localhost,bar,public]"); 170 EXPECT_EQ(Policy("cat", "", false), "[cat,_1_foo,private]"); 171 EXPECT_EQ(Policy("cat", "", true), "[cat,foo,public]"); 172 EXPECT_EQ(Policy("cat", "bar", false), "[cat,bar,public]"); 173 EXPECT_EQ(Policy("cat", "bar", true), "[cat,bar,public]"); 174 EXPECT_EQ(Policy("cat.0-dog", "bar", true), "[cat.0-dog,bar,public]"); 175 EXPECT_EQ(Policy(".cat", "bar", true), "[.cat,bar,public]"); 176 } 177 178 Status WrongPolicy(const string& attr_container, const string& attr_shared_name, 179 bool use_node_name_as_default) { 180 string dbg; 181 auto s = ComputePolicy(attr_container, attr_shared_name, 182 use_node_name_as_default, &dbg); 183 CHECK(!s.ok()); 184 return s; 185 } 186 187 TEST(ContainerInfo, Error) { 188 // Missing attribute. 189 HasError(WrongPolicy("none", "", false), "No attr"); 190 HasError(WrongPolicy("", "none", false), "No attr"); 191 HasError(WrongPolicy("none", "none", false), "No attr"); 192 193 // Invalid container. 194 HasError(WrongPolicy("12$%", "", false), "container contains invalid char"); 195 HasError(WrongPolicy("-cat", "", false), "container contains invalid char"); 196 197 // Invalid shared name. 198 HasError(WrongPolicy("", "_foo", false), "shared_name cannot start with '_'"); 199 } 200 201 // Stub DeviceBase subclass which only sets a device name, for testing resource 202 // handles. 203 class StubDevice : public DeviceBase { 204 public: 205 explicit StubDevice(const string& name) : DeviceBase(nullptr) { 206 attr_.set_name(name); 207 } 208 209 Allocator* GetAllocator(AllocatorAttributes) override { 210 return cpu_allocator(); 211 } 212 213 const DeviceAttributes& attributes() const override { return attr_; } 214 215 private: 216 DeviceAttributes attr_; 217 }; 218 219 // Empty stub resource for testing resource handles. 220 class StubResource : public ResourceBase { 221 public: 222 string DebugString() override { return ""; } 223 int value_{0}; 224 }; 225 226 TEST(ResourceHandleTest, CRUD) { 227 ResourceMgr resource_mgr(""); 228 OpKernelContext::Params params; 229 params.resource_manager = &resource_mgr; 230 StubDevice device("device_name"); 231 params.device = &device; 232 OpKernelContext ctx(¶ms, 0); 233 234 ResourceHandle p = 235 MakeResourceHandle<StubResource>(&ctx, "container", "name"); 236 237 { 238 auto* r = new StubResource(); 239 r->value_ = 42; 240 TF_EXPECT_OK(CreateResource(&ctx, p, r)); 241 } 242 { 243 StubResource* r = nullptr; 244 TF_ASSERT_OK(LookupResource(&ctx, p, &r)); 245 ASSERT_TRUE(r != nullptr); 246 EXPECT_EQ(r->value_, 42); 247 r->Unref(); 248 } 249 { 250 TF_EXPECT_OK(DeleteResource<StubResource>(&ctx, p)); 251 StubResource* unused = nullptr; 252 EXPECT_FALSE(LookupResource(&ctx, p, &unused).ok()); 253 } 254 } 255 256 TEST(ResourceHandleTest, DifferentDevice) { 257 ResourceMgr resource_mgr(""); 258 OpKernelContext::Params params; 259 params.resource_manager = &resource_mgr; 260 StubDevice device("device_name"); 261 params.device = &device; 262 OpKernelContext ctx(¶ms, 0); 263 264 ResourceHandle p = 265 MakeResourceHandle<StubResource>(&ctx, "container", "name"); 266 267 ResourceMgr other_resource_mgr(""); 268 OpKernelContext::Params other_params; 269 other_params.resource_manager = &other_resource_mgr; 270 StubDevice other_device("other_device_name"); 271 other_params.device = &other_device; 272 OpKernelContext other_ctx(&other_params, 0); 273 274 auto* r = new StubResource(); 275 ASSERT_FALSE(CreateResource(&other_ctx, p, r).ok()); 276 r->Unref(); 277 } 278 279 // Other stub resource to test type-checking of resource handles. 280 class OtherStubResource : public ResourceBase { 281 public: 282 string DebugString() override { return ""; } 283 }; 284 285 TEST(ResourceHandleTest, DifferentType) { 286 ResourceMgr resource_mgr(""); 287 OpKernelContext::Params params; 288 params.resource_manager = &resource_mgr; 289 StubDevice device("device_name"); 290 params.device = &device; 291 OpKernelContext ctx(¶ms, 0); 292 293 ResourceHandle p = 294 MakeResourceHandle<StubResource>(&ctx, "container", "name"); 295 296 auto* r = new OtherStubResource; 297 ASSERT_FALSE(CreateResource(&ctx, p, r).ok()); 298 r->Unref(); 299 } 300 301 TEST(ResourceHandleTest, DeleteUsingResourceHandle) { 302 ResourceMgr resource_mgr(""); 303 OpKernelContext::Params params; 304 params.resource_manager = &resource_mgr; 305 StubDevice device("device_name"); 306 params.device = &device; 307 OpKernelContext ctx(¶ms, 0); 308 309 ResourceHandle p = 310 MakeResourceHandle<StubResource>(&ctx, "container", "name"); 311 312 StubResource* r = new StubResource; 313 TF_EXPECT_OK(CreateResource(&ctx, p, r)); 314 315 StubResource* lookup_r = nullptr; 316 TF_EXPECT_OK(LookupResource<StubResource>(&ctx, p, &lookup_r)); 317 EXPECT_EQ(lookup_r, r); 318 319 TF_EXPECT_OK(DeleteResource(&ctx, p)); 320 EXPECT_NE(LookupResource<StubResource>(&ctx, p, &lookup_r).ok(), true); 321 r->Unref(); 322 } 323 324 } // end namespace tensorflow 325
