|
23 | 23 | PresignedLink, |
24 | 24 | ) |
25 | 25 | from models_library.generics import Envelope |
| 26 | +from models_library.projects_nodes import NodeID, NodeState |
26 | 27 | from models_library.projects_pipeline import ComputationTask, PipelineDetails, TaskID |
27 | 28 | from models_library.projects_state import RunningState |
28 | 29 | from models_library.utils.fastapi_encoders import jsonable_encoder |
@@ -89,29 +90,39 @@ def create_computation_cb(url, **kwargs) -> CallbackResult: |
89 | 90 | pipeline = {} |
90 | 91 | node_states = {} |
91 | 92 | for node_id in body.get("subgraph"): |
92 | | - pipeline[node_id] = [ |
93 | | - "62237c33-8d6c-4709-aa92-c3cf693dd6d2", |
94 | | - "0bdf824f-57cb-4e38-949e-fd12c184f000", |
| 93 | + pipeline[NodeID(node_id)] = [ |
| 94 | + NodeID("62237c33-8d6c-4709-aa92-c3cf693dd6d2"), |
| 95 | + NodeID("0bdf824f-57cb-4e38-949e-fd12c184f000"), |
95 | 96 | ] |
96 | | - node_states[node_id] = {"state": {"modified": True, "dependencies": []}} |
97 | | - node_states["62237c33-8d6c-4709-aa92-c3cf693dd6d2"] = { |
98 | | - "modified": True, |
99 | | - "dependencies": ["2f493631-30b4-4ad8-90f2-a74e4b46fe73"], |
100 | | - } |
101 | | - node_states["0bdf824f-57cb-4e38-949e-fd12c184f000"] = { |
102 | | - "modified": True, |
103 | | - "dependencies": [ |
104 | | - "2f493631-30b4-4ad8-90f2-a74e4b46fe73", |
105 | | - "62237c33-8d6c-4709-aa92-c3cf693dd6d2", |
106 | | - ], |
107 | | - } |
| 97 | + node_states[NodeID(node_id)] = NodeState.model_construct( |
| 98 | + **{"state": {"modified": True, "dependencies": []}} |
| 99 | + ) |
| 100 | + node_states[ |
| 101 | + NodeID("62237c33-8d6c-4709-aa92-c3cf693dd6d2") |
| 102 | + ] = NodeState.model_construct( |
| 103 | + **{ |
| 104 | + "modified": True, |
| 105 | + "dependencies": {NodeID("2f493631-30b4-4ad8-90f2-a74e4b46fe73")}, |
| 106 | + } |
| 107 | + ) |
| 108 | + node_states[ |
| 109 | + NodeID("0bdf824f-57cb-4e38-949e-fd12c184f000") |
| 110 | + ] = NodeState.model_construct( |
| 111 | + **{ |
| 112 | + "modified": True, |
| 113 | + "dependencies": { |
| 114 | + NodeID("2f493631-30b4-4ad8-90f2-a74e4b46fe73"), |
| 115 | + NodeID("62237c33-8d6c-4709-aa92-c3cf693dd6d2"), |
| 116 | + }, |
| 117 | + } |
| 118 | + ) |
108 | 119 | returned_computation = ComputationTask.model_validate( |
109 | 120 | ComputationTask.model_config["json_schema_extra"]["examples"][0] |
110 | 121 | ).model_copy( |
111 | 122 | update={ |
112 | 123 | "id": TaskID(f"{kwargs['json']['project_id']}"), |
113 | 124 | "state": state, |
114 | | - "pipeline_details": PipelineDetails( |
| 125 | + "pipeline_details": PipelineDetails.model_construct( |
115 | 126 | **{ |
116 | 127 | "adjacency_list": pipeline, |
117 | 128 | "node_states": node_states, |
@@ -142,7 +153,7 @@ def get_computation_cb(url, **kwargs) -> CallbackResult: |
142 | 153 | update={ |
143 | 154 | "id": TaskID(Path(url.path).name), |
144 | 155 | "state": state, |
145 | | - "pipeline_details": PipelineDetails( |
| 156 | + "pipeline_details": PipelineDetails.model_construct( |
146 | 157 | **{ |
147 | 158 | "adjacency_list": pipeline, |
148 | 159 | "node_states": node_states, |
|
0 commit comments