WIP enhancement(controller): add type for step spec

client/starwhale/api/_impl/evaluation/__init__.py

@@ -4,6 +4,7 @@
 
 from starwhale.consts import DecoratorInjectAttr
 from starwhale.base.type import PredictLogMode
+from starwhale.base.client.models.models import StepType
 
 from .log import Evaluation
 from .pipeline import PipelineHandler
@@ -109,6 +110,7 @@ def _register_predict(
             dataset_uris=datasets,
         ),
         built_in=True,
+        typ=StepType.evaluation,
     )(func)
 
 
@@ -173,6 +175,7 @@ def _register_evaluate(
             predict_auto_log=use_predict_auto_log,
         ),
         built_in=True,
+        typ=StepType.evaluation,
     )(func)
 
 

client/starwhale/api/_impl/experiment.py

@@ -11,7 +11,7 @@
 from starwhale.base.context import Context
 from starwhale.api._impl.model import build as build_starwhale_model
 from starwhale.api._impl.dataset import Dataset
-from starwhale.base.client.models.models import FineTune
+from starwhale.base.client.models.models import FineTune, StepType
 
 
 # TODO: support arguments
@@ -140,5 +140,6 @@ def _register_ft(
             require_train_datasets=require_train_datasets,
             require_validation_datasets=require_validation_datasets,
         ),
+        typ=StepType.fine_tune,
     )(func)
     setattr(func, DecoratorInjectAttr.FineTune, True)

client/starwhale/api/_impl/job/handler.py

@@ -20,6 +20,7 @@
 from starwhale.api._impl.evaluation import PipelineHandler
 from starwhale.base.client.models.models import (
     FineTune,
+    StepType,
     RuntimeResource,
     ParameterSignature,
 )
