Commit 609407de authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

feat: 简化任务提交前端并重构优先级逻辑

parent facd6a99

README.md

+3 −3
Original line number Diff line number Diff line
@@ -189,9 +189,9 @@ Worker 节点通常不暴露核心业务端口，仅暴露监控相关端口供 


推荐业务方直接对接 API Gateway（避免 Airflow 侧鉴权变化带来的耦合，同时提供 JSONB 查询与批量触发能力）。接口与示例见 [api.md](file:///home/cham/PycharmProjects/csst-airflow/api.md)。



- 批量提交接口支持可选 `priority=high|normal|low`，默认 `normal`。

- Gateway 会将优先级映射到实际调度 DAG：`<dag_id>`、`<dag_id>__high`、`<dag_id>__low`。

- 当前主要工厂 DAG 已提供 `high/low` 变体，并为任务设置 `priority_weight` 与 `weight_rule=absolute`。

- 批量提交接口支持可选 `priority=high|low`，默认 `low`。

- Gateway 会将优先级映射到实际调度 DAG：`<dag_id>`、`<dag_id>__high`。

- 当前主要工厂 DAG 已提供 `high` 变体，并为任务设置 `priority_weight` 与 `weight_rule=absolute`。



## 相关文档

docker-celery-3.0.5/api_gateway/README.md

+1 −1
Original line number Diff line number Diff line
@@ -7,7 +7,7 @@ 
2. **高性能批量触发**: 支持瞬间接收 100,000+ 个任务，存入 Postgres 后通过后台异步推送给 Airflow。

3. **状态与日志聚合**: 结合 Airflow Callback，任务的状态、输入、输出和简短日志全部聚合在同一张 Postgres 表中。

4. **异构 JSON 查询**: 利用 PostgreSQL 的 JSONB 和 GIN 索引，支持对任何异构的输入参数字段进行快速反向查询。

5. **提交时优先级选择**: `POST /api/tasks/batch-submit` 支持可选 `priority=high|normal|low`，默认 `normal`；网关会将其映射到实际调度 DAG：`<dag_id>`、`<dag_id>__high`、`<dag_id>__low`。

5. **提交时优先级选择**: `POST /api/tasks/batch-submit` 支持可选 `priority=high|low`，默认 `low`；网关会将其映射到实际调度 DAG：`<dag_id>`、`<dag_id>__high`。

6. **请求 DAG 与调度 DAG 分离**: 查询接口保留 `dag_id`/`requested_dag_id` 表示原始请求 DAG，额外返回 `scheduled_dag_id` 表示最终实际触发的 DAG 变体。

7. **Harbor 镜像查询**: 支持仓库名模糊搜索与 tag 列表查询（用于前端/工具做版本选择）。

docker-celery-3.0.5/api_gateway/main.py

+16 −12
Original line number Diff line number Diff line
@@ -77,9 +77,10 @@ AIRFLOW_USER = os.getenv("AIRFLOW_USER", "airflow") 
AIRFLOW_PASS = os.getenv("AIRFLOW_PASS", "airflow")

ELASTICSEARCH_URL = os.getenv("ELASTICSEARCH_URL", "http://elasticsearch:9200")



SUPPORTED_PRIORITIES = {"high", "normal", "low"}

SUPPORTED_PRIORITIES = {"high", "low"}

PRIORITY_ALIASES = {

    "medium": "normal",

    "medium": "low",

    "normal": "low",

}



def extract_worker_label(worker_name: str) -> str:
@@ -115,14 +116,13 @@ def load_worker_inventory_hosts() -> Dict[str, str]: 
    return hosts

PRIORITY_DAG_SUFFIXES = {

    "high": "__high",

    "normal": "",

    "low": "__low",

    "low": "",

}



class TaskRequest(BaseModel):

    task_id: str = Field(..., description="业务方提供的唯一任务 ID，将作为 Airflow 的 dag_run_id")

    dag_id: str

    priority: str = Field("normal", description="可选优先级：high | normal | low")

    priority: str = Field("low", description="可选优先级：high | low")

    inputs: Dict[str, Any]



class BatchSubmitResponse(BaseModel):
@@ -150,13 +150,17 @@ def submit_to_airflow(task_id: str, dag_id: str, inputs: dict): 
        # In Airflow 3 REST API v2, the logical_date must be unique

        logical_date = datetime.now(timezone.utc).isoformat(timespec="microseconds").replace("+00:00", "Z")

        

        conf = {

            "task_id": str(task_id),

            **(inputs or {}),

        }

        if "input" not in conf:

            conf["input"] = json.dumps(conf, ensure_ascii=False, separators=(",", ":"))



        payload = {

            "dag_run_id": str(task_id),

            "logical_date": logical_date,

            "conf": {

                "task_id": str(task_id),

                **inputs

            }

            "conf": conf,

        }

        

        response = requests.post(
@@ -183,12 +187,12 @@ def submit_to_airflow(task_id: str, dag_id: str, inputs: dict): 




def normalize_priority(priority: Optional[str]) -> str:

    normalized = str(priority or "normal").strip().lower()

    normalized = str(priority or "low").strip().lower()

    normalized = PRIORITY_ALIASES.get(normalized, normalized)

    if normalized not in SUPPORTED_PRIORITIES:

        raise HTTPException(

            status_code=400,

            detail="priority must be one of: high, normal, low",

            detail="priority must be one of: high, low",

        )

    return normalized
@@ -197,7 +201,7 @@ def split_priority_variant_dag_id(dag_id: str) -> Tuple[str, str]: 
    for priority, suffix in PRIORITY_DAG_SUFFIXES.items():

        if suffix and dag_id.endswith(suffix):

            return dag_id[: -len(suffix)], priority

    return dag_id, "normal"

    return dag_id, "low"





def resolve_scheduled_dag_id(dag_id: str, priority: Optional[str]) -> Tuple[str, str, str]:

docker-celery-3.0.5/dags/common/dag_factory.py

+120 −9
Original line number Diff line number Diff line
@@ -15,13 +15,11 @@ logger = logging.getLogger(__name__) 
HARBOR = os.environ.get("HARBOR", "csu-harbor.csst.nao:10443")

PRIORITY_WEIGHTS = {

    "high": 100,

    "normal": 10,

    "low": 1,

}

PRIORITY_DAG_SUFFIXES = {

    "high": "__high",

    "normal": "",

    "low": "__low",

    "low": "",

}



def stringify_runtime_env_value(value: object) -> str:
@@ -78,12 +76,18 @@ class RuntimeEnvDockerOperator(DockerOperator): 
    """



    def __init__(self, *args, **kwargs):

        self._pipeline_image_name = str(kwargs.pop("pipeline_image_name", "") or "").strip()

        self._static_environment = dict(kwargs.get("environment") or {})

        super().__init__(*args, **kwargs)



    def execute(self, context):

        runtime_env = load_dag_runtime_env()

        self.environment = {**self._static_environment, **runtime_env}



        if self._pipeline_image_name:

            resolved_tag = resolve_pipeline_image_tag(context, self._pipeline_image_name)

            self.image = f"{HARBOR}/csst/{self._pipeline_image_name}:{resolved_tag}"



        logger.info(

            "Resolved %s runtime env vars for task %s at execution time",

            len(runtime_env),
@@ -92,6 +96,66 @@ class RuntimeEnvDockerOperator(DockerOperator): 
        return super().execute(context)





def _extract_str(value: object) -> str:

    if value is None:

        return ""

    if isinstance(value, str):

        return value.strip()

    return str(value).strip()





def _load_json_object(raw: object) -> dict:

    if raw is None:

        return {}

    if isinstance(raw, dict):

        return raw

    if isinstance(raw, str):

        try:

            parsed = json.loads(raw)

        except Exception:

            return {}

        return parsed if isinstance(parsed, dict) else {}

    return {}





def resolve_pipeline_image_tag(context, image_name: str) -> str:

    params = context.get("params") or {}



    docker_images = params.get("docker_images")

    if isinstance(docker_images, dict):

        mapped = _extract_str(docker_images.get(image_name))

        if mapped:

            return mapped



    version = _extract_str(params.get("version") or params.get("image_tag") or params.get("tag"))

    if version:

        return version



    params_input = params.get("input")

    payload = _load_json_object(params_input)

    if payload:

        docker_images = payload.get("docker_images")

        if isinstance(docker_images, dict):

            mapped = _extract_str(docker_images.get(image_name))

            if mapped:

                return mapped



        version = _extract_str(payload.get("version") or payload.get("image_tag") or payload.get("tag"))

        if version:

            return version



    dag_run = context.get("dag_run")

    conf = getattr(dag_run, "conf", None) if dag_run is not None else None

    conf_obj = _load_json_object(conf)

    docker_images = conf_obj.get("docker_images")

    if isinstance(docker_images, dict):

        mapped = _extract_str(docker_images.get(image_name))

        if mapped:

            return mapped



    return "latest"





def sync_status_to_db(context, status, error_msg=None):

    """

    统一状态同步回调函数，更新 csst_task_records 表
@@ -181,7 +245,7 @@ def create_csst_pipeline_dag( 
    default_params_input: str,

    output_subpath: str = "",

    log_tasks: list = None,

    priority_weight: int = PRIORITY_WEIGHTS["normal"],

    priority_weight: int = PRIORITY_WEIGHTS["low"],

    weight_rule: str = "absolute",

) -> DAG:

    """
@@ -243,8 +307,8 @@ def create_csst_pipeline_dag( 
        # 2. 运行 Docker 容器

        run_pipeline = RuntimeEnvDockerOperator(

            task_id=f"run_{dag_id.replace('-', '_')}",

            # 动态提取 tag，如果没有获取到则默认是 'latest'

            image=f"{HARBOR}/csst/{image_name}:{{{{ dag_run.conf.get('docker_images', {{}}).get('{image_name}', 'latest') }}}}",

            pipeline_image_name=image_name,

            image=f"{HARBOR}/csst/{image_name}:latest",

            command="run {{ params.input }}",

            docker_url="unix://var/run/docker.sock",

            network_mode="bridge",
@@ -287,7 +351,7 @@ def create_csst_pipeline_dag_variants( 
    default_params_input: str,

    output_subpath: str = "",

    log_tasks: list = None,

    include_priorities: tuple[str, ...] = ("normal", "high", "low"),

    include_priorities: tuple[str, ...] = ("low", "high"),

) -> dict[str, DAG]:

    dags = {}
@@ -295,12 +359,12 @@ def create_csst_pipeline_dag_variants( 
        dag_suffix = PRIORITY_DAG_SUFFIXES[priority]

        variant_dag_id = f"{dag_id}{dag_suffix}"

        variant_tags = list(tags)

        if priority != "normal":

        if priority != "low":

            variant_tags.append(priority)



        dags[variant_dag_id] = create_csst_pipeline_dag(

            dag_id=variant_dag_id,

            description=description if priority == "normal" else f"{description} ({priority})",

            description=description if priority == "low" else f"{description} ({priority})",

            tags=variant_tags,

            image_name=image_name,

            default_params_input=build_priority_default_params_input(
@@ -315,3 +379,50 @@ def create_csst_pipeline_dag_variants( 
        )



    return dags





def create_single_docker_dag(

    dag_id: str,

    description: str,

    tags: list,

    image_name: str,

    default_params_input: str,

    schedule: str = "@once",

) -> DAG:

    args = COMMON_DEFAULT_ARGS.copy()



    try:

        parsed_input = json.loads(default_params_input)

    except Exception:

        parsed_input = {}



    args["params"] = {

        "input": default_params_input,

        **(parsed_input if isinstance(parsed_input, dict) else {}),

    }



    with DAG(

        dag_id=dag_id,

        default_args=args,

        description=description,

        start_date=datetime(2023, 1, 1),

        schedule=schedule,

        catchup=False,

        tags=tags,

        on_success_callback=on_success_callback,

        on_failure_callback=on_failure_callback,

    ) as dag:

        RuntimeEnvDockerOperator(

            task_id=f"run_{dag_id.replace('-', '_')}",

            pipeline_image_name=image_name,

            image=f"{HARBOR}/csst/{image_name}:latest",

            command="run {{ params.input }}",

            docker_url="unix://var/run/docker.sock",

            network_mode="bridge",

            auto_remove="force",

            force_pull=True,

            mount_tmp_dir=False,

            environment={},

        )



    return dag

docker-celery-3.0.5/dags/csst-echo.py

0 → 100644
+16 −0
Original line number Diff line number Diff line 
import sys

import os



sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from common.dag_factory import create_single_docker_dag



PARAMS_INPUT = '{"A":1,"B":2}'



dag = create_single_docker_dag(

    dag_id="csst-echo",

    description="Echo test DAG",

    tags=["csst", "echo", "test"],

    image_name="csst-echo",

    default_params_input=PARAMS_INPUT,

)