Merge pull request #677 from DingPengfei/add_fair_scheduler_parser

add parser of fair scheduler from yarn scheduler API
DTStack · Aug 5, 2022 · 6acb463 · 6acb463
2 parents 9fca9eb + 4e5dea4
commit 6acb463
Showing 1 changed file with 39 additions and 8 deletions.
diff --git a/.../taier-plugins/schedules/yarn-core/src/main/java/com/dtstack/taier/yarn/DtYarnClient.java b/.../taier-plugins/schedules/yarn-core/src/main/java/com/dtstack/taier/yarn/DtYarnClient.java
@@ -266,16 +266,26 @@ private List<JSONObject> getQueueResource(YarnClient yarnClient) throws Exceptio
             return null;
         }
         JSONObject schedulerInfoJson = schedulerJson.getJSONObject("schedulerInfo");
-        if (!schedulerInfoJson.containsKey("queues")) {
-            LOG.error("get yarn queueInfo error! Miss queues field");
-            return null;
+
+        String schedulerType = schedulerInfoJson.getString("type");
+        if (StringUtils.equalsIgnoreCase(schedulerType, ConfigConstrant.CAPACITYSCHEDULER_TPYE)) {
+
+            if (!schedulerInfoJson.containsKey("queues")) {
+                LOG.error("get yarn queueInfo error! Miss queues field");
+                return null;
+            }
+            JSONObject queuesJson = schedulerInfoJson.getJSONObject("queues");
+            return modifyCapQueueInfo(null, queuesJson);
+        }
+        if (StringUtils.equalsIgnoreCase(schedulerType, ConfigConstrant.FAIRSCHEDULER_TPYE)) {
+            JSONObject rootQueueJson = schedulerInfoJson.getJSONObject("rootQueue");
+            JSONObject queuesJson = rootQueueJson.getJSONObject("childQueues");
+            return modifyFairQueueInfo(queuesJson);
         }
-        JSONObject queuesJson = schedulerInfoJson.getJSONObject("queues");
-        List<JSONObject> modifyQueueInfos = modifyQueueInfo(null, queuesJson);
-        return modifyQueueInfos;
+        return null;
     }
 
-    private List<JSONObject> modifyQueueInfo(String parentName, JSONObject queueInfos) {
+    private List<JSONObject> modifyCapQueueInfo(String parentName, JSONObject queueInfos) {
         List<JSONObject> queues = new ArrayList<>();
         if (!queueInfos.containsKey("queue")) {
             return null;
@@ -288,7 +298,7 @@ private List<JSONObject> modifyQueueInfo(String parentName, JSONObject queueInfo
             String queueNewName = parentName + queueName;
 
             if (queueInfo.containsKey("queues")) {
-                List<JSONObject> childQueues = modifyQueueInfo(queueNewName, queueInfo.getJSONObject("queues"));
+                List<JSONObject> childQueues = modifyCapQueueInfo(queueNewName, queueInfo.getJSONObject("queues"));
                 if (childQueues != null) {
                     queues.addAll(childQueues);
                 }
@@ -303,6 +313,27 @@ private List<JSONObject> modifyQueueInfo(String parentName, JSONObject queueInfo
         }
         return queues;
     }
+
+    private List<JSONObject> modifyFairQueueInfo(JSONObject queueInfos) {
+        List<JSONObject> queues = new ArrayList<>();
+        if (!queueInfos.containsKey("queue")) {
+            return null;
+        }
+
+        for (Object ob : queueInfos.getJSONArray("queue")) {
+            JSONObject queueInfo = (JSONObject) ob;
+
+            if (queueInfo.containsKey("childQueues")) {
+                List<JSONObject> childQueues = modifyFairQueueInfo(queueInfo.getJSONObject("queues"));
+                if (childQueues != null) {
+                    queues.addAll(childQueues);
+                }
+            } else {
+                queues.add(queueInfo);
+            }
+        }
+        return queues;
+    }
 
     private void retainCapacity(JSONObject queueInfo) {
         Double capacity = queueInfo.getDouble("capacity");