jason810496 commented on code in PR #75190:
URL: https://github.com/apache/airflow/pull/75190#discussion_r4334504089
##########
ts-sdk/src/sdk/dag-graph.ts:
##########
@@ -149,13 +155,18 @@ export function buildDagGraph(dag: Dag): DagGraph {
return sets;
};
- for (const { upstream, downstream } of orderEdges) {
+ for (const { upstream, downstream, label } of orderEdges) {
const upstreamIsGroup = groups.has(upstream);
const downstreamIsGroup = groups.has(downstream);
if (!upstreamIsGroup && !downstreamIsGroup) continue;
const from = tasksAt(upstream, "upstream");
const to = tasksAt(downstream, "downstream");
+ if (!upstreamIsGroup && label !== undefined) {
Review Comment:
This only covers the direction where the upstream is a task. The
group-to-task direction gets no task edge labels at all.
```ts
// x and y are in g
reported.after(Label.of("L").on(g));
// edge_info: { "g.downstream_join_id": { report: "L" } }
```
Python writes the join entry plus `g.x -> report` and `g.y -> report`, both
labelled `"L"`. In `DAGNode._set_relatives` (task-sdk `_internal/node.py`),
`relatives = task_object.leaves` picks up the group's leaves and the loop calls
`edge_modifier.add_edge_info` for each of them, the same way it does for roots
in the task-to-group direction.
Would it be better to handle the mirror case here with
`ends.leaves(upstream)`?
##########
ts-sdk/src/sdk/dag.ts:
##########
@@ -327,48 +327,105 @@ export function isTaskRef(value: unknown): value is
TaskRef {
const conditionTasks = new WeakMap<object, TaskRef>();
-const labelledNodes = new WeakMap<object, { readonly node: Node; readonly
text: string }>();
+const labelledNodes = new WeakMap<
+ object,
+ { readonly nodes: readonly Node[]; readonly text: string }
+>();
function resolveNode(node: Node): Node {
if (typeof node !== "object" || node === null) return node;
- return labelledNodes.get(node)?.node ?? conditionTasks.get(node) ?? node;
+ return conditionTasks.get(node) ?? node;
}
/**
- * Label the edge that `before` or `after` draws to `node`, as Python's
`Label` does.
- * The result only goes into `before` or `after`; calling either on it throws.
- * `text` cannot be empty.
+ * A text label for an edge. The Airflow UI shows the label on the edge in the
graph, as it does
+ * for Python's `Label`. The text comes before the task that the label is put
on. That way the
+ * text sits between the two tasks of the edge in the code, just as the label
sits between them in
+ * the graph:
*
* ```ts
- * checked.before(label(processed, "rows found"), label(emptyNotice, "no
rows"));
+ * // load >> Label("when empty") >> report_empty
+ * loaded.before(Label.of("when empty").on(reportedEmpty));
+ * // load >> Label("always") >> cleanup
+ * cleaned.after(Label.of("always").on(loaded));
* ```
*/
-export function label(node: Node, text: string): Node {
- if (typeof text !== "string") {
- throw new Error(`label() takes a string to name the edge with, not
${describeValue(text)}`);
+export class Label {
+ readonly #text: string;
+
+ // The constructor is private only in TypeScript. Plain JavaScript can still
call
+ // `new Label("")`, so the text check lives here and not in `of`. A check in
`of` would not run
+ // for that call.
+ private constructor(text: string) {
+ if (typeof text !== "string") {
+ throw new Error(
+ `Label.of() takes a string to name the edge with, not
${describeValue(text)}`,
+ );
+ }
+ if (text.trim() === "") {
+ throw new Error(
+ "Label.of() takes the text to put on the edge, not " +
+ (text === "" ? "an empty string" : "a blank string"),
+ );
+ }
+ this.#text = text;
}
- if (text === "") {
- throw new Error("label() takes the text to put on the edge, not an empty
string");
+
+ /**
+ * Make a label with the given `text`. Use `on` to put the label on edges.
`text` must not be
+ * empty or only whitespace.
+ */
+ static of(text: string): Label {
+ return new Label(text);
}
- const target = resolveNode(node);
- const id = nodeId(target);
- if (id === undefined) {
- throw new Error(`label() takes a task or task group reference, not
${describeValue(node)}`);
+
+ /**
+ * Put the label on each edge that a `before` or `after` call draws between
its receiver and
+ * `node` or a node in `more`.
+ *
+ * So each edge of a fan-out can have its own label, and one label can go on
several edges:
+ *
+ * ```ts
+ * checked.before(Label.of("rows found").on(processed), Label.of("no
rows").on(reportedEmpty));
+ * // load >> Label("forks") >> [report_empty, notify]
+ * loaded.before(Label.of("forks").on(reportedEmpty, notified));
+ * ```
+ *
+ * Pass the value that `on` returns to `before` or `after`. Calling `before`
or `after` on that
+ * value throws an error.
+ */
+ on(node: Node, ...more: readonly Node[]): Node {
+ const text = this.#text;
+ const nodes = [node, ...more].flatMap((each) => {
+ // Labelling a labelled node again replaces its label.
+ const labelled = labelledNodes.get(each);
+ if (labelled !== undefined) return labelled.nodes;
+ const target = resolveNode(each);
+ if (nodeId(target) === undefined) {
+ throw new Error(
+ `Label.of(${JSON.stringify(text)}).on() takes a task or task group
reference, ` +
+ `not ${describeValue(each)}`,
+ );
+ }
+ return [target];
+ });
+ const { dagId } = nodes[0]!;
Review Comment:
`dagId` comes from the first node and the rest are never checked, so
`Label.of("x").on(taskFromDagA, taskFromDagB)` returns a node tagged with Dag
A, and the misuse message names only Dag A. The mismatch then surfaces inside
`#drawOrderEdge`, after part of the fan-out is already recorded. A check here
that every node has the same `dagId` would catch it before any edge is drawn.
##########
ts-sdk/src/coordinator/serde.ts:
##########
@@ -87,15 +87,31 @@ const TASK_MODULE = "airflow.sdk.coordinators.node";
*/
const TASK_LANGUAGE = "typescript";
-/** Python's `DAG.edge_info`, keying a group end by its join node as Python
does. */
-function serializeEdgeInfo(dag: Dag): SerializedValue {
+/**
+ * Python's `DAG.edge_info`, keying a group end by its join node as Python
does.
+ *
+ * This writes the labels in the order they were last set. Python writes a
label when it draws the
+ * edge. So when two labels land on one entry, the label set last is the one
kept.
+ */
+function serializeEdgeInfo(dag: Dag, graph: DagGraph): SerializedValue {
const groups = getDagTaskGroups(dag);
+ const grouped = new Set([...groups.values()].flatMap((group) =>
group.taskIds));
const info: Record<string, Record<string, SerializedValue>> = {};
- for (const { upstream, downstream, label } of getDagOrderEdges(dag)) {
+ const put = (from: string, to: string, label: string): void => {
+ (info[from] ??= {})[to] = { label };
+ };
+ for (const { upstream, downstream, label } of getDagLabelledEdges(dag)) {
if (label === undefined) continue;
const from = groups.has(upstream) ? upstream + DOWNSTREAM_JOIN_SUFFIX :
upstream;
const to = groups.has(downstream) ? downstream + UPSTREAM_JOIN_SUFFIX :
downstream;
- (info[from] ??= {})[to] = { label };
+ put(from, to, label);
+ // Python's `extract >> Label("x") >> group` also labels the edge from
extract to each root of
+ // the group. Like the Java SDK, this code labels the edges to the roots
for an edge drawn with
+ // `before` or `after`, but only when no group holds extract.
+ if (grouped.has(upstream)) continue;
+ for (const root of graph.groupRootsReached.get(upstream)?.get(downstream)
?? []) {
+ put(upstream, root, label);
Review Comment:
`groupRootsReached` is keyed by upstream and downstream only, so the roots
get labelled whichever verb drew the edge. Python labels them only when the
task is the receiver.
```ts
staging.after(Label.of("rows").on(extracted));
// here: extract -> staging.upstream_join_id, extract -> staging.stage,
extract -> staging.audit
// Python: extract -> staging.upstream_join_id (only)
```
`TaskGroup._set_relatives` (task-sdk `taskgroup.py:389-395`) passes
`edge_modifier` to `update_relative`, which writes the join entry, but then
calls `task.set_upstream(task_or_task_list)` for each root with no modifier, so
the root edges stay unlabelled.
##########
ts-sdk/src/sdk/dag-graph.ts:
##########
@@ -49,6 +49,11 @@ export interface DagGraph {
/** Each task's downstream task IDs, with every group endpoint expanded. */
readonly downstreamTaskIds: Map<string, Set<string>>;
readonly groupEdges: Map<string, GroupEdgeSets>;
+ /**
+ * For each labelled edge from a task to a group, the roots the group had
when the edge was
+ * expanded, keyed by the task ID and then by the group ID.
Review Comment:
`GroupEnds` is built after every task edge is linked, so
`ends.roots(downstream)` returns the roots of the finished task graph, not the
roots at the time the edge was declared. The two differ when a root is added
after the labelled edge:
```ts
extract.before(Label.of("L").on(g));
x.before(y); // x and y both in g
```
Python works the roots out when `>>` runs, so it labels `extract -> y`. This
labels only `extract -> x`.
```suggestion
* For each labelled edge from a task to a group, the roots of that group
in the finished task
* graph, keyed by the task ID and then by the group ID.
```
##########
ts-sdk/src/sdk/dag.ts:
##########
@@ -1189,7 +1252,22 @@ export class Dag {
}
#addOrderEdge(upstreamNode: Node, downstreamNode: Node, verb: "before" |
"after"): void {
- const text = labelledNodes.get(verb === "before" ? downstreamNode :
upstreamNode)?.text;
+ // Calling before() or after() on a labelled node throws, so only the
argument can be a
+ // labelled node.
+ const argument = verb === "before" ? downstreamNode : upstreamNode;
+ const labelled = labelledNodes.get(argument);
+ for (const node of labelled?.nodes ?? [argument]) {
Review Comment:
Each node is drawn in turn, and `#drawOrderEdge` validates and records one
edge at a time, so a bad node later in the list throws after the earlier edges
are already in `#orderEdges` and `#labelOrder`:
```ts
a.before(Label.of("x").on(b, taskFromAnotherDag)); // a -> b is kept, then
it throws
```
A Dag module that catches the error, or a loader that carries on past it, is
left with a partial fan-out. Validating every node before drawing any edge
would keep the call all-or-nothing.
##########
ts-sdk/src/coordinator/serde.ts:
##########
@@ -87,15 +87,31 @@ const TASK_MODULE = "airflow.sdk.coordinators.node";
*/
const TASK_LANGUAGE = "typescript";
-/** Python's `DAG.edge_info`, keying a group end by its join node as Python
does. */
-function serializeEdgeInfo(dag: Dag): SerializedValue {
+/**
+ * Python's `DAG.edge_info`, keying a group end by its join node as Python
does.
+ *
+ * This writes the labels in the order they were last set. Python writes a
label when it draws the
+ * edge. So when two labels land on one entry, the label set last is the one
kept.
+ */
+function serializeEdgeInfo(dag: Dag, graph: DagGraph): SerializedValue {
const groups = getDagTaskGroups(dag);
+ const grouped = new Set([...groups.values()].flatMap((group) =>
group.taskIds));
const info: Record<string, Record<string, SerializedValue>> = {};
- for (const { upstream, downstream, label } of getDagOrderEdges(dag)) {
+ const put = (from: string, to: string, label: string): void => {
+ (info[from] ??= {})[to] = { label };
+ };
+ for (const { upstream, downstream, label } of getDagLabelledEdges(dag)) {
if (label === undefined) continue;
const from = groups.has(upstream) ? upstream + DOWNSTREAM_JOIN_SUFFIX :
upstream;
const to = groups.has(downstream) ? downstream + UPSTREAM_JOIN_SUFFIX :
downstream;
- (info[from] ??= {})[to] = { label };
+ put(from, to, label);
+ // Python's `extract >> Label("x") >> group` also labels the edge from
extract to each root of
+ // the group. Like the Java SDK, this code labels the edges to the roots
for an edge drawn with
+ // `before` or `after`, but only when no group holds extract.
+ if (grouped.has(upstream)) continue;
Review Comment:
Python does not drop the label for a grouped upstream, it moves the edge:
`_convert_streams_to_task_groups` (task-sdk `edges.py:120-128`) replaces the
task with its group, so `e.before(Label.of("L").on(g2))` with `e` in `g1` gives
Python
```
{"g1.downstream_join_id": {"g2.upstream_join_id": "L"}}
```
where this writes `{"g1.e": {"g2.upstream_join_id": "L"}}`. That follows the
SDK rule that a label never changes the ends of its edge, so it reads
deliberate, and the new test at line 528 pins it. Worth saying so in the
comment, since "continue" alone looks like the roots case was just not worked
out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
