feat(job): custom job scheduler

.gitignore

@@ -100,3 +100,4 @@ docker/charts/starwhale-*.tgz
 
 # ansible
 bootstrap/inventory
+/client/.idea/

client/scripts/sw-docker-entrypoint

@@ -9,6 +9,7 @@ PIP_CACHE_DIR=${SW_PIP_CACHE_DIR:=/root/.cache/pip}
 _MANIFEST_RUNTIME=$(cat ${WORKDIR}/_manifest.yaml| grep "python:" | awk '{print $2}' | awk -F '.' '{print $1"."$2}') || true
 _MODEL_RUNTIME=$(cat ${WORKDIR}/model.yaml | grep 'runtime' | awk '{print $2}') || true
 VERBOSE="-vvvv"
+TASK_INDEX=${SW_TASK_INDEX:=0}
 
 _update_python_alter() {
     echo "--> set python/python3 to $1 ..."
@@ -59,11 +60,6 @@ pre_check() {
         echo "${WORKDIR} is not starwhale target dir, will exit"
         exit 1
     fi
-
-    if [ ! -f "${SW_TASK_INPUT_CONFIG}" ]; then
-        echo "${SW_TASK_INPUT_CONFIG} not found, please set env and volume file into container"
-        exit 1
-    fi
 }
 
 set_python() {
@@ -80,7 +76,7 @@ set_python() {
         _update_python_alter "python3.8"
     fi
 }
-
+# TODO:restore when processing evaluation(eval run use param: 'runtime URI', and it can be uri or dir path )
 restore_activate_runtime() {
     echo '--> restore python env ...'
     export PYTHONWARNINGS="ignore:Unverified HTTPS request"
@@ -91,14 +87,14 @@ restore_activate_runtime() {
     eval "$(./activate.sw)"
 }
 
-run_ppl() {
-    echo "--> start to run swmp ppl, use $(which swcli) cli..."
-    swcli ${VERBOSE} model ppl ${WORKDIR}/src
+run_all() {
+    echo "--> start to run swmp all, use $(which swcli) cli..."
+    swcli ${VERBOSE} model eval ${WORKDIR} --type=all --dataset=${SW_DATASET_URI} --version=${SW_EVALUATION_VERSION}
 }
 
-run_cmp() {
-    echo "--> start to run swmp cmp, use $(which swcli) cli..."
-    swcli ${VERBOSE} model cmp ${WORKDIR}/src
+run_single() {
+    echo "--> start to run swmp single, use $(which swcli) cli..."
+    swcli ${VERBOSE} model eval ${WORKDIR} --type=single --dataset=${SW_DATASET_URI} --step=${SW_TASK_STEP} --task-index=${TASK_INDEX} --version=${SW_EVALUATION_VERSION}
 }
 
 welcome() {
@@ -124,11 +120,11 @@ case "$1" in
     pre_config)
         pre_config
         ;;
-    cmp)
-        eval_task_prepare $1 && run_cmp
+    run_all)
+        eval_task_prepare $1 && run_all
         ;;
-    ppl)
-        eval_task_prepare $1 && run_ppl
+    run_single)
+        eval_task_prepare $1 && run_single
         ;;
     *)
         exec "$@"

client/starwhale/api/_impl/data_store.py

@@ -15,9 +15,11 @@
 import pyarrow as pa  # type: ignore
 import requests
 import pyarrow.parquet as pq  # type: ignore
+from loguru import logger
 from typing_extensions import Protocol
 
 from starwhale.utils.fs import ensure_dir
+from starwhale.consts.env import SWEnv
 from starwhale.utils.config import SWCliConfigMixed
 
 try:
@@ -714,7 +716,7 @@ def get_instance() -> "LocalDataStore":
 
     def __init__(self, root_path: str) -> None:
         self.root_path = root_path
