Migrate ML Worker explain method to WebSocket

backend/src/main/java/ai/giskard/ml/dto/MLWorkerWSExplainDTO.java

@@ -0,0 +1,11 @@
+package ai.giskard.ml.dto;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.Getter;
+
+import java.util.Map;
+
+@Getter
+public class MLWorkerWSExplainDTO implements MLWorkerWSBaseDTO {
+    Map<String, MLWorkerWSExplanationDTO> explanations;
+}

backend/src/main/java/ai/giskard/ml/dto/MLWorkerWSExplainParamDTO.java

@@ -0,0 +1,16 @@
+package ai.giskard.ml.dto;
+
+import lombok.Getter;
+import lombok.Setter;
+
+import java.util.Map;
+
+@Getter
+@Setter
+public class MLWorkerWSExplainParamDTO implements MLWorkerWSBaseDTO {
+    private MLWorkerWSArtifactRefDTO model;
+
+    private MLWorkerWSArtifactRefDTO dataset;
+
+    private Map<String, String> columns;
+}

backend/src/main/java/ai/giskard/ml/dto/MLWorkerWSExplanationDTO.java

@@ -0,0 +1,12 @@
+package ai.giskard.ml.dto;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.Getter;
+
+import java.util.Map;
+
+@Getter
+public class MLWorkerWSExplanationDTO {
+    @JsonProperty("per_feature")
+    Map<String, Float> perFeature;
+}

backend/src/main/java/ai/giskard/service/ModelService.java

@@ -105,18 +105,32 @@ public boolean shouldDrop(String columnDtype, String value) {
             ((columnDtype.startsWith("int") || columnDtype.startsWith("float")) && Strings.isBlank(value));
     }
 
