Repository: sql-machine-learning/elasticdl

Issue:
Provide default gpu resource name and validation
Currently users can only pass "gpu" as part of the resource name. However, k8s requires it to be either `"nvidia.com/gpu"` or `"amd.com/gpu"` if AMD plugin is enabled. There are other different vendors as well but a pattern to use for validation would be `"<vendor>.com/gpu"`.

We should consider adding `"nvidia.com/gpu"` as the default and validate for user provided gpu resource name based on the pattern `"<vendor>.com/gpu"`.

Reference: https://kubernetes.io/docs/tasks/manage-gpus/scheduling-gpus/


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/common/k8s_utils.py b/elasticdl/python/elasticdl/common/k8s_utils.py
--- a/elasticdl/python/elasticdl/common/k8s_utils.py
+++ b/elasticdl/python/elasticdl/common/k8s_utils.py
@@ -2,6 +2,10 @@
 
 
 _ALLOWED_RESOURCE_TYPES = ["memory", "disk", "ephemeral-storage", "cpu", "gpu"]
+# Any domain name is (syntactically) valid if it's a dot-separated list of
+# identifiers, each no longer than 63 characters, and made up of letters,
+# digits and dashes (no underscores).
+_GPU_VENDOR_REGEX_STR = r"^[a-zA-Z\d-]{,63}(\.[a-zA-Z\d-]{,63})*/gpu$"
 
 
 def _is_numeric(n):
@@ -57,7 +61,14 @@ def parse_resource(resource_str):
             _valid_mem_spec(v)
         elif k == "cpu":
             _valid_cpu_spec(v)
-        elif k == "gpu":
+        elif "gpu" in k:
+            if k == "gpu":
+                k = "nvidia.com/gpu"
+            elif not re.compile(_GPU_VENDOR_REGEX_STR).match(k):
+                raise ValueError(
+                    "gpu resource name does not have a valid vendor name: %s"
+                    % k
+                )
             _valid_gpu_spec(v)
         else:
             raise ValueError(


Test patch:
diff --git a/elasticdl/python/tests/k8s_utils_test.py b/elasticdl/python/tests/k8s_utils_test.py
--- a/elasticdl/python/tests/k8s_utils_test.py
+++ b/elasticdl/python/tests/k8s_utils_test.py
@@ -11,7 +11,7 @@ def test_parse_resource(self):
                 "cpu": "250m",
                 "memory": "32Mi",
                 "disk": "64Mi",
-                "gpu": "1",
+                "nvidia.com/gpu": "1",
                 "ephemeral-storage": "32Mi",
             },
             parse_resource(
@@ -20,17 +20,32 @@ def test_parse_resource(self):
         )
         # When cpu is non-numeric, parse_resource works as expected
         self.assertEqual(
-            {"cpu": "250m", "memory": "32Mi", "disk": "64Mi", "gpu": "1"},
+            {
+                "cpu": "250m",
+                "memory": "32Mi",
+                "disk": "64Mi",
+                "nvidia.com/gpu": "1",
+            },
             parse_resource("cpu=250m,memory=32Mi,disk=64Mi,gpu=1"),
         )
         # When cpu is integer, parse_resource works as expected
         self.assertEqual(
-            {"cpu": "1", "memory": "32Mi", "disk": "64Mi", "gpu": "1"},
+            {
+                "cpu": "1",
+                "memory": "32Mi",
+                "disk": "64Mi",
+                "nvidia.com/gpu": "1",
+            },
             parse_resource("cpu=1,memory=32Mi,disk=64Mi,gpu=1"),
         )
         # When cpu is float, parse_resource works as expected
         self.assertEqual(
-            {"cpu": "0.1", "memory": "32Mi", "disk": "64Mi", "gpu": "1"},
+            {
+                "cpu": "0.1",
+                "memory": "32Mi",
+                "disk": "64Mi",
+                "nvidia.com/gpu": "1",
+            },
             parse_resource("cpu=0.1,memory=32Mi,disk=64Mi,gpu=1"),
         )
         # When cpu is non-numeric, raise an error
@@ -54,6 +69,45 @@ def test_parse_resource(self):
             parse_resource,
             "cpu=2,memory=32Mi,disk=64Mi,gpu=0.1",
         )
+        # When gpu resource name has a valid vendor name,
+        # parse_resource works as expected
+        self.assertEqual(
+            {
+                "cpu": "0.1",
+                "memory": "32Mi",
+                "disk": "64Mi",
+                "amd.com/gpu": "1",
+            },
+            parse_resource("cpu=0.1,memory=32Mi,disk=64Mi,amd.com/gpu=1"),
+        )
+        # When gpu resource name does not have a valid vendor name,
+        # raise an error
+        self.assertRaisesRegex(
+            ValueError,
+            "gpu resource name does not have a valid vendor name: blah-gpu",
+            parse_resource,
+            "cpu=2,memory=32Mi,disk=64Mi,blah-gpu=1",
+        )
+        # When gpu resource name does not have a valid vendor name,
+        # raise an error
+        self.assertRaisesRegex(
+            ValueError,
+            "gpu resource name does not have a valid vendor name: @#/gpu",
+            parse_resource,
+            "cpu=2,memory=32Mi,disk=64Mi,@#/gpu=1",
+        )
+        self.assertRaisesRegex(
+            ValueError,
+            "gpu resource name does not have a valid vendor name: a_d.com/gpu",
+            parse_resource,
+            "cpu=2,memory=32Mi,disk=64Mi,a_d.com/gpu=1",
+        )
+        self.assertRaisesRegex(
+            ValueError,
+            "gpu resource name does not have a valid vendor name: *",
+            parse_resource,
+            "cpu=2,memory=32Mi,disk=64Mi," + "a" * 64 + "/gpu=1",
+        )
         # When memory does not contain expected regex, raise an error
         self.assertRaisesRegex(
             ValueError,