Jsonify customized metadata on management API (#3059)

* jsonify customized metadata * correct class for logger * static final logger * java lint :server:verifyJava --------- Co-authored-by: Geeta Chauhan <4461127+chauhang@users.noreply.github.com>
pytorch · Apr 4, 2024 · f9e5862 · f9e5862
1 parent 5a097f3
commit f9e5862
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/frontend/server/src/main/java/org/pytorch/serve/http/messages/DescribeModelResponse.java b/frontend/server/src/main/java/org/pytorch/serve/http/messages/DescribeModelResponse.java
@@ -1,11 +1,16 @@
 package org.pytorch.serve.http.messages;
 
+import com.google.gson.JsonObject;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Date;
 import java.util.List;
+import org.pytorch.serve.util.JsonUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class DescribeModelResponse {
+    private static final Logger logger = LoggerFactory.getLogger(DescribeModelResponse.class);
 
     private String modelName;
     private String modelVersion;
@@ -22,7 +27,7 @@ public class DescribeModelResponse {
     private List<Worker> workers;
     private Metrics metrics;
     private JobQueueStatus jobQueueStatus;
-    private String customizedMetadata;
+    private JsonObject customizedMetadata;
 
     public DescribeModelResponse() {
         workers = new ArrayList<>();
@@ -160,10 +165,16 @@ public void setJobQueueStatus(JobQueueStatus jobQueueStatus) {
     }
 
     public void setCustomizedMetadata(byte[] customizedMetadata) {
-        this.customizedMetadata = new String(customizedMetadata, Charset.forName("UTF-8"));
+        String stringMetadata = new String(customizedMetadata, Charset.forName("UTF-8"));
+        try {
+            this.customizedMetadata = JsonUtils.GSON.fromJson(stringMetadata, JsonObject.class);
+        } catch (com.google.gson.JsonSyntaxException ex) {
+            logger.warn("Customized metadata should be a dictionary.");
+            this.customizedMetadata = new JsonObject();
+        }
     }
 
-    public String getCustomizedMetadata() {
+    public JsonObject getCustomizedMetadata() {
         return customizedMetadata;
     }
 

diff --git a/frontend/server/src/test/java/org/pytorch/serve/ModelServerTest.java b/frontend/server/src/test/java/org/pytorch/serve/ModelServerTest.java
@@ -890,7 +890,7 @@ public void testNoopCustomized() throws InterruptedException {
         DescribeModelResponse[] resp =
                 JsonUtils.GSON.fromJson(TestUtils.getResult(), DescribeModelResponse[].class);
         Assert.assertEquals(
-                resp[0].getCustomizedMetadata(), "{\n  \"data1\": \"1\",\n  \"data2\": \"2\"\n}");
+                resp[0].getCustomizedMetadata().toString(), "{\"data1\":\"1\",\"data2\":\"2\"}");
 
         testUnregisterModel("noop-customized", "1.0");
     }