Commit 39a9f352 authored by BO ZHANG's avatar BO ZHANG 🏀
Browse files

refactor(dag): 解耦 Brick DAG 的搜索后端与 data_model 注册层级

parent 560e34fc

README.md

+9 −1
Original line number Diff line number Diff line
@@ -100,12 +100,19 @@ DagRunGroup.trigger( 
说明：



- `data_model="raw"` 时，底层会走 level0 数据查询。

- 非 `raw` 时，会走 level1 数据查询。

- 通用 `data.find()` 对非 `raw` 数据会按 `dfs/data_model.toml` 中注册的存储层级分流到 level1 或 level2。

- `batch_id` 顶层参数表示当前新生成的 `DagRunGroup` / 输出产品所属批次。

- `data.source_batch_id` 表示检索上游 level1 数据时使用的批次条件。

- 当 `data.source_batch_id` 为空时，框架会默认回退到顶层 `batch_id`。

- 为兼容旧调用方式，若调用方仍传 `data.batch_id`，当前实现会将其视为 `source_batch_id` 的别名。



Brick / healpix 查询补充说明：



- 对 `FileBrickDAG` 来说，`data_model` 只是上游数据的业务筛选条件，不要求名字体现 `l1` / `l2`。

- `FileBrickDAG` 在生成 DagRun 前查询唯一 `healpix` 时，始终固定走 `level1.find_unique_healpix()`。

- 对 `CatalogBrickDAG` 也是同样逻辑，`data_model` 不要求体现层级，但唯一 `healpix` 查询始终固定走 `level2.find_unique_healpix()`。

- 也就是说，Brick DAG 的 unique healpix 搜索后端由 DAG 类型决定，而不是由 `data_model` 名称或注册层级决定。



#### `proc` 支持的参数



- `pmapname`
@@ -180,6 +187,7 @@ submission: 
- `submission.data.<field>` 用于配置 `data` 组字段，例如 `data_model`、`dataset`、`source_batch_id`、`instrument`、`obs_type`、`healpix`、`custom_id`。

- `submission.proc.<field>` 用于配置 `proc` 组字段，例如 `pmapname`、`ref_cat`、`extra_kwargs`。

- `data_model` 虽然属于 `data` 检索条件，但通常和 DAG 定义强绑定，因此推荐在 YAML 中显式给出。

- 对 Brick DAG 来说，这里的 `data_model` 表示“生成任务时要检索的上游数据产品”，不是“决定查询访问 level1 还是 level2 的开关”。

- `source_batch_id` 建议只在非 `raw` 的 level1 / L2 检索场景中使用；raw 数据场景通常不需要该字段。

- 大部分 L1 DAG 推荐固定为 `data_model.value: "raw"` 且 `fixed: true`。

- 某些 L2 / brick DAG 需要固定到特定 level1 模型，例如：

csst_dag/v1/dag_utils.py

+6 −6
Original line number Diff line number Diff line
@@ -42,20 +42,20 @@ def load_default_dag_run_message(): 
    return template





def get_dag_message_template(dag_id):

def get_dag_message_template(dag):

    """Get the dag_cfg message template for a given dag_cfg."""

    if dag_id not in DAG_LIST:

        raise ValueError(f"Unknown dag_cfg: {dag_id}")

    with open(os.path.join(DAG_MESSAGE_TEMPLATE_DIRECTORY, f"{dag_id}.json"), "r") as f:

    if dag not in DAG_LIST:

        raise ValueError(f"Unknown dag_cfg: {dag}")

    with open(os.path.join(DAG_MESSAGE_TEMPLATE_DIRECTORY, f"{dag}.json"), "r") as f:

        template = json.load(f)

    return template





def generate_dag_run_message(dag_id, **kwargs):

def generate_dag_run_message(dag, **kwargs):

    """Generate a DAG run message."""

    DAG_RUN_ID_DIGITS = 6

    this_dag_run_id = generate_dag_run_id(DAG_RUN_ID_DIGITS)

    msg = get_dag_message_template(dag_id)

    msg = get_dag_message_template(dag)



    for k, v in kwargs.items():

        assert msg.get(k, None) is not None, f"Unknown key: {k}"

csst_dag/v2/dag/__init__.py

+1 −1
Original line number Diff line number Diff line
@@ -60,7 +60,7 @@ CSST_DAGS = { 
    "csst-msc-l2-mbi-fphot": FileBrickDAG(

        "csst-msc-l2-mbi-fphot",

    ),

    "csst-msc-l2-mbi-photoz": FileBrickDAG(

    "csst-msc-l2-mbi-photoz": CatalogBrickDAG(

        "csst-msc-l2-mbi-photoz",

    ),

    "csst-msc-l2-mbi-xcat": CatalogBrickDAG(

csst_dag/v2/dag/brick.py

+28 −15
Original line number Diff line number Diff line
@@ -7,7 +7,6 @@ from astropy.table import Column, Table 


from ..config_loader import load_dag_config

from ..dfs import data as dfs_data

from ..dfs.model_registry import get_storage_level

from ..models import DagRun, DagRunGroup

from .base import BaseDAG, DependencyBundle
@@ -22,9 +21,10 @@ class BaseBrickDAG(BaseDAG): 
    调用 DFS 的 `find_unique_healpix()` 检索唯一 healpix 列表。

    随后以 healpix 为基本分组键生成 DagRun；若 DAG 配置中声明了

    `brick_multipliers`，则进一步生成 `healpix × multiplier` 的排列组合。

    这里选择哪个 DFS 后端由 DAG 类型决定，而不是由 `data_model`

    的字符串形式或注册层级决定。

    """



    source_level = ""

    UNIQUE_QUERY_KEYS = (

        "data_model",

        "dataset",
@@ -55,13 +55,8 @@ class BaseBrickDAG(BaseDAG): 
        data_model = str(data_query.get("data_model") or "").strip()

        if not data_model or data_model.lower() == "raw":

            raise ValueError(f"{self.__class__.__name__} 需要在 data 中显式提供非 raw 的 data_model。")



        actual_level = get_storage_level(data_model)

        if actual_level != self.source_level:

            raise ValueError(

                f"{self.__class__.__name__} 需要 `{self.source_level}` data_model，"

                f"当前 `{data_model}` 注册在 `{actual_level}`。"

            )

        # 对 Brick DAG 来说，data_model 只是上游数据筛选条件，

        # 不能再用它反推 unique_healpix 应该访问 level1 还是 level2。

        return plan_query, data_query, upstream_query



    @staticmethod
@@ -83,8 +78,8 @@ class BaseBrickDAG(BaseDAG): 
        }



    def resolve_unique_healpix(self, data_query: dict[str, Any]) -> list[str]:

        """统一通过 v2 DFS wrapper 查询唯一 healpix。"""

        return dfs_data.find_unique_healpix(**data_query)

        """按 DAG 类型选择固定的 unique_healpix 搜索后端。"""

        return self.find_unique_healpix(**data_query)



    def resolve_data_dependency(self, data_query: dict[str, Any]) -> Table:

        unique_healpix = self.resolve_unique_healpix(self._build_unique_healpix_query(data_query))
@@ -212,15 +207,33 @@ class BaseBrickDAG(BaseDAG): 




class FileBrickDAG(BaseBrickDAG):

    """基于 `level1.find_unique_healpix()` 的文件型 Brick DAG。"""

    """

    基于 `level1.find_unique_healpix()` 的文件型 Brick DAG。



    Notes

    -----

    `data_model` 可以是任意业务数据模型名，不要求名字体现 `l1`/`l2`。

    FileBrickDAG 的 unique healpix 搜索后端始终固定为 level1。

    """



    source_level = "level1"

    def find_unique_healpix(self, **data_query) -> list[str]:

        """文件型 Brick 固定通过 level1 检索唯一 healpix。"""

        return dfs_data.find_level1_unique_healpix(**data_query)





class CatalogBrickDAG(BaseBrickDAG):

    """基于 `level2.find_unique_healpix()` 的星表型 Brick DAG。"""

    """

    基于 `level2.find_unique_healpix()` 的星表型 Brick DAG。



    Notes

    -----

    `data_model` 可以是任意业务数据模型名，不要求名字体现 `l1`/`l2`。

    CatalogBrickDAG 的 unique healpix 搜索后端始终固定为 level2。

    """



    source_level = "level2"

    def find_unique_healpix(self, **data_query) -> list[str]:

        """星表型 Brick 固定通过 level2 检索唯一 healpix。"""

        return dfs_data.find_level2_unique_healpix(**data_query)





class BrickDAG(FileBrickDAG):

csst_dag/v2/dag/file.py

+17 −10
Original line number Diff line number Diff line
@@ -108,14 +108,17 @@ class FileDAG(BaseDAG): 
    def schedule(

        self,

        *,

        plan_basis: Table,

        data_basis: Table,

        deps: DependencyBundle,

        data: dict[str, Any],

        proc: dict[str, Any],

        dag_group: DagRunGroup,

        force_success: bool = False,

        return_data_list: bool = True,

        dag_group: DagRunGroup,

        proc: dict[str, Any],

    ) -> DagRunGroup:

        task_list = self.dispatcher(plan_basis, data_basis)

        if deps.plan is None or deps.data is None:

            return dag_group



        task_list = self.dispatcher(deps.plan, deps.data)

        for this_task in task_list:

            if not (force_success or this_task.get("success")):

                continue
@@ -127,7 +130,11 @@ class FileDAG(BaseDAG): 
                priority=dag_group.get("priority", "low"),

                data_model=task.get("data_model"),

                dataset=task.get("dataset"),

                source_batch_id=data.get("source_batch_id"),

                source_batch_id=(

                    str(data.get("source_batch_id"))

                    if data.get("source_batch_id") not in (None, "")

                    else None

                ),

                instrument=task.get("instrument"),

                obs_type=task.get("obs_type"),

                obs_group=task.get("obs_group"),
@@ -165,11 +172,11 @@ class FileDAG(BaseDAG): 
            return dag_group



        dag_group = self.schedule(

            plan_basis=deps.plan,

            data_basis=deps.data,

            deps=deps,

            data=data,

            proc=proc,

            dag_group=dag_group,

            force_success=bool(data.get("force_success", False)),

            return_data_list=bool(data.get("return_data_list", True)),

            dag_group=dag_group,

            proc=proc,

        )

        return dag_group