Support GCP custom instance types family extraction (#3376)

epam · Oct 4, 2023 · 3e66c1b · 3e66c1b
1 parent 2713cb1
commit 3e66c1b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 10 deletions.
diff --git a/workflows/pipe-common/pipeline/hpc/utils.py b/workflows/pipe-common/pipeline/hpc/utils.py
@@ -38,14 +38,22 @@ def get_run_id_from_host(self, host):
 
     def extract_family_from_instance_type(self, cloud_provider, instance_type):
         if cloud_provider == CloudProvider.aws():
-            search = re.search('^(\w+)\..*', instance_type)
-            return search.group(1) if search else None
+            search = re.search('^(\\w+)\\..*', instance_type)
+            if search:
+                return search.group(1)
         elif cloud_provider == CloudProvider.gcp():
-            search = re.search('^(?!custom)(\w+-(?!custom)\w+)-?.*', instance_type)
-            return search.group(1) if search else None
+            search = re.search('^(\\w*-?custom)-\\w+-\\w+$', instance_type)
+            if search:
+                return search.group(1)
+            search = re.search('^(gpu-\\w*-?custom)-\\w+-\\w+(-\\w+)-\\w+$', instance_type)
+            if search:
+                return search.group(1) + search.group(2)
+            search = re.search('^(\\w+-\\w+)-?\\w*(-?\\w*)', instance_type)
+            if search:
+                return search.group(1) + search.group(2)
         elif cloud_provider == CloudProvider.azure():
             # will return Bms for Standard_B1ms or Dsv3 for Standard_D2s_v3 instance types
-            search = re.search('^([a-zA-Z]+)\d+(.*)', instance_type.split('_', 1)[1].replace('_', ''))
-            return search.group(1) + search.group(2) if search else None
-        else:
-            return None
+            search = re.search('^([a-zA-Z]+)\\d+(.*)', instance_type.split('_', 1)[1].replace('_', ''))
+            if search:
+                return search.group(1) + search.group(2)
+        return None
diff --git a/workflows/pipe-common/test/test_extract_family_from_instance_type.py b/workflows/pipe-common/test/test_extract_family_from_instance_type.py
@@ -33,8 +33,14 @@ def test_gcp_familes():
     assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "n2-highcpu-2") == "n2-highcpu"
     assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "n2d-highcpu-128") == "n2d-highcpu"
     assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "e2-small") == "e2-small"
-    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "custom-12-16") is None
-    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "gpu-custom-4-16384-k80-1") is None
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "c3-standard-4") == "c3-standard"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "c3-standard-4-lssd") == "c3-standard-lssd"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "a2-highgpu-1g") == "a2-highgpu"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "a2-ultragpu-4g") == "a2-ultragpu"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "custom-12-16384") == "custom"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "n2-custom-12-16384") == "n2-custom"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "gpu-custom-4-16384-k80-1") == "gpu-custom-k80"
+    assert utils.extract_family_from_instance_type(CloudProvider.gcp(), "gpu-n2-custom-4-16384-k80-1") == "gpu-n2-custom-k80"
 
 
 def test_azure_familes():