51
51
class TaskSpecification (TypedDict ):
52
52
# The task time
53
53
time : float
54
- # The optional collection id
55
- collection_id : NotRequired [int ]
56
-
57
- # Additional parameters
54
+ # The task's sequential ID
58
55
seq_id : int
56
+ # Whether the task is migratable or not
59
57
migratable : bool
58
+ # The task's home rank
60
59
home : int
60
+ # The task's current node
61
61
node : int
62
- user_defined : dict
62
+ # The collection id
63
+ collection_id : NotRequired [int ]
64
+ # User-defined parameters
65
+ user_defined : NotRequired [dict ]
63
66
64
67
65
68
class SharedBlockSpecification (TypedDict ):
@@ -182,15 +185,21 @@ def dict_merge(a, b):
182
185
data .get ("tasks" , []),
183
186
lambda t : TaskSpecification (dict_merge (
184
187
{ "time" : t .get ("time" , 0.0 ) },
185
- { "collection_id" : t .get ("collection_id" , None )} if "collection_id" in t else {}
188
+ { "seq_id" : t .get ("seq_id" , None )} if "seq_id" in t else {},
189
+ { "migratable" : t .get ("migratable" , True )} if "migratable" in t else {},
190
+ { "home" : t .get ("home" , 0 )} if "home" in t else {},
191
+ { "node" : t .get ("node" , 0 )} if "node" in t else {},
192
+ { "collection_id" : t .get ("collection_id" , None )} if "collection_id" in t else {},
193
+ { "user_defined" : t .get ("user_defined" , {})} if "user_defined" in t else {}
186
194
))
187
195
),
188
196
"shared_blocks" : self .__normalize_member (
189
197
data .get ("shared_blocks" , []),
190
198
lambda b : SharedBlockSpecification ({
191
199
"size" : b .get ("size" , 0.0 ),
192
200
"tasks" : set (b .get ("tasks" , {})),
193
- "home_rank" : b .get ("home_rank" )
201
+ "home_rank" : b .get ("home_rank" ),
202
+ "shared_id" : b .get ("shared_id" )
194
203
})
195
204
),
196
205
"communications" : self .__normalize_member (
@@ -200,7 +209,9 @@ def dict_merge(a, b):
200
209
"ranks" : self .__normalize_member (
201
210
data .get ("ranks" , []),
202
211
lambda r : RankSpecification ({
203
- "tasks" : set (r .get ("tasks" , []))
212
+ "tasks" : set (r .get ("tasks" , [])),
213
+ "id" : set (r .get ("id" , None )),
214
+ "user_defined" : set (r .get ("user_defined" , {}) if "user_defined" in r else {})
204
215
})
205
216
)
206
217
})
0 commit comments