Merge pull request #1232 from atalman/pick_values_from_release_matrix

Add architecture extraction logic to getting started page generation

Author: atalman

.github/workflows/update-quick-start-module.yml

@@ -107,5 +107,5 @@ jobs:
           commit-message: Modify published_versions.json file
           title: '[Getting Started Page] Modify published_versions.json file'
           body: >
-            This PR is auto-generated Gettins Started page update
+            This PR is auto-generated. It updates Getting Started page
           labels: automated pr

scripts/gen_quick_start_module.py

@@ -30,11 +30,8 @@ class OperatingSystem(Enum):
 ENABLE = "enable"
 DISABLE = "disable"
 
-# Mapping json to release matrix is here for now
-# TBD drive the mapping via:
-#  1. Scanning release matrix and picking 2 latest cuda versions and 1 latest rocm
-#  2. Possibility to override the scanning algorithm with arguments passed from workflow
-acc_arch_ver_map = {
+# Mapping json to release matrix default values
+acc_arch_ver_default = {
     "nightly": {
         "accnone": ("cpu", ""),
         "cuda.x": ("cuda", "11.6"),
@@ -49,6 +46,11 @@ class OperatingSystem(Enum):
         }
     }
 
+# Initialize arch version to default values
+# these default values will be overwritten by
+# extracted values from the release marix
+acc_arch_ver_map = acc_arch_ver_default
+
 LIBTORCH_DWNL_INSTR = {
         PRE_CXX11_ABI: "Download here (Pre-cxx11 ABI):",
         CXX11_ABI: "Download here (cxx11 ABI):",
@@ -163,6 +165,26 @@ def gen_install_matrix(versions) -> Dict[str, str]:
                     result[key] = "<br />".join(lines)
     return result
 
+# This method is used for extracting two latest verisons of cuda and
+# last verion of rocm. It will modify the acc_arch_ver_map object used
+# to update getting started page.
+def extract_arch_ver_map(release_matrix):
+    def gen_ver_list(chan, gpu_arch_type):
+        return {
+            x["desired_cuda"]: x["gpu_arch_version"]
+            for x in release_matrix[chan]["linux"]
+            if x["gpu_arch_type"] == gpu_arch_type
+        }
+
+    for chan in ("nightly", "release"):
+        cuda_ver_list = gen_ver_list(chan, "cuda")
+        rocm_ver_list = gen_ver_list(chan, "rocm")
+        cuda_list = sorted(cuda_ver_list.values())[-2:]
+        acc_arch_ver_map[chan]["rocm5.x"] = ("rocm", max(rocm_ver_list.values()))
+        for cuda_ver, label in zip(cuda_list, ["cuda.x", "cuda.y"]):
+            acc_arch_ver_map[chan][label] = ("cuda", cuda_ver)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--autogenerate', dest='autogenerate', action='store_true')
@@ -178,6 +200,7 @@ def main():
             for osys in OperatingSystem:
                 release_matrix[val][osys.value] = read_matrix_for_os(osys, val)
 
+        extract_arch_ver_map(release_matrix)
         for val in ("nightly", "release"):
             update_versions(versions, release_matrix[val], val)