-    public ExplainResponse explain(ProjectModel model, Dataset dataset, Map<String, String> features) {
-        try (MLWorkerClient client = mlWorkerService.createClient(model.getProject().isUsingInternalWorker())) {
-            ExplainRequest request = ExplainRequest.newBuilder()
-                .setModel(grpcMapper.createRef(model))
-                .setDataset(grpcMapper.createRef(dataset, false))
-                .putAllColumns(features.entrySet().stream()
-                    .filter(entry -> !shouldDrop(dataset.getColumnDtypes().get(entry.getKey()), entry.getValue()))
-                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)))
-                .build();
-
-            return client.getBlockingStub().explain(request);
+    public MLWorkerWSExplainDTO explain(ProjectModel model, Dataset dataset, Map<String, String> features) {
+        MLWorkerWSExplainDTO response = null;
+        MLWorkerID workerID = model.getProject().isUsingInternalWorker() ? MLWorkerID.INTERNAL : MLWorkerID.EXTERNAL;
+        if (mlWorkerWSService.isWorkerConnected(workerID)) {
+            MLWorkerWSArtifactRefDTO modelRef = new MLWorkerWSArtifactRefDTO();
+            modelRef.setProjectKey(model.getProject().getKey());
+            modelRef.setId(model.getId().toString());
+
+            MLWorkerWSArtifactRefDTO datasetRef = new MLWorkerWSArtifactRefDTO();
+            datasetRef.setProjectKey(dataset.getProject().getKey());
+            datasetRef.setId(dataset.getId().toString());
+            datasetRef.setSample(false);
+
+            MLWorkerWSExplainParamDTO param = new MLWorkerWSExplainParamDTO();
+            param.setModel(modelRef);
+            param.setDataset(datasetRef);
+            param.setColumns(features.entrySet().stream()
+                .filter(entry -> !shouldDrop(dataset.getColumnDtypes().get(entry.getKey()), entry.getValue()))
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+
+            MLWorkerWSBaseDTO result = mlWorkerWSCommService.performAction(workerID, MLWorkerWSAction.explain, param);
+            if (result != null && result instanceof MLWorkerWSExplainDTO) {
+                response = (MLWorkerWSExplainDTO) result;
+            }
         }
+        return response;
     }
 
     public ExplainTextResponse explainText(ProjectModel model, Dataset dataset, InspectionSettings inspectionSettings, String featureName, Map<String, String> features) throws IOException {

backend/src/main/java/ai/giskard/service/ml/MLWorkerWSCommService.java

@@ -3,6 +3,7 @@
 import ai.giskard.ml.MLWorkerID;
 import ai.giskard.ml.MLWorkerWSAction;
 import ai.giskard.ml.dto.MLWorkerWSBaseDTO;
+import ai.giskard.ml.dto.MLWorkerWSExplainDTO;
 import ai.giskard.ml.dto.MLWorkerWSGetInfoDTO;
 import ai.giskard.ml.dto.MLWorkerWSRunModelForDataFrameDTO;
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -54,7 +55,7 @@ public MLWorkerWSBaseDTO performAction(MLWorkerID workerID, MLWorkerWSAction act
                 case runTestSuite -> null;
                 case runModel -> null;
                 case runModelForDataFrame -> objectMapper.readValue(result, MLWorkerWSRunModelForDataFrameDTO.class);
-                case explain -> null;
+                case explain -> objectMapper.readValue(result, MLWorkerWSExplainDTO.class);
                 case explainText -> null;
                 case echo -> null;
                 case generateTestSuite -> null;

backend/src/main/java/ai/giskard/web/rest/controllers/ModelController.java

@@ -5,6 +5,7 @@
 import ai.giskard.domain.ml.Dataset;
 import ai.giskard.domain.ml.ModelType;
 import ai.giskard.domain.ml.ProjectModel;
+import ai.giskard.ml.dto.MLWorkerWSExplainDTO;
 import ai.giskard.ml.dto.MLWorkerWSRunModelForDataFrameDTO;
 import ai.giskard.repository.ProjectRepository;
 import ai.giskard.repository.ml.DatasetRepository;
@@ -70,10 +71,12 @@ public ExplainResponseDTO explain(@PathVariable @NotNull UUID modelId, @PathVari
         ProjectModel model = modelRepository.getMandatoryById(modelId);
         permissionEvaluator.validateCanReadProject(model.getProject().getId());
         Dataset dataset = datasetRepository.getMandatoryById(datasetId);
-        ExplainResponse explanations = modelService.explain(model, dataset, data.getFeatures());
+        MLWorkerWSExplainDTO explanations = modelService.explain(model, dataset, data.getFeatures());
         ExplainResponseDTO result = new ExplainResponseDTO();
-        explanations.getExplanationsMap().forEach((label, perFeatureExplanations) ->
-            result.getExplanations().put(label, perFeatureExplanations.getPerFeatureMap()));
+        if (explanations != null) {
+            explanations.getExplanations().forEach((label, perFeatureExplanations) ->
+                result.getExplanations().put(label, perFeatureExplanations.getPerFeature()));
+        }
         return result;
     }
 

python-client/giskard/ml_worker/server/ml_worker_service.py

@@ -730,3 +730,16 @@ def runModelForDataFrame(ml_worker, params: dict, *args, **kwargs):
             "prediction": list(predictions.prediction.astype(str)),
             "raw_prediction": list(predictions.prediction),
         }
+
+
+@websocket_actor(MLWorkerAction.explain)
+def explain_ws(ml_worker, params: dict, *args, **kwargs):
+    model = BaseModel.download(ml_worker.client, params["model"]["project_key"], params["model"]["id"])
+    dataset = Dataset.download(
+        ml_worker.client, params["dataset"]["project_key"], params["dataset"]["id"], params["dataset"]["sample"]
+    )
+    explanations = explain(model, dataset, params["columns"])
+
+    logger.info(explanations)
+
+    return {"explanations": {str(k): {"per_feature": v} for k, v in explanations["explanations"].items()}}