-        self.name_pattern = re.compile(r"^[A-Za-z0-9-_/]+$")
+        self.name_pattern = re.compile(r"^[A-Za-z0-9-_/ ]+$")
         self.tables: Dict[str, MemoryTable] = {}
 
     def update_table(
@@ -778,6 +780,7 @@ def __init__(
                 self.columns = columns
                 self.keep_none = keep_none
 
+        logger.debug(f"scan enter, table size:{len(tables)}")
         infos: List[TableInfo] = []
         for table_desc in tables:
             table = self.tables.get(table_desc.table_name, None)
@@ -837,6 +840,7 @@ def __init__(
                         )
                     )
             else:
+                logger.debug(f"scan by disk table{info.name}")
                 iters.append(
                     _scan_table(
                         table_path,
@@ -851,7 +855,9 @@ def __init__(
             yield record
 
     def dump(self) -> None:
-        for table in self.tables.values():
+        logger.debug(f"start dump, tables size:{len(self.tables.values())}")
+        for table in list(self.tables.values()):
+            logger.debug(f"dump {table.table_name} to {self.root_path}")
             table.dump(self.root_path)
 
 
@@ -977,17 +983,19 @@ def scan_tables(
 
 
 def get_data_store() -> DataStore:
-    instance = os.getenv("SW_INSTANCE")
-    if instance is None or instance == "local":
+    instance_uri = os.getenv(SWEnv.instance_uri)
+    if instance_uri is None:
+        instance_uri = SWCliConfigMixed()._current_instance_obj["uri"]
+    if instance_uri == "local":
         return LocalDataStore.get_instance()
     else:
-        return RemoteDataStore(instance)
+        return RemoteDataStore(instance_uri)
 
 
 def _flatten(record: Dict[str, Any]) -> Dict[str, Any]:
     def _new(key_prefix: str, src: Dict[str, Any], dest: Dict[str, Any]) -> None:
         for k, v in src.items():
-            k = key_prefix + k
+            k = key_prefix + str(k)
             if type(v) is dict:
                 _new(k + "/", v, dest)
             dest[k] = v
@@ -1031,7 +1039,13 @@ def insert(self, record: Dict[str, Any]) -> None:
         record = _flatten(record)
         for k in record:
             for ch in k:
-                if not ch.isalnum() and ch != "-" and ch != "_" and ch != "/":
+                if (
+                    not ch.isalnum()
+                    and ch != "-"
+                    and ch != "_"
+                    and ch != "/"
+                    and not ch.isspace()
+                ):
                     raise RuntimeError(f"invalid field {k}")
         self._insert(record)
 

client/starwhale/api/_impl/dataset.py

@@ -126,6 +126,7 @@ def __init__(
         self.name = name
         self.version = version
         self.table_name = f"{name}/{version[:VERSION_PREFIX_CNT]}/{version}"
+        logger.debug(f"dataset table name:{self.table_name}")
         self._ds_wrapper = DatastoreWrapperDataset(self.table_name, project)
 
         self.start = start

client/starwhale/api/_impl/job.py

@@ -0,0 +1,216 @@
+import copy
+import typing as t
+from pathlib import Path
+from collections import defaultdict
+
+import yaml
+from loguru import logger
+
+from starwhale.utils import load_yaml
+from starwhale.consts import DEFAULT_EVALUATION_JOB_NAME, DEFAULT_EVALUATION_RESOURCE
+from starwhale.utils.fs import ensure_file
+from starwhale.utils.load import load_module
+
+
+def step(
+    job_name: str = DEFAULT_EVALUATION_JOB_NAME,
+    resources: t.Optional[t.List[str]] = None,
+    concurrency: int = 1,
+    task_num: int = 1,
+    needs: t.Optional[t.List[str]] = None,
+) -> t.Any:
+    _resources = resources or [
+        DEFAULT_EVALUATION_RESOURCE,
+    ]
+    _needs = needs or []
+
+    def decorator(func: t.Any) -> t.Any:
+        if Parser.is_parse_stage():
+            _step = dict(
+                job_name=job_name,
+                step_name=func.__qualname__,
+                resources=_resources,
+                concurrency=concurrency,
+                task_num=task_num,
+                needs=_needs,
+            )
+            Parser.add_job(job_name, _step)
+
+        return func
+
+    return decorator
+
+
+# Runtime concept
+class Context:
+    def __init__(
+        self,
+        workdir: Path,
+        src_dir: Path,
+        step: str = "",
+        total: int = 1,
+        index: int = 0,
+        dataset_uris: t.List[str] = [],
+        version: str = "",
+        project: str = "",
+        kw: t.Dict[str, t.Any] = {},
+    ):
+        self.project = project
+        self.version = version
+        self.step = step
+        self.total = total
+        self.index = index
+        self.dataset_uris = dataset_uris
+        self.workdir = workdir
+        self.src_dir = src_dir
+        self.kw = copy.deepcopy(kw)
+
+    def get_param(self, name: str) -> t.Any:
+        return self.kw.get(name)
+
+    def put_param(self, name: str, value: t.Any) -> None:
+        if not self.kw:
+            self.kw = {}
+        self.kw.setdefault(name, value)
+
+    def __repr__(self) -> str:
+        return "step:{}, total:{}, index:{}".format(self.step, self.total, self.index)
+
+
+class Step:
+    def __init__(
+        self,
+        job_name: str,
+        step_name: str,
+        resources: t.List[str],
+        needs: t.List[str],
+        concurrency: int = 1,
+        task_num: int = 1,
+        status: str = "",
+    ):
+        self.job_name = job_name
+        self.step_name = step_name
+        self.resources = resources
+        self.concurrency = concurrency
+        self.task_num = task_num
+        self.needs = needs
+        self.status = status
+
+    def __repr__(self) -> str:
+        return (
+            "%s(job_name=%r, step_name=%r, resources=%r, needs=%r, concurrency=%r, task_num=%r, status=%r)"
+            % (
+                self.__class__.__name__,
+                self.job_name,
+                self.step_name,
+                self.resources,
+                self.needs,
+                self.concurrency,
+                self.task_num,
+                self.status,
+            )
+        )
+
+
+class ParseConfig:
+    def __init__(self, is_parse_stage: bool, jobs: t.Dict[str, t.List[t.Dict]]):
+        self.parse_stage = is_parse_stage
+        self.jobs = jobs
+
+    def clear(self) -> None:
+        self.jobs = {}
+
+
+# shared memory, not thread safe
+parse_config = ParseConfig(False, {})
+
+
+class Parser:
+    @staticmethod
+    def set_parse_stage(parse_stage: bool) -> None:
+        parse_config.parse_stage = parse_stage
+
+    @staticmethod
+    def is_parse_stage() -> bool:
+        return parse_config.parse_stage
+
+    @staticmethod
+    def add_job(job_name: str, step: t.Dict) -> None:
+        _jobs = parse_config.jobs
+        if job_name not in _jobs:
+            parse_config.jobs[job_name] = []
+
+        parse_config.jobs[job_name].append(step)
+
+    @staticmethod
+    def get_jobs() -> t.Dict[str, t.List[t.Dict]]:
+        return parse_config.jobs
+
+    # load is unique,so don't need to think multi load and clean
+    @staticmethod
+    def clear_config() -> None:
+        global parse_config
+        parse_config.clear()
+
+    @staticmethod
+    def parse_job_from_module(module: str, path: Path) -> t.Dict[str, t.List[t.Dict]]:
+        """
+        parse @step from module
+        :param module: module name
+        :param path: abs path
+        :return: jobs
+        """
+        Parser.set_parse_stage(True)
+        # parse DAG
+        logger.debug(f"parse @step for module:{module}")
+        load_module(module, path)
+        _jobs = Parser.get_jobs().copy()
+        Parser.clear_config()
+        return _jobs
+
+    @staticmethod
+    def generate_job_yaml(module: str, path: Path, target_file: Path) -> None:
+        """
+        generate job yaml
+        :param target_file: yaml target path
+        :param module: module name
+        :param path: abs path
+        :return: None
+        """
+        _jobs = Parser.parse_job_from_module(module, path)
+        # generate DAG
+        logger.debug("generate DAG")
+        if Parser.check(_jobs):
+            # dump to target
+            ensure_file(target_file, yaml.safe_dump(_jobs, default_flow_style=False))
+            logger.debug("generator DAG success!")
+        else:
+            logger.error("generator DAG error! reason: check is failed.")
+
+    @staticmethod
+    def check(jobs: t.Dict[str, t.List[t.Dict]]) -> bool:
+        checks = []
+        logger.debug(f"jobs:{jobs}")
+        for job in jobs.items():
+            all_steps = []
+            needs = []
+            for _step in job[1]:
+                all_steps.append(_step["step_name"])
+                for d in _step["needs"]:
+                    if d:
+                        needs.append(d)
+            logger.debug(f"all steps:{all_steps}, length:{len(all_steps)}")
+            _check = all(item in all_steps for item in needs)
+            if not _check:
+                logger.error(f"job:{job[0]} check error!")
+            checks.append(_check)
+
+        return all(checks)
+
+    @staticmethod
+    def parse_job_from_yaml(file_path: str) -> t.Dict[str, t.List[Step]]:
+        _jobs = load_yaml(file_path)
+        rt = defaultdict(list)
+        for k, v in _jobs.items():
+            rt[k] = [Step(**_v) for _v in v]
+        return rt