chore(controller): add draft field to model version vo (star-whale#2982)

mfq1020 · Nov 16, 2023 · 0b768e5 · 0b768e5
1 parent 855b540
commit 0b768e5
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 1 deletion.
diff --git a/client/starwhale/base/client/models/models.py b/client/starwhale/base/client/models/models.py
@@ -793,6 +793,7 @@ class ModelVersionVo(SwBaseModel):
     owner: Optional[UserVo] = None
     shared: bool
     built_in_runtime: Optional[str] = Field(None, alias='builtInRuntime')
+    draft: bool
 
 
 class ModelVo(SwBaseModel):

diff --git a/client/tests/core/test_model.py b/client/tests/core/test_model.py
@@ -775,6 +775,7 @@ def test_list_with_project(self, req: Mocker, mock_list: MagicMock):
                 alias="v2",
                 created_time=5,
                 shared=False,
+                draft=False,
             ),
         )
         mock_list.return_value = ([model_foo, model_foo], {})
@@ -1200,6 +1201,7 @@ def test_run(self, rm: Mocker) -> None:
                         "createdTime": 1697042331000,
                         "owner": None,
                         "shared": True,
+                        "draft": False,
                         "stepSpecs": [
                             {
                                 "name": "mnist.evaluator:MNISTInference.predict",
@@ -1247,6 +1249,7 @@ def test_run(self, rm: Mocker) -> None:
                     "versionTag": "t1",
                     "createdTime": 1697042331000,
                     "shared": 1,
+                    "draft": False,
                     "id": "1",
                 }
             },

diff --git a/client/tests/data/job/job_list_resp.json b/client/tests/data/job/job_list_resp.json
@@ -66,6 +66,7 @@
             "createdTime": 1685593961000,
             "owner": null,
             "shared": 0,
+            "draft": false,
             "builtInRuntime": null
           }
         },

diff --git a/client/tests/sdk/test_job_sdk.py b/client/tests/sdk/test_job_sdk.py
@@ -162,6 +162,7 @@ def test_list_jobs_from_server(
                     alias="v2",
                     created_time=789,
                     shared=False,
+                    draft=False,
                 ),
             ),
             runtime=RuntimeVo(
@@ -286,6 +287,7 @@ def test_create_remote_job(self, rm: Mocker, mock_load_conf: MagicMock) -> None:
                         "createdTime": 1697042331000,
                         "owner": None,
                         "shared": True,
+                        "draft": False,
                         "stepSpecs": [
                             {
                                 "name": "mnist.evaluator:MNISTInference.predict",
@@ -333,6 +335,7 @@ def test_create_remote_job(self, rm: Mocker, mock_load_conf: MagicMock) -> None:
                     "versionTag": "t1",
                     "createdTime": 1697042331000,
                     "shared": 1,
+                    "draft": False,
                     "id": "1",
                 }
             },

diff --git a/server/controller/src/main/java/ai/starwhale/mlops/api/protocol/model/ModelVersionVo.java b/server/controller/src/main/java/ai/starwhale/mlops/api/protocol/model/ModelVersionVo.java
@@ -70,7 +70,11 @@ public class ModelVersionVo implements Serializable {
     @NotNull
     private List<StepSpec> stepSpecs;
 
+    @NotNull
+    @JsonProperty("draft")
+    private Boolean draft;
+
     public static ModelVersionVo empty() {
-        return new ModelVersionVo("", "", "", false, List.of(), 0L, -1L, UserVo.empty(), false, "", List.of());
+        return new ModelVersionVo("", "", "", false, List.of(), 0L, -1L, UserVo.empty(), false, "", List.of(), false);
     }
 }
diff --git a/...ller/src/main/java/ai/starwhale/mlops/domain/model/converter/ModelVersionVoConverter.java b/...ller/src/main/java/ai/starwhale/mlops/domain/model/converter/ModelVersionVoConverter.java
@@ -60,6 +60,7 @@ public ModelVersionVo convert(
                     .stepSpecs(jobSpecParser.parseAndFlattenStepFromYaml(entity.getJobs()))
                     .size(entity.getStorageSize())
                     .shared(entity.getShared())
+                    .draft(entity.getDraft())
                     .build();
         } catch (JsonProcessingException e) {
             log.error("convert ModelVersionVo error", e);