Add util function to verify if override resource have a subset of dev…

…ice kind of base resource. This is not used by OSS. PiperOrigin-RevId: 595245998
tensorflow · Jan 3, 2024 · 06ff18d · 06ff18d
1 parent 8fc366a
commit 06ff18d
Show file tree

Hide file tree

Showing 4 changed files with 86 additions and 0 deletions.
diff --git a/tensorflow_serving/resources/BUILD b/tensorflow_serving/resources/BUILD
@@ -46,6 +46,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":resources_cc_proto",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@org_tensorflow//tensorflow/core:lib",
     ],
 )

diff --git a/tensorflow_serving/resources/resource_util.cc b/tensorflow_serving/resources/resource_util.cc
@@ -23,9 +23,12 @@ limitations under the License.
 #include <vector>
 
 #include "google/protobuf/wrappers.pb.h"
+#include "absl/container/flat_hash_set.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/types.h"
+#include "tsl/platform/errors.h"
 
 namespace tensorflow {
 namespace serving {
@@ -144,6 +147,28 @@ Status ResourceUtil::VerifyResourceValidity(const Resource& resource) const {
   return result;
 }
 
+Status ResourceUtil::VerifyOverrideDeviceValidity(
+    const ResourceAllocation& base_allocation,
+    const ResourceAllocation& override_allocation) const {
+  absl::flat_hash_set<std::pair<std::string, std::string>>
+      base_device_kind_pairs;
+  for (const auto& entry : base_allocation.resource_quantities()) {
+    base_device_kind_pairs.insert(
+        {entry.resource().device(), entry.resource().kind()});
+  }
+  for (const auto& entry : override_allocation.resource_quantities()) {
+    if (base_device_kind_pairs.find(
+            {entry.resource().device(), entry.resource().kind()}) ==
+        base_device_kind_pairs.end()) {
+      return errors::InvalidArgument(
+          "Invalid resource allocation: device-kind from override "
+          "resource was not found in base resource: ",
+          entry.resource().DebugString());
+    }
+  }
+  return Status();
+}
+
 ResourceAllocation ResourceUtil::Normalize(
     const ResourceAllocation& allocation) const {
   return NormalizeResourceAllocation(allocation);

diff --git a/tensorflow_serving/resources/resource_util.h b/tensorflow_serving/resources/resource_util.h
@@ -49,6 +49,12 @@ class ResourceUtil {
   // have undefined behavior otherwise), and guarantee to produce valid outputs.
   virtual Status VerifyValidity(const ResourceAllocation& allocation) const;
 
+  // Determine if the override_allocation's device set is a subset of
+  // base_allocation. Only used by controller
+  Status VerifyOverrideDeviceValidity(
+      const ResourceAllocation& base_allocation,
+      const ResourceAllocation& override_allocation) const;
+
   // Verifies whether 'resource' is valid, i.e. it only refers to valid devices,
   // i.e. those supplied via Options.
   Status VerifyResourceValidity(const Resource& resource) const;

diff --git a/tensorflow_serving/resources/resource_util_test.cc b/tensorflow_serving/resources/resource_util_test.cc
@@ -1818,6 +1818,60 @@ TEST_F(ResourceUtilTest, MinUnbound) {
                           "} "));
 }
 
+TEST_F(ResourceUtilTest, OverrideDeviceValidity) {
+  const auto base_resource = CreateProto<ResourceAllocation>(
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'main' "
+      "    kind: 'processing' "
+      "  } "
+      "  quantity: 100 "
+      "} "
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'tpu' "
+      "    kind: 'hbm' "
+      "  } "
+      "  quantity: 8 "
+      "} ");
+  const auto valid_override_resource = CreateProto<ResourceAllocation>(
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'main' "
+      "    kind: 'processing' "
+      "  } "
+      "  quantity: 300 "
+      "} "
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'tpu' "
+      "    kind: 'hbm' "
+      "  } "
+      "  quantity: 4 "
+      "} ");
+  const auto invalid_override_resource = CreateProto<ResourceAllocation>(
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'main' "
+      "    kind: 'processing' "
+      "  } "
+      "  quantity: 300 "
+      "} "
+      "resource_quantities { "
+      "  resource { "
+      "    device: 'gpu' "
+      "    kind: 'ram' "
+      "  } "
+      "  quantity: 4 "
+      "} ");
+  TF_EXPECT_OK(util_.VerifyOverrideDeviceValidity(base_resource,
+                                                  valid_override_resource));
+  EXPECT_FALSE(util_
+                   .VerifyOverrideDeviceValidity(base_resource,
+                                                 invalid_override_resource)
+                   .ok());
+}
+
 }  // namespace
 }  // namespace serving
 }  // namespace tensorflow