Add checkpoint_map to parent_config

This commit is contained in:
Nuno Campos
2024-10-14 15:44:18 -07:00
parent b818bf2fba
commit 233bd78ee4
4 changed files with 140 additions and 4 deletions
+2 -2
View File
@@ -484,7 +484,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
patch_checkpoint_map(saved.config, saved.metadata),
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
patch_checkpoint_map(saved.parent_config, saved.metadata),
tasks_w_writes(
next_tasks.values(),
saved.pending_writes,
@@ -565,7 +565,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
patch_checkpoint_map(saved.config, saved.metadata),
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
patch_checkpoint_map(saved.parent_config, saved.metadata),
tasks_w_writes(
next_tasks.values(),
saved.pending_writes,
+1 -2
View File
@@ -179,8 +179,7 @@ def map_debug_checkpoint(
"step": step,
"payload": {
"config": patch_checkpoint_map(config, metadata),
"parent_config": parent_config,
# "parent_config": patch_checkpoint_map(parent_config, metadata),
"parent_config": patch_checkpoint_map(parent_config, metadata),
"values": read_channels(channels, stream_channels),
"metadata": metadata,
"next": [t.name for t in tasks],
+79
View File
@@ -9092,6 +9092,9 @@ def test_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
),
@@ -9250,6 +9253,9 @@ def test_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),),
@@ -9279,6 +9285,9 @@ def test_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(
@@ -9707,6 +9716,13 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
)
@@ -9770,6 +9786,15 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(
re.compile(r"child:.+|child1:")
): AnyStr(),
}
),
}
},
),
@@ -9798,6 +9823,9 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
),
@@ -10056,6 +10084,9 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(),
@@ -10085,6 +10116,9 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(
@@ -10170,6 +10204,13 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(),
@@ -10208,6 +10249,13 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(
@@ -10253,6 +10301,13 @@ def test_doubly_nested_graph_state(
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(
@@ -10444,6 +10499,12 @@ def test_send_to_nested_graphs(
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
@@ -10476,6 +10537,12 @@ def test_send_to_nested_graphs(
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
@@ -10931,6 +10998,12 @@ def test_weather_subgraph(
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
tasks=(
@@ -11020,6 +11093,12 @@ def test_weather_subgraph(
"thread_id": "14",
"checkpoint_ns": AnyStr("weather_graph:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("weather_graph:"): AnyStr(),
}
),
}
},
tasks=(),
+58
View File
@@ -7738,6 +7738,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
),
@@ -7903,6 +7906,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("inner:"): AnyStr()}
),
}
},
tasks=(
@@ -7934,6 +7940,9 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("inner:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("inner:"): AnyStr()}
),
}
},
tasks=(
@@ -8328,6 +8337,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
).tasks[0]
@@ -8374,6 +8386,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
)
@@ -8439,6 +8458,15 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(
re.compile(r"child:.+|child1:")
): AnyStr(),
}
),
}
},
),
@@ -8467,6 +8495,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
),
@@ -8732,6 +8763,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(),
@@ -8761,6 +8795,9 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{"": AnyStr(), AnyStr("child:"): AnyStr()}
),
}
},
tasks=(
@@ -8850,6 +8887,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(),
@@ -8888,6 +8932,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(
@@ -8933,6 +8984,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
}
),
}
},
tasks=(