From d31eb2a5e60d493cfc9c6dc50f2fa52e85e2a218 Mon Sep 17 00:00:00 2001
From: Yifan Xiong <yifan.xiong@microsoft.com>
Date: Mon, 21 Sep 2020 13:48:25 +0800
Subject: [PATCH] [Webportal] Support dynamic sku types for different vc
 (#4900)

Support dynamic sku types for different virtual clusters on webportal.
---
 src/rest-server/src/controllers/v2/cluster.js | 35 +++++++++++++++++--
 .../job-submission/job-submission-page.jsx    |  4 +--
 .../src/app/job-submission/utils/conn.js      | 17 ++++++---
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/src/rest-server/src/controllers/v2/cluster.js b/src/rest-server/src/controllers/v2/cluster.js
index dd772ba9e7..89c6e505b8 100644
--- a/src/rest-server/src/controllers/v2/cluster.js
+++ b/src/rest-server/src/controllers/v2/cluster.js
@@ -16,11 +16,40 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 // module dependencies
+const axios = require('axios');
 const status = require('statuses');
+const createError = require('@pai/utils/error');
 const { resourceUnits } = require('@pai/config/vc');
+const { hivedWebserviceUri } = require('@pai/config/launcher');
+const asyncHandler = require('@pai/middlewares/v2/asyncHandler');
 
-const getSkuTypes = (req, res) => {
-  res.status(status('OK')).json(resourceUnits);
-};
+const getSkuTypes = asyncHandler(async (req, res) => {
+  if (req.query.vc) {
+    let vcStatus;
+    try {
+      vcStatus = (
+        await axios.get(
+          `${hivedWebserviceUri}/v1/inspect/clusterstatus/virtualclusters/${req.query.vc}`,
+        )
+      ).data;
+    } catch (error) {
+      throw createError(
+        'Not Found',
+        'NoVirtualClusterError',
+        `Cannot get sku types for virtual clyster ${req.query.vc}.`,
+      );
+    }
+    const leafCellTypes = new Set(vcStatus.map((cell) => cell.leafCellType));
+    const skuTypes = Object.keys(resourceUnits)
+      .filter((key) => leafCellTypes.has(key))
+      .reduce((obj, key) => {
+        obj[key] = resourceUnits[key];
+        return obj;
+      }, {});
+    res.status(status('OK')).json(skuTypes);
+  } else {
+    res.status(status('OK')).json(resourceUnits);
+  }
+});
 
 module.exports = { getSkuTypes };
diff --git a/src/webportal/src/app/job-submission/job-submission-page.jsx b/src/webportal/src/app/job-submission/job-submission-page.jsx
index 5c779c8280..86833192a0 100644
--- a/src/webportal/src/app/job-submission/job-submission-page.jsx
+++ b/src/webportal/src/app/job-submission/job-submission-page.jsx
@@ -349,12 +349,12 @@ export const JobSubmissionPage = ({
   }, []);
 
   useEffect(() => {
-    listHivedSkuTypes()
+    listHivedSkuTypes(jobInformation.virtualCluster)
       .then(hivedSkuTypes => {
         setHivedSkuTypes(hivedSkuTypes);
       })
       .catch(alert);
-  }, []);
+  }, [jobInformation.virtualCluster]);
 
   const onToggleAdvanceFlag = useCallback(() => {
     setAdvanceFlag(!advanceFlag);
diff --git a/src/webportal/src/app/job-submission/utils/conn.js b/src/webportal/src/app/job-submission/utils/conn.js
index 0582b1c2dd..fd3013f951 100644
--- a/src/webportal/src/app/job-submission/utils/conn.js
+++ b/src/webportal/src/app/job-submission/utils/conn.js
@@ -6,6 +6,7 @@ import { clearToken } from '../../user/user-logout/user-logout.component.js';
 import config from '../../config/webportal.config';
 import yaml from 'js-yaml';
 import { get } from 'lodash';
+import urljoin from 'url-join';
 
 const token = cookies.get('token');
 
@@ -52,16 +53,22 @@ export async function listUserVirtualClusters(user) {
   return get(userInfo, 'virtualCluster', []);
 }
 
-export async function listHivedSkuTypes() {
+export async function listHivedSkuTypes(virtualCluster) {
   if (config.launcherScheduler !== 'hivedscheduler') {
     return {};
   }
   return wrapper(async () =>
-    (await fetch(`${config.restServerUri}/api/v2/cluster/sku-types`, {
-      headers: {
-        Authorization: `Bearer ${token}`,
+    (await fetch(
+      urljoin(
+        config.restServerUri,
+        `/api/v2/cluster/sku-types?vc=${virtualCluster}`,
+      ),
+      {
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
       },
-    })).json(),
+    )).json(),
   );
 }