@@ -105,6 +106,7 @@ def register(
         require_dataset: bool = False,
         built_in: bool = False,
         fine_tune: FineTune | None = None,
+        typ: StepType | None = None,
     ) -> t.Callable:
         """Register a function as a handler. Enable the function execute by needs handler, run with gpu/cpu/mem resources in server side,
         and control replicas of handler run.
@@ -125,6 +127,8 @@ def register(
             built_in: [bool, optional] A special flag to distinguish user defined args in handler function from the StarWhale ones.
               This should always be False unless you know what it does.
             fine_tune: [FineTune, optional The fine tune config for the handler. Default is None.
+            typ:
+              - [str, optional] The type of the handler. Default is None.
 
         Example:
         ```python
@@ -202,6 +206,7 @@ def decorator(func: t.Callable) -> t.Callable:
                 parameters_sig=parameters_sig,
                 ext_cmd_args=ext_cmd_args,
                 fine_tune=fine_tune,
+                step_type=typ,
             )
 
             cls._register(_handler, func)
@@ -347,6 +352,7 @@ def _preload_registering_handlers(
                         name="predict",
                         require_dataset=True,
                         built_in=True,
+                        typ=StepType.evaluation,
                     )
 
                     evaluate_register = partial(
@@ -355,6 +361,7 @@ def _preload_registering_handlers(
                         name="evaluate",
                         built_in=True,
                         replicas=1,
+                        typ=StepType.evaluation,
                     )
 
                     run_info = getattr(_cls, "_registered_run_info", None)

client/starwhale/base/client/models/models.py

@@ -779,6 +779,12 @@ class RuntimeResource(SwBaseModel):
     limit: Optional[float] = None
 
 
+class StepType(Enum):
+    evaluation = 'EVALUATION'
+    fine_tune = 'FINE_TUNE'
+    serving = 'SERVING'
+
+
 class DatasetVersionViewVo(SwBaseModel):
     id: str
     version_name: str = Field(..., alias='versionName')
@@ -1492,6 +1498,7 @@ class StepSpec(SwBaseModel):
     virtual: Optional[bool] = None
     job_name: Optional[str] = None
     show_name: str
+    step_type: Optional[StepType] = None
     require_dataset: Optional[bool] = None
     fine_tune: Optional[FineTune] = None
     container_spec: Optional[ContainerSpec] = None

client/starwhale/core/model/model.py

@@ -86,7 +86,7 @@
 from starwhale.core.runtime.model import StandaloneRuntime
 from starwhale.base.client.api.job import JobApi
 from starwhale.base.client.api.model import ModelApi
-from starwhale.base.client.models.models import JobRequest, ModelInfoVo
+from starwhale.base.client.models.models import StepType, JobRequest, ModelInfoVo
 
 
 @unique
@@ -271,6 +271,7 @@ def _gen_model_serving(self, search_modules: t.List[str], workdir: Path) -> None
         cls_name, _, func_name = func.__qualname__.rpartition(".")
         h = StepSpecClient(
             name="serving",
+            step_type=StepType.serving,
             show_name="virtual handler for model serving",
             func_name=func.__qualname__,
             module_name=func.__module__,

client/starwhale/utils/json.py

@@ -1,5 +1,6 @@
 import json
 import dataclasses
+from enum import Enum
 from typing import Any
 
 from pydantic import BaseModel
@@ -11,6 +12,8 @@ def default(self, o: Any) -> Any:
             return o.decode("utf-8")
         if isinstance(o, BaseModel):
             return json.loads(o.json(exclude_unset=True))
+        if isinstance(o, Enum):
+            return o.value
         if dataclasses.is_dataclass(o):
             return dataclasses.asdict(o)
         return super().default(o)

client/tests/core/test_model.py

@@ -62,6 +62,7 @@
 from starwhale.base.client.models.models import (
     UserVo,
     ModelVo,
+    StepType,
     JobRequest,
     ModelVersionVo,
     ResponseMessageString,
@@ -120,6 +121,7 @@ def test_build_workflow(
             show_name="t1",
             func_name="test_func1",
             module_name="mock",
+            step_type="EVALUATION",
         )
         Handler._registered_handlers["depend"] = StepSpecClient(
             name="depend",
@@ -191,6 +193,7 @@ def test_build_workflow(
                     func_name="test_func1",
                     module_name="mock",
                     replicas=1,
+                    step_type=StepType.evaluation,
                 )
             ],
             "depend": [
@@ -200,6 +203,7 @@ def test_build_workflow(
                     func_name="test_func1",
                     module_name="mock",
                     replicas=1,
+                    step_type=StepType.evaluation,
                 ),
                 StepSpecClient(
                     name="depend",
@@ -208,11 +212,13 @@ def test_build_workflow(
                     module_name="mock",
                     replicas=1,
                     needs=["base"],
+                    step_type=None,
                 ),
             ],
             "serving": [
                 StepSpecClient(
                     name="serving",
+                    step_type=StepType.serving,
                     show_name="virtual handler for model serving",
                     func_name="StandaloneModel._serve_handler",
                     module_name="starwhale.core.model.model",

client/tests/sdk/test_job_handler.py

@@ -15,6 +15,7 @@
 from starwhale.base.models.model import JobHandlers, StepSpecClient
 from starwhale.base.client.models.models import (
     FineTune,
+    StepType,
     RuntimeResource,
     ParameterSignature,
 )
@@ -242,6 +243,7 @@ def video_evaluate_handler(*args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
             StepSpecClient(
                 cls_name="",
@@ -257,6 +259,7 @@ def video_evaluate_handler(*args, **kwargs): ...
                 require_dataset=False,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
         ]
 
@@ -282,6 +285,7 @@ def video_evaluate_handler(*args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
             StepSpecClient(
                 cls_name="",
@@ -297,6 +301,7 @@ def video_evaluate_handler(*args, **kwargs): ...
                 require_dataset=False,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
         ]
 
@@ -346,6 +351,7 @@ def evaluate_handler(*args, **kwargs): ...
                     require_dataset=True,
                     parameters_sig=[],
                     ext_cmd_args="",
+                    step_type=StepType.evaluation,
                 ),
                 StepSpecClient(
                     cls_name="",
@@ -361,6 +367,7 @@ def evaluate_handler(*args, **kwargs): ...
                     require_dataset=False,
                     parameters_sig=[],
                     ext_cmd_args="",
+                    step_type=StepType.evaluation,
                 ),
             ],
             "mock_user_module:predict_handler": [
@@ -385,6 +392,7 @@ def evaluate_handler(*args, **kwargs): ...
                     require_dataset=True,
                     parameters_sig=[],
                     ext_cmd_args="",
+                    step_type=StepType.evaluation,
                 )
             ],
         }
@@ -511,6 +519,7 @@ def evaluate(self, *args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
             StepSpecClient(
                 cls_name="MockHandler",
@@ -528,6 +537,7 @@ def evaluate(self, *args, **kwargs): ...
                 require_dataset=False,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
         ]
         assert jobs_info["mock_user_module:MockHandler.predict"] == [
@@ -546,6 +556,7 @@ def evaluate(self, *args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             )
         ]
         _, steps = Step.get_steps_from_yaml(
@@ -614,6 +625,7 @@ def evaluate_handler(self, *args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             )
         ]
 
@@ -639,6 +651,7 @@ def evaluate_handler(self, *args, **kwargs): ...
                 require_dataset=True,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
             StepSpecClient(
                 cls_name="MockHandler",
@@ -654,6 +667,7 @@ def evaluate_handler(self, *args, **kwargs): ...
                 require_dataset=False,
                 parameters_sig=[],
                 ext_cmd_args="",
+                step_type=StepType.evaluation,
             ),
         ]
 
@@ -793,6 +807,7 @@ def ft2(): ...
                         require_train_datasets=True,
                         require_validation_datasets=True,
                     ),
+                    step_type=StepType.fine_tune,
                 ),
             ],
             "mock_user_module:ft2": [
@@ -816,6 +831,7 @@ def ft2(): ...
                         require_train_datasets=True,
                         require_validation_datasets=True,
                     ),
+                    step_type=StepType.fine_tune,
                 ),
                 StepSpecClient(
                     name="mock_user_module:ft2",
@@ -843,6 +859,7 @@ def ft2(): ...
                         require_train_datasets=False,
                         require_validation_datasets=False,
                     ),
+                    step_type=StepType.fine_tune,
                 ),
             ],
         }

console/src/api/server/data-contracts.ts

@@ -1305,6 +1305,7 @@ export interface IStepSpec {
     virtual?: boolean
     job_name?: string
     show_name: string
+    step_type?: 'EVALUATION' | 'FINE_TUNE' | 'SERVING'
     require_dataset?: boolean
     fine_tune?: IFineTune
     container_spec?: IContainerSpec

server/controller/src/main/java/ai/starwhale/mlops/api/JobController.java

@@ -234,7 +234,7 @@ public ResponseEntity<ResponseMessage<String>> createJob(
         } else {
             jobId = jobServiceForWeb.createJob(userJobConverter.convert(projectUrl, jobRequest));
         }
-        if (jobId == null)  {
+        if (jobId == null) {
             throw new StarwhaleApiException(
                     new SwValidationException(ValidSubject.JOB, "job request is invalid"),
                     HttpStatus.BAD_REQUEST

server/controller/src/main/java/ai/starwhale/mlops/api/protocol/model/ModelVersionViewVo.java

@@ -16,7 +16,7 @@
 
 package ai.starwhale.mlops.api.protocol.model;
 
-import ai.starwhale.mlops.domain.job.spec.StepSpec;
+import ai.starwhale.mlops.domain.job.spec.step.StepSpec;
 import io.swagger.v3.oas.annotations.media.Schema;
 import java.util.List;
 import javax.validation.constraints.NotNull;

server/controller/src/main/java/ai/starwhale/mlops/api/protocol/model/ModelVersionVo.java

@@ -17,7 +17,7 @@
 package ai.starwhale.mlops.api.protocol.model;
 
 import ai.starwhale.mlops.api.protocol.user.UserVo;
-import ai.starwhale.mlops.domain.job.spec.StepSpec;
+import ai.starwhale.mlops.domain.job.spec.step.StepSpec;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import io.swagger.v3.oas.annotations.media.Schema;
 import java.io.Serializable;

server/controller/src/main/java/ai/starwhale/mlops/domain/dataset/DatasetService.java

@@ -56,9 +56,9 @@
 import ai.starwhale.mlops.domain.job.JobType;
 import ai.starwhale.mlops.domain.job.bo.Job;
 import ai.starwhale.mlops.domain.job.bo.JobCreateRequest;
-import ai.starwhale.mlops.domain.job.spec.Env;
 import ai.starwhale.mlops.domain.job.spec.JobSpecParser;
-import ai.starwhale.mlops.domain.job.spec.StepSpec;
+import ai.starwhale.mlops.domain.job.spec.step.Env;
+import ai.starwhale.mlops.domain.job.spec.step.StepSpec;
 import ai.starwhale.mlops.domain.job.step.VirtualJobLoader;
 import ai.starwhale.mlops.domain.project.ProjectService;
 import ai.starwhale.mlops.domain.project.bo.Project;

server/controller/src/main/java/ai/starwhale/mlops/domain/ft/FineTuneAppService.java

@@ -44,9 +44,9 @@
 import ai.starwhale.mlops.domain.job.converter.UserJobConverter;
 import ai.starwhale.mlops.domain.job.mapper.JobMapper;
 import ai.starwhale.mlops.domain.job.po.JobEntity;
-import ai.starwhale.mlops.domain.job.spec.Env;
 import ai.starwhale.mlops.domain.job.spec.JobSpecParser;
-import ai.starwhale.mlops.domain.job.spec.StepSpec;
+import ai.starwhale.mlops.domain.job.spec.step.Env;
+import ai.starwhale.mlops.domain.job.spec.step.StepSpec;
 import ai.starwhale.mlops.domain.job.status.JobStatusMachine;
 import ai.starwhale.mlops.domain.model.ModelDao;
 import ai.starwhale.mlops.domain.model.ModelService;

server/controller/src/main/java/ai/starwhale/mlops/domain/job/ModelServingService.java

@@ -25,10 +25,10 @@
 import ai.starwhale.mlops.domain.job.bo.UserJobCreateRequest;
 import ai.starwhale.mlops.domain.job.mapper.ModelServingMapper;
 import ai.starwhale.mlops.domain.job.po.ModelServingEntity;
-import ai.starwhale.mlops.domain.job.spec.Env;
 import ai.starwhale.mlops.domain.job.spec.JobSpecParser;
 import ai.starwhale.mlops.domain.job.spec.ModelServingSpec;
-import ai.starwhale.mlops.domain.job.spec.StepSpec;
+import ai.starwhale.mlops.domain.job.spec.step.Env;
+import ai.starwhale.mlops.domain.job.spec.step.StepSpec;
 import ai.starwhale.mlops.domain.job.status.JobStatus;
 import ai.starwhale.mlops.domain.job.step.VirtualJobLoader;
 import ai.starwhale.mlops.domain.model.ModelDao;