handle denpency

Signed-off-by: wmqwxb <[email protected]>
FederatedAI · Nov 28, 2023 · a4397bb · a4397bb
1 parent 4292723
commit a4397bb
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 3 deletions.
diff --git a/src/main/java/org/fedai/fate/board/global/Dict.java b/src/main/java/org/fedai/fate/board/global/Dict.java
@@ -134,7 +134,7 @@ public class Dict {
         logTypeMap.put("partyWarning", "task_warning");
         logTypeMap.put("partyInfo", "task_info");
         logTypeMap.put("partyDebug", "task_debug");
-        logTypeMap.put("componentInfo", "componentInfo");
+        logTypeMap.put("componentInfo", "task_info");
     }
 
     //the fields for job search

diff --git a/src/main/java/org/fedai/fate/board/services/JobWebSocketService.java b/src/main/java/org/fedai/fate/board/services/JobWebSocketService.java
@@ -16,6 +16,7 @@
 package org.fedai.fate.board.services;
 
 import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
 import com.google.common.collect.Maps;
 import org.fedai.fate.board.controller.JobDetailController;
@@ -32,6 +33,7 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 public class JobWebSocketService implements Runnable {
     private static final Logger logger = LoggerFactory.getLogger(JobWebSocketService.class);
@@ -128,6 +130,8 @@ public void run() {
                     ResponseResult<Map<String, Object>> mapResponseResult = responseResultListenableFuture.get();
                     if (0 == mapResponseResult.getCode()) {
                         flushToWebData.put(Dict.SUMMARY_DATA, mapResponseResult.getData());
+                        JSONObject jsonObject = handleDependencies(mapResponseResult.getData());
+                        responseResult.getData().put("newDependencies",jsonObject);
                     } else {
                         throw new IllegalArgumentException("summary parameter error");
                     }
@@ -150,7 +154,6 @@ public void run() {
                             break;
 
 
-
                         }
                         Thread.sleep(500);
                     } else {
@@ -166,7 +169,6 @@ public void run() {
                     }
                     break;
                 }
-
             }
         } catch (Exception e) {
             logger.error("web socket error", e);
@@ -182,4 +184,68 @@ public void run() {
             }
         }
     }
+
+    private JSONObject handleDependencies(Map<String, Object> map) {
+        JSONObject dataset = (JSONObject) map.get("dataset");
+        JSONObject dags = dataset.getJSONObject("dag");
+        JSONObject dag = dags.getJSONObject("dag");
+        JSONObject tasks = dag.getJSONObject("tasks");
+        Set<String> components = tasks.keySet();
+        JSONObject result = new JSONObject();
+        for (String component : components) {
+            JSONArray jsonArray = new JSONArray();
+
+            JSONObject info = tasks.getJSONObject(component);
+            JSONObject inputs = info.getJSONObject("inputs");
+
+            JSONObject data = inputs.getJSONObject("data");
+            Set<String> datas = data.keySet();
+            for (String dataKey : datas) {
+                JSONObject jsonObject = new JSONObject();
+                JSONObject task_output_artifact = data.getJSONObject("task_output_artifact");
+                String type = dataKey;
+                String modelType = "data";
+                boolean dataSource = false;
+                String componentName = task_output_artifact.getString("producer_task");
+                String outputInfo = task_output_artifact.getString("output_artifact_key");
+
+                jsonObject.put("conponent_name", componentName);
+                jsonObject.put("data_source", dataSource);
+                jsonObject.put("model_type", modelType);
+                jsonObject.put("type", type);
+
+                JSONArray upOutPutInfo = new JSONArray();
+                upOutPutInfo.add(outputInfo);
+                jsonObject.put("up_output_info", upOutPutInfo);
+                jsonArray.add(jsonObject);
+            }
+
+            JSONObject model = inputs.getJSONObject("model");
+            Set<String> models = model.keySet();
+            for (String modelKey : models) {
+                JSONObject jsonObject = new JSONObject();
+                JSONObject task_output_artifact = data.getJSONObject("task_output_artifact");
+                String type = modelKey;
+                String modelType = "model";
+                boolean dataSource = false;
+                String componentName = task_output_artifact.getString("producer_task");
+                String outputInfo = task_output_artifact.getString("output_artifact_key");
+
+                jsonObject.put("conponent_name", componentName);
+                jsonObject.put("data_source", dataSource);
+                jsonObject.put("model_type", modelType);
+                jsonObject.put("type", type);
+
+                JSONArray upOutPutInfo = new JSONArray();
+                upOutPutInfo.add(outputInfo);
+
+                jsonObject.put("up_output_info", upOutPutInfo);
+
+                jsonArray.add(jsonObject);
+            }
+            result.put(component, jsonArray);
+
+        }
+        return result;
+    }
 }