Repository: atomind-ai/mlip-arena (language: python)

Issue:
Add EOS
Refer to [Matcalc EOS](https://github.com/materialsvirtuallab/matcalc/blob/8d3b7348f975f04606e1a83b8c0ff2ecf101865a/matcalc/eos.py#L20) to implement prefect job and flow


Discussion / hints:

Fix patch:
diff --git a/mlip_arena/tasks/combustion/H256O128.extxyz b/mlip_arena/tasks/combustion/H256O128.extxyz
index 9105354..59fcd94 100644
--- a/mlip_arena/tasks/combustion/H256O128.extxyz
+++ b/mlip_arena/tasks/combustion/H256O128.extxyz
@@ -1,386 +1,386 @@
 384
 Lattice="30.0 0.0 0.0 0.0 30.0 0.0 0.0 0.0 30.0" Properties=species:S:1:pos:R:3 Built=T with=T Packmol=T pbc="T T T"
-H        6.71669500      27.40330700      13.68784000
-H        6.56413900      26.94187400      13.13356600
-H        4.01874500      24.54731500       2.42614400
-H        4.06326800      24.60306800       1.69243900
-H       19.81990800      11.97306900      11.66638000
-H       19.80020800      12.70644500      11.59436300
-H       28.98607000      12.86336400      12.19767400
-H       29.00971400      12.88360500      12.93418200
-H       10.37632500       7.60336700       9.09607600
-H       10.09412100       7.01607100       9.44083300
-H       10.21969400      27.00047900      10.37176700
-H        9.56226800      26.66727600      10.35833100
-H        3.51642300       4.64618300      13.26100400
-H        3.19087000       4.01330300      13.06893100
-H       17.07272200      26.43633800       3.97721600
-H       17.76156200      26.60077000       4.18185100
-H        3.58254200       8.40601100       1.85885100
-H        3.32028600       8.99119200       1.49526000
-H       20.93526200      14.99894700      19.94272400
-H       20.92651600      14.99223100      20.67980700
-H       26.78606500      14.69290300      18.64689300
-H       26.43088200      14.16219600      19.01514500
-H       23.31077000      25.02192200      25.42458800
-H       23.23137700      24.98837100      26.15669800
-H        4.23870900      22.06352800       2.24193700
-H        4.13605700      21.89474500       2.95214000
-H       17.19183000      13.50635300      20.97020900
-H       17.33285700      13.55720300      20.24844800
-H       11.40128700      21.01758100      13.20015300
-H       11.23939200      21.24064200      13.88385500
-H        8.08553400      17.35614800       4.83985100
-H        7.85698800      18.00868300       4.58416100
-H       18.23851700      17.87971000      18.64072400
-H       18.18586700      17.93400600      19.37400000
-H        7.67448900       9.27045600       9.58857400
-H        8.38197500       9.28217700       9.38183700
-H       22.81810900       7.01891100      11.62916000
-H       23.02343500       6.88860700      10.93326000
-H        7.00678500      28.41965000      16.15063200
-H        6.61194800      29.00429800      16.36441700
-H       10.22790300      12.62515700      19.02217900
-H       10.06283300      12.03600000      19.43334700
-H        5.67633000      19.26890500       3.09689400
-H        6.16859600      19.35120500       3.63940200
-H        7.83689900      28.57325000       6.00799800
-H        8.20904000      29.01838100       5.55326500
-H       24.24356500       1.95484800      23.83583200
-H       24.33048700       1.23652000      23.97677000
-H        8.15609200      12.56884200       2.83826200
-H        8.47473700      12.51613700       2.17561500
-H        1.98944400      22.63517100       9.38706600
-H        1.99816600      22.52032700       8.65895300
-H        3.21198500      13.34481400      26.27085600
-H        2.83315800      12.78596400      26.56681200
-H       23.13687400      14.90616700      17.31808300
-H       23.60063200      15.47147100      17.41174900
-H       24.37760200      23.67694500       6.16965600
-H       24.26969800      23.54491000       5.45248300
-H       12.12534300       0.92505100       5.56915800
-H       12.59580200       1.08149500       6.11469100
-H       14.77057800      23.04450600      14.78597800
-H       14.88153400      22.32226700      14.88331100
-H       10.12077400       6.02995200      19.27346700
-H       10.56701000       5.58085700      19.65109300
-H        8.88571400      26.92992500      13.97009200
-H        9.13262800      27.42418800      13.48208400
-H       13.61412900      25.32932600       5.89573200
-H       14.17465000      25.02676400       5.52467300
-H       17.25064000       9.39374700      16.24924400
-H       17.23228900       9.42911200      16.98533300
-H       12.66108800      13.49816800       7.35969900
-H       12.75204000      13.66251500       6.64686500
-H       19.85540500      12.40447200      18.67913500
-H       19.87736300      13.08567100      18.39824500
-H       21.16039700      14.54404800      12.22178400
-H       20.77918300      14.69154600      11.60832300
-H        1.89709500      13.01516500      11.87415000
-H        1.61578300      12.78775700      11.23183900
-H       22.48727800      20.55335700      26.81374300
-H       22.08316200      19.93770500      26.84656000
-H       18.09641200      11.03175000      20.02888200
-H       18.08987100      11.14550300      19.30057500
-H        9.03487100      28.27705600      17.67592300
-H        9.02850300      28.47186100      18.38685500
-H       21.61332100      18.16863300      14.82530000
-H       20.91922800      18.33149900      15.01271300
-H        5.06777300       6.91012100      12.14248900
-H        5.14167800       7.53010400      12.53437400
-H       15.24240300       1.00156500       8.79894400
-H       15.15041500       1.70061000       8.58379800
-H        7.38156900      16.75420300      20.84490000
-H        7.49809200      17.42858000      20.57094300
-H        5.80395600       5.23791900      23.23164500
-H        5.59378400       5.32074900      23.93334300
-H       25.86193300      21.36196900      25.55815300
-H       26.43947000      21.14184700      25.15639300
-H       15.99876300      19.78997000       1.87456100
-H       15.41808700      19.79298200       2.32867600
-H       13.76283200      27.95478100       8.97671100
-H       13.72689900      27.99668400       8.24161500
-H        3.90503500      26.10922400      10.77631300
-H        3.36520200      25.83210000      11.19487700
-H        3.70175700       2.02893500      10.84822800
-H        3.98998200       2.68818200      11.00864200
-H       19.61345400      10.74535100      22.22723700
-H       19.65167500      10.72936900      22.96323800
-H       21.15335200      11.20594200      26.99475000
-H       21.53029400      10.81346000      27.49202900
-H       20.09434500       6.25424400      17.87873800
-H       20.01076700       6.21479100      18.61008800
-H        3.69537200      12.39851000       8.93646700
-H        4.09765200      11.81750500       9.14627000
-H       22.83093000      16.73250700      24.19709400
-H       23.31256100      17.17900100      23.86229600
-H        3.36198000      12.89766000      23.38943200
-H        3.45885500      13.00681500      24.11200700
-H        2.04611800      15.71757800       8.59329600
-H        1.43234600      15.77875600       8.18961900
-H        6.12945900      15.14026700      20.02887800
-H        6.02278200      14.57365100      20.48820400
-H        5.08911400      27.18139700       2.77337500
-H        5.06662100      27.18501900       3.51018800
-H       14.70758200       4.01177700      13.31696800
-H       14.89664300       4.24570600      12.64395500
-H        3.65814300      29.03356400      14.88286700
-H        3.73378000      28.95725100      14.15357400
-H       16.99951200      11.42586000      13.86881500
-H       17.03417100      11.12538100      13.19656200
-H        2.06797500      18.08136200      13.54583400
-H        1.48136600      18.09474000      13.99206700
-H       14.60250500       7.57683000       5.97130200
-H       15.33793200       7.61271900       6.00698800
-H       26.83711700       3.23048000       3.03884500
-H       26.55008700       3.90471800       3.11903900
-H       22.99777600      19.34731400      20.31560600
-H       22.45982800      19.77772100      20.57785400
-H       16.41966600       2.80599400      11.89128000
-H       16.56901900       2.08412200      11.88831400
-H       11.73017700      14.43811800      22.41741700
-H       11.12423500      14.82954800      22.56917700
-H       17.38249100      23.87139500      12.76358200
-H       17.38350900      23.82842900      12.02767000
-H        1.00061100      21.01224200       1.31470000
-H        1.44321900      21.03360900       1.90381400
-H        8.09004900       6.72936800       0.99999800
-H        8.19345600       7.35764200       1.37147100
-H        3.89822700      15.37766400       9.73104000
-H        4.54288400      15.70539800       9.87394700
-H       22.22552500      28.69576500      19.77030500
-H       21.65346500      29.00000300      19.41873800
-H        4.65573400      15.45943800      27.22827700
-H        5.35397400      15.56981400      27.43730400
-H       12.98406500       8.83317400      11.91500600
-H       12.47937800       9.06790600      11.43167900
-H       21.24968700       1.72985200      23.51877400
-H       21.91940400       2.02865000      23.44385400
-H       11.59386400      22.25408800      27.15665400
-H       11.73855200      21.54498900      27.29686000
-H       21.77510800       4.83204700      19.45360300
-H       21.74760500       4.86959100      18.71790800
-H       28.68304100       9.21677300      24.30293500
-H       28.79634600       9.49180700      23.62844900
-H       18.79159300       1.71351500       2.69284300
-H       18.77465700       2.18174600       3.26195300
-H       19.17586000      13.02896000      14.86997500
-H       19.71644200      13.51755200      14.75833800
-H        6.02966400      25.52593900      10.73507600
-H        5.80325300      25.22399900      10.10184400
-H        7.38506300       1.43983700      15.49859600
-H        6.79506700       1.06865500      15.73847300
-H       27.58382000      24.42158300      19.38106600
-H       27.42251200      25.04349100      19.01964600
-H        9.93092500      22.73321600      10.93929900
-H       10.22059100      22.24094300      11.40531900
-H       15.50084100      11.06308800       6.45439200
-H       15.87813400      11.34578600       5.88769500
-H       16.10016600      28.54356000       2.40068100
-H       16.19875700      27.92618400       2.01011700
-H       17.12622300       9.66589700       8.78315600
-H       16.48701000       9.96416700       8.99729500
-H       25.00042100      21.05189700      28.69399300
-H       25.35829500      20.90402200      28.06671900
-H       11.72392500      14.35901600      14.78824400
-H       11.48529200      14.43413500      14.09482800
-H       27.80088200      28.99951400      23.03485400
-H       27.35755700      28.41175300      23.07244700
-H        2.81120100      21.97327000       6.33202600
-H        3.02278700      22.64853600       6.53857400
-H       11.16091600      21.16245100      20.56587100
-H       10.97600400      21.57060200      19.98052200
-H       25.63314200       8.24689800      11.42146200
-H       25.84695500       7.98978900      10.76450500
-H        5.21035300      25.88590800      28.03985400
-H        4.61973200      25.44924700      28.10235800
-H       14.83558600      10.14445700      28.76566400
-H       14.45468700      10.77548600      28.75415200
-H       24.19325000      21.19802200      21.61625300
-H       24.60150400      21.42553100      22.18632500
-H        3.38383900      29.00000600       6.51473400
-H        3.38610800      28.83091700       7.23224200
-H       17.13203800      26.72357300      21.49248700
-H       17.08711600      27.41871500      21.25129500
-H       21.43007400      22.30567900      19.46406700
-H       20.99498700      21.73213800      19.30543700
-H       19.31644700      11.15485400      16.07821400
-H       19.68199100      11.18838700      16.71748400
-H       10.93924200       5.99489900      26.27955500
-H       10.41495800       5.47711900      26.30060500
-H       14.95008000      15.52212300       7.01442400
-H       14.72471900      15.69369600       7.69500400
-H       18.99727600      22.68120300      25.90377400
-H       19.40957700      23.25149200      25.68424600
-H        8.88447900       1.79024300      17.17555900
-H        8.86324600       1.77770600      17.91231200
-H       12.15121200       1.54535800      11.42100700
-H       11.71339400       1.67546500      11.99962700
-H        2.22952700      10.78313900      28.21835700
-H        2.22169400      10.63858800      27.49554500
-H        4.71118900      24.88843500       6.95444600
-H        4.75449800      25.06231100       6.23939000
-H        1.28060100       1.44032300      19.29701200
-H        1.83528100       1.36120500      18.81796700
-H       12.45357900       9.77922200      25.48060700
-H       11.74916000       9.99016300      25.53267400
-H        8.23385400      11.13248500      18.56147300
-H        8.05531600      10.83169400      19.21036700
-H       24.48638300      19.81957100      23.69282200
-H       25.13594000      19.58627200      23.95178100
-H       26.70376200       9.82401100      27.55103500
-H       26.90704500       9.33315600      28.06206400
-H        8.00838700       4.22559800      18.28897400
-H        8.42532000       3.95184600      17.74616600
-H       14.94452000       3.51127400       4.84402200
-H       14.89449900       3.42688600       5.57463200
-H       19.83606200       2.22162900      13.82277900
-H       20.02512800       2.59599500      13.21654700
-H        6.67903100       3.41714500       3.27561700
-H        5.96308500       3.29632200       3.40304000
-H        3.40442700       8.43161300      18.68125100
-H        3.47980300       8.38947800      19.41334200
-H       11.59437400      13.32368500      26.26415800
-H       12.23470100      13.62606800      26.46899700
-H       14.10053700       2.78998700       1.85965400
-H       14.21152000       2.83848100       2.58680300
-H       20.73807400      13.46815700      24.67708300
-H       21.38650500      13.80030400      24.56470100
-H       20.60974700      10.79901200      13.69614100
-H       20.72497000      10.71348400      14.41920500
-H        3.17495000      22.77164300      26.22314900
-H        3.62633300      22.80052800      25.64105600
-H       28.64961800      27.52039100      17.72169200
-H       28.65554600      27.60539900      18.45391600
-H        7.29178200      28.18007600       2.36813400
-H        7.28037500      28.22151000       3.10404600
-H        6.81956500      10.50511700       4.98278800
-H        6.16534900      10.84336000       4.95106300
-O       28.65774000      26.05310300       1.73187800
-O       27.82042800      25.18551200       1.41787100
-O        7.89599500      15.38902200      14.65745800
-O        7.88882100      15.60518500      13.43041800
-O       11.37704800      22.17970800       3.53597600
-O       10.62568400      23.16642300       3.65535900
-O       13.21571800      10.10612800       5.39037400
-O       13.81380800      10.31472700       4.31744300
-O       12.93699800      12.19119900      15.32963100
-O       13.36779300      11.62125600      16.35040800
-O       27.62157300       9.98656400       1.83057900
-O       26.68111300      10.54405800       1.23296700
-O        6.77274700      17.50780000       1.73044000
-O        5.79918400      17.87923300       1.04734400
-O       26.06597600      11.96086100      28.39447800
-O       26.32478700      12.75375600      27.46887600
-O        1.16753100       1.35249300      15.49177600
-O        1.18059100       1.20596000      14.25453600
-O       15.08210300      17.91068800      23.23158100
-O       15.17178200      18.46551400      22.11958700
-O       18.60542300      22.96513600      18.41436700
-O       19.71179100      23.45867400      18.70553300
-O        1.55709700      26.18859700       1.16408800
-O        2.61399600      26.49796400       1.74689400
-O       13.99292900      14.15190600      21.10355700
-O       14.04751700      14.18255600      19.85917400
-O       10.38541200      18.23372400      12.93027500
-O        9.91188600      17.49941400      12.04203600
-O       24.58403200      25.06294600      21.97592600
-O       25.05025500      26.20923800      21.83081500
-O       12.79523200      10.91568200      21.34176200
-O       12.59216100      11.07454900      22.56074900
-O        7.88339000      21.12603300       4.34091800
-O        8.23639900      20.09896500       3.73024400
-O       14.29669600      11.48202900      10.12026500
-O       14.65658800      11.78505600       8.96655000
-O       24.45593800      10.09229400      18.89934500
-O       23.27212000      10.31487400      19.21783900
-O       10.14103800      28.74148400       8.55297800
-O        9.09630200      28.47275800       9.17645200
-O        6.53750100      18.89849400      13.13483400
-O        7.18188100      19.75253500      12.49624700
-O        5.99854200      21.81797900       5.25108200
-O        5.99224800      23.06318800       5.20840100
-O        9.03156900       6.67285800      22.94269500
-O        7.94989000       6.26883900      23.41082700
-O       11.37806100      15.86060900      17.89431700
-O       10.58973300      15.69880900      16.94312500
-O       18.69252800      21.49332500      14.52289300
-O       19.04391300      21.05544000      13.41060100
-O       10.88775200      25.62526500      15.90813900
-O       10.56171400      25.66761900      17.10993500
-O        5.83323700       6.64702500      10.08718000
-O        5.78880400       6.30641500       8.88950900
-O        7.71126200       8.59122800      18.86977200
-O        7.73147600       8.60812400      17.62409400
-O       20.74284900       2.12287500      21.37159100
-O       20.54629600       2.47368100      20.19230800
-O        9.33194100      10.52852700      15.25052100
-O        8.99841100       9.39518300      14.85467600
-O       15.28732400      14.68259500       9.86397700
-O       14.24814200      15.19621500      10.32081500
-O        8.79509800       2.36395800      13.47414000
-O        8.02117200       2.64244400      12.53825000
-O       22.07041200      27.82260300      11.86453400
-O       22.47004200      27.29076900      10.81103700
-O       25.64900400      10.71424900      24.59372700
-O       25.14662500      10.86773900      25.72353300
-O       17.46952600      13.14388400      17.59581300
-O       17.02295500      12.33034400      16.76446800
-O       12.27907200       3.72064200       4.88171600
-O       12.34240900       3.77544700       3.63857800
-O       19.50544800      19.26652500      26.41877500
-O       19.46811300      19.42938700      27.65347700
-O       16.09394700      23.74355500       8.66223700
-O       16.43146400      24.61823500       7.84161300
-O       12.88585100      27.95040100      27.15844000
-O       13.33429000      28.28825200      28.27072000
-O       17.90265700      22.35366300       6.91603700
-O       18.19323300      21.88298700       5.79959800
-O       23.38376600       4.75998600       8.13216000
-O       23.39868600       4.98386700       9.35774600
-O       16.75294800       1.89397000       5.04381400
-O       17.32240900       2.99348800       5.18231300
-O       17.56519600       7.26680100      21.77071000
-O       18.31705700       8.25843700      21.83209000
-O       11.99835900      11.23426800      18.12324800
-O       12.36029500      11.59229700      19.26044800
-O        7.58959900      13.32747400      11.00503600
-O        6.37218600      13.53282600      10.83727300
-O       18.59862600      24.35057600       9.30150100
-O       18.94728500      24.62991400       8.13839600
-O       23.48008200      18.54852600      17.99906900
-O       23.36944900      17.83172600      16.98597300
-O       22.77455400      28.18135800      14.92714700
-O       23.86580000      27.65620500      14.63422500
-O       20.78948500       3.67260900       1.83948100
-O       21.56215500       3.51027100       2.80334500
-O        9.05637000      15.36340100      27.31637100
-O       10.18282500      14.88236800      27.54465200
-O        5.36124000      14.96264300      23.52556700
-O        4.24760100      14.87097200      22.97437700
-O       22.64266600      13.26344800      26.40413500
-O       22.64331400      14.50641200      26.49043800
-O        1.32756400      12.17561700      20.06774300
-O        2.10349400      12.35107600      21.02667600
-O        1.18281700       3.80763100       4.94719100
-O        1.01891800       2.82667000       5.69769700
-O       12.59019600      17.94027500      24.45964400
-O       12.50958200      18.08831700      23.22514300
-O        9.77496800      21.51131100       7.49102600
-O       10.45743600      20.53412700       7.85400700
-O        1.69369700       6.90651400       3.72886700
-O        1.19082500       6.11102100       4.54539400
-O       26.88872500      10.54198400      11.12473400
-O       27.32702000      10.28414700       9.98727000
-O        1.72700700      26.73221500      14.58446600
-O        2.56724200      25.91897900      15.01465200
-O       10.15694600      18.69534200      19.15618300
-O        9.99999700      19.29683100      20.23599100
-O        1.99644200      10.71276900      12.34430700
-O        2.09478800      10.26022600      13.50100000
-O        2.90497500       4.19153100       2.82131400
-O        1.84632400       4.84452600       2.74877200
-O        5.05591600      16.45745000      12.34283200
-O        6.13192000      15.97608900      11.93920500
-O        6.49320100      16.95653500      24.96253500
-O        6.59211000      17.17298400      26.18555300
+H       16.87897000      14.36340900       5.78639500
+H       17.43576500      14.42974000       6.26492700
+H        3.55483600       6.05817200      17.69903100
+H        3.93076700       5.56617900      17.29899100
+H        9.49657700      24.94365100      23.70565700
+H        8.79155100      24.79399100      23.55088300
+H       23.86076500      25.42809100       5.93695500
+H       24.14102300      25.10486100       6.53728200
+H        3.01569800      26.19656300      15.08020000
+H        2.75520000      25.98067600      14.42525900
+H        9.67788700      17.96928400       5.86481000
+H       10.22656200      18.46063200       5.83403500
+H       12.57870500      16.09171300      25.13668100
+H       12.52257400      16.32590900      24.43996300
+H       12.70913000       6.58341200      15.79997200
+H       13.14484100       6.28817300      16.31611500
+H       21.57221500      26.22007700       8.20775100
+H       21.49858000      25.76305400       7.63406000
+H        5.23150000       2.79202900      17.94668200
+H        5.34160000       2.30038000      17.40856200
+H        3.20140800      24.30032700      16.20596600
+H        3.01866600      23.59120400      16.29060400
+H       24.14000600      28.98000500      16.44843000
+H       24.06795800      28.97179200      15.71483900
+H        5.06521500      20.71716800       5.13534300
+H        4.35233500      20.56737200       5.02230400
+H       20.35575900      22.80574000      20.92112100
+H       20.33775300      22.82449400      20.18441300
+H        9.44001100      14.49775500      28.22965700
+H        8.79780000      14.40101700      27.88091500
+H        9.85111200      22.18687800       5.65961300
+H       10.23873200      22.52486700       5.13147700
+H        2.43073100      14.23882800       9.88843100
+H        2.41513200      14.74593300      10.42323500
+H       24.12085900      18.37439500      13.96560000
+H       24.11116500      17.72761900      14.31915600
+H        8.20116600      27.96832200      10.13677000
+H        7.93113800      28.48718300      10.58541500
+H        9.12140400      26.33764100       3.10957700
+H        9.77423300      26.15812100       3.40112700
+H        4.61212100       6.00445300       4.16492600
+H        4.64955500       6.71113500       4.37135300
+H       11.25926100      17.21528300      14.18800300
+H       11.45743100      17.84358800      14.51872400
+H        8.30883000      10.16305200      12.29874000
+H        7.72273600      10.57641900      12.12834100
+H       27.50695500      17.01494000       3.57254600
+H       26.93803700      16.58868000       3.76760800
+H        5.77076100       2.49322200      14.46871700
+H        5.62762700       2.75411500      15.14315000
+H        6.51732300      25.84318300      26.82939800
+H        6.55807500      25.12027600      26.69098100
+H       12.12323800      24.13458100      24.41591800
+H       12.45973600      24.45908800      23.84593700
+H        5.17331100      23.16438400      16.63650900
+H        4.88027500      22.87390600      17.24738200
+H        8.10959100      21.29987000       4.46261500
+H        7.63531700      20.73605600       4.43829600
+H       22.74497200      17.55560800       1.89494000
+H       22.79182400      17.39285800       2.61238800
+H       10.18128000      19.85571100      21.25732700
+H       10.60903100      20.45053300      21.17590400
+H        7.36765600      22.32656800      11.62091000
+H        7.36366200      22.32242700      12.35805300
+H       22.11051700      17.35754100      10.70421900
+H       21.63342400      16.93939000      10.32879200
+H       17.11800800      14.46985000      18.41272300
+H       17.76175400      14.14922200      18.57459600
+H       23.88598700      25.00383500      10.24699000
+H       24.23686000      25.54091700       9.88388500
+H       23.07268800      28.13033100      13.42606000
+H       22.86236700      28.70230100      13.84082400
+H       19.71929900      26.69691800      16.96391200
+H       19.77256100      26.69242100      16.22868700
+H       10.20897000       5.12509300      26.31992400
+H       10.27351800       5.67934800      25.83820900
+H        7.74945200      21.75294300      26.96569600
+H        7.85650000      21.94562800      26.26225700
+H       28.99846900       5.82412900       4.51948000
+H       28.84661100       5.17571900       4.20338400
+H       14.00669600      21.28067600      10.51077400
+H       14.02815700      20.57927500      10.28496600
+H       24.56688200       8.44802400      10.18169000
+H       24.87884700       8.74970200      10.77757700
+H       19.82097000       6.71060100      21.14119900
+H       20.20300600       6.57543700      20.52541300
+H       12.50680000      16.29448300      21.20002800
+H       12.58574100      16.94571300      21.53630400
+H        2.88750900       6.48194900      12.41992700
+H        2.63814400       5.97929100      11.94184100
+H       18.31924100      27.95926100      11.93343800
+H       18.08127600      28.28411700      12.55089700
+H        8.80421900      21.80835300      22.53574400
+H        8.19015200      22.18697500      22.38414900
+H       16.51146900      26.58804000      12.13737900
+H       16.51397100      26.60180800      11.40034500
+H        3.01154500       1.61542600      13.17298500
+H        2.31442300       1.37599500      13.16260300
+H       14.38141000      16.74073800      22.79765000
+H       14.54283900      16.49896600      23.47507200
+H       27.32062300      13.13890000       3.91277400
+H       27.68551200      13.68559000       4.24654100
+H       14.17603100      27.61355700       6.47229000
+H       14.22488300      28.10975000       5.92931700
+H        9.46812300       7.70071300       3.45368100
+H        9.61903800       8.06415600       4.07701700
+H        2.18384900       3.88122900      12.55886400
+H        2.16274500       3.79907900      11.82659400
+H        8.39628800      28.36448200      27.01582200
+H        7.83759500      27.92933600      26.81106900
+H       12.15799700      14.43608700      11.46881200
+H       11.91985600      14.40143500      10.77203200
+H       13.07421100      20.72377800      27.23797600
+H       12.38307500      20.94338900      27.37032300
+H        5.74890500      20.92045300      10.67404300
+H        6.07702100      21.30697500      10.13892200
+H        1.20101800       3.91867000       5.21703800
+H        1.33343400       3.87199400       5.94071000
+H       15.87910900       3.60421600      17.51325200
+H       16.55088000       3.81033800      17.29042700
+H       12.77118400      24.89708700      14.04640100
+H       12.93989800      25.26872400      13.43253200
+H        9.17302000       6.59818400       6.60346000
+H        9.20243000       6.69949200       7.33303900
+H        3.82873400      21.34798500      18.48337600
+H        4.44631900      21.08569500      18.78867100
+H        6.64065400      26.32535100       2.13715800
+H        7.07217200      25.92015200       1.69781700
+H       23.61875500      28.68096800       2.12527200
+H       24.20152800      28.50119700       1.71118000
+H       28.97955000      28.31110000      20.77510000
+H       28.51757300      27.79654500      21.03049100
+H        5.03608000       8.09600900      15.70374700
+H        5.01029200       8.77247600      15.41195800
+H        6.38181200      23.23894300      14.67726800
+H        6.72455700      23.67837800      15.15980100
+H        4.17225800      15.87846200      26.32328700
+H        3.60033300      16.10995900      25.91989500
+H        9.98548800      27.03736900      11.34257400
+H        9.89827200      26.31071900      11.43082100
+H       11.08624000       3.40906900      27.43653900
+H       11.03887400       3.46845400      28.16978100
+H       18.17462600      26.60855300      18.64341900
+H       17.84284900      27.25533700      18.76592900
+H       24.56041800      23.73132300      15.45710400
+H       24.15216800      23.55566200      16.04522700
+H       12.42149000       8.51852900      26.92738800
+H       12.85187200       8.54739800      26.32959800
+H       12.58494400      23.66894000      26.51841800
+H       11.87423100      23.85271600      26.58569800
+H       25.66776400      11.38987800       3.60979300
+H       25.04312400      11.04530700       3.79556200
+H       26.02410100      25.69683300       2.37428800
+H       26.49306900      25.51267100       1.83617300
+H        1.17207300      13.56847700      17.81787500
+H        1.09749900      13.56957200      17.08449200
+H       16.34427300      15.52582100      25.01154500
+H       16.21113800      15.84954200      25.66030800
+H        8.34844300       4.44909300      22.93722500
+H        7.86249000       4.49217500      23.48986100
+H       24.87485300       5.31864800       1.40835600
+H       25.19256300       5.84446800       1.00094100
+H        3.42570400      18.20815000      27.19415500
+H        2.77521600      18.34972600      26.87755500
+H       26.13269700      10.18845100      14.93319100
+H       25.85987900       9.71045000      14.44278400
+H       16.37800300       2.82461800      14.94014800
+H       16.88769100       2.31071600      14.80037500
+H        2.68461500      21.47767100      25.65109500
+H        2.86206000      21.90699200      25.07872200
+H       12.04855900       6.95917200      11.68418900
+H       12.39801200       7.51168600      11.34356500
+H       17.65835900      17.90689000      21.24729900
+H       17.57119400      17.94728700      20.51642000
+H       19.13123000       3.51526100      22.77599000
+H       18.50245100       3.89609300      22.72106000
+H       17.59014300      22.68251700       7.89677900
+H       17.94157500      23.30906000       7.73138900
+H       17.80517200      27.73224100      26.14671300
+H       17.17556300      27.67636700      25.76739200
+H        1.23383900      16.30952500      11.48348600
+H        1.21410800      16.78863900      10.92359800
+H        5.14106900       3.12858300      22.32613900
+H        4.84773300       3.33214200      22.97106700
+H        1.54627400      27.32284900      21.64647900
+H        1.89536100      27.35916600      22.29473300
+H       23.61494100       4.81278600      28.82197500
+H       23.15278300       4.71015800      28.25691600
+H       16.55009200       9.41232000       9.23163100
+H       16.39011400       8.95635500       9.78833300
+H        9.24672000      17.04310600      22.97075500
+H        9.69788100      16.57902500      23.32359500
+H        5.92999000       6.70344600      12.54324900
+H        6.12013700       6.02943400      12.77339200
+H       20.55806500      20.91865300      27.16038800
+H       20.17869900      21.08398200      27.77043800
+H        6.86904800      16.04280700      13.61109400
+H        7.39742300      16.55244000      13.54396000
+H        2.56041500      20.50409800       7.85659000
+H        1.94972600      20.22580200       8.16158900
+H       10.30946400      26.76199200      22.35478000
+H        9.74930900      26.93466800      22.80179500
+H        9.37796100       7.80951200      12.91974400
+H        8.83394300       8.02560600      12.47167900
+H        8.87347500      15.06285800       6.33146400
+H        8.47422200      14.92537200       6.93570600
+H        7.76353600       8.45873000      15.11360900
+H        7.75143800       8.55689900      14.38310900
+H        3.37293300      18.43420600      23.45136000
+H        3.96942500      18.08124700      23.70242100
+H       26.76422700      16.17472200      17.31888500
+H       26.40881800      16.38505300      17.92950700
+H       23.72340800      16.16827600      17.16088600
+H       23.54079100      15.74444300      16.58605400
+H        9.04422200      24.03668900      13.69386900
+H        8.77587200      23.93037800      14.37217600
+H       26.63067400      17.02388700       6.06368500
+H       26.82086300      16.77080800       6.72941200
+H       11.89740000       7.80260100      21.08513000
+H       12.35463200       8.38081400      21.08994900
+H       18.88034700      13.04076800       4.44410900
+H       18.99432900      13.09305500       5.17053000
+H       12.33984600      11.69259200      22.86388500
+H       12.74815800      11.21413800      23.24830300
+H        5.80287000      16.81823400      16.93960100
+H        5.49624100      17.28191100      17.42374600
+H       28.49793800       7.81327600      23.13893400
+H       28.47600900       7.97048500      22.41906000
+H       12.83557500      11.24373500      25.45085300
+H       12.41881800      11.24066000      26.05889700
+H        8.27635600      24.04023300       1.68334600
+H        7.96194000      23.42057200       1.43722300
+H       19.77199600      11.69701700       8.46538200
+H       19.78074500      11.70307100       9.20247100
+H       11.31673000      12.37398400      14.55721300
+H       10.84297500      12.47166100      14.00094900
+H       10.79125000      20.69161200       3.40245000
+H       11.27250900      20.38956300       3.87209900
+H       12.69025500      11.48337700       5.00482800
+H       11.98627900      11.67412000       4.89783200
+H       26.72371500       1.50329200       2.91105000
+H       26.51636700       2.14519300       3.20834400
+H        9.52738300      16.34429400      15.41539000
+H        8.94769200      16.45923000      15.85602700
+H       22.92343200      13.91655100      28.55448400
+H       22.91144200      13.99598900      27.82170900
+H       19.42780900       1.69485400       3.41368800
+H       18.91557000       1.36746900       3.83063400
+H        7.19465000      19.03417100      20.30585200
+H        6.81579600      18.83473100      20.90594100
+H        5.76403400       2.66353000      26.49024900
+H        6.17055300       2.09977300      26.24461900
+H       22.76304500       7.37905700      21.38041300
+H       22.27422500       7.45776800      21.92655800
+H       20.87103900       4.95557800      21.95393100
+H       20.54276300       4.55466100      21.42960800
+O        2.79836800       3.68628500      15.85501600
+O        2.81086000       3.61004300      17.09857400
+O       16.47940800       2.53818900      22.88295600
+O       16.11951200       2.71624200      24.06243800
+O       16.11955600      27.75767800      21.06500300
+O       16.53710300      26.62354600      21.36799700
+O        2.28922000       5.17821400      26.49769500
+O        2.38796000       5.21061100      25.25608000
+O       23.35821600      23.78530000       2.30561500
+O       22.62749600      23.99025500       3.29376900
+O        8.01645400       1.87918100      22.34390100
+O        7.71565300       1.00097600      23.17497300
+O       24.09262600      23.80245300      18.50597300
+O       25.13119000      23.75358600      19.19256200
+O       23.43919000       8.32924800       7.70010800
+O       22.65907800       8.89107400       8.49268800
+O       23.97254400      22.00311700       5.43698700
+O       23.89487400      22.49778100       6.57789900
+O       27.56254300       8.43488800       7.35581200
+O       27.43055900       8.25837300       6.12950500
+O       20.87006100      27.63986300      11.38162200
+O       20.86380900      27.63798600      12.62756100
+O        1.73533900      16.60794700       2.29549200
+O        1.55322100      16.46304200       1.07146500
+O        2.03508400      26.37263200       7.32515600
+O        1.51280300      27.39855300       7.80172200
+O        9.85776800      20.59545000      27.48500700
+O        9.78700000      20.58172200      26.24113800
+O        6.50961300      12.26791300      10.82158600
+O        6.77142400      11.85977400       9.67385600
+O       27.41406500      13.33114600      25.89152900
+O       27.64803100      14.38606200      25.27119400
+O       10.65841500      26.52178800      25.03188900
+O       11.32208500      27.05107500      25.94392100
+O       19.37531500      27.41779100      27.77152100
+O       19.50655800      28.65269600      27.87247500
+O       15.16713800      28.62232100      25.08493400
+O       15.07154300      28.26860400      26.27579700
+O       17.00751400      16.63266600      17.58492000
+O       17.59502600      16.71087600      16.48896400
+O       11.43827100      14.97615400       3.72092700
+O       10.81544400      15.61953100       4.58727500
+O       15.49533900      18.92456600       3.73508000
+O       15.41648300      19.38131500       2.57854700
+O       24.31615500       8.87176200       3.90336400
+O       24.24717100       8.63758100       2.68155900
+O        7.28827800      12.82799800      18.71906300
+O        6.55542500      13.39545800      19.55172300
+O       11.08744900       6.55884500      22.72049900
+O       12.06607400       6.27858300      23.43893700
+O        6.38492100      21.01048900      22.84221300
+O        6.12927500      19.99335000      23.51488000
+O       24.06443800       6.57624000       5.98481500
+O       24.52114300       5.76588200       6.81375800
+O       10.04719800       9.05734900       6.05599500
+O       10.94084300       9.28468500       6.89392400
+O       14.96821700      27.18796600      18.05445400
+O       15.07988500      27.85573600      17.00850000
+O        3.86508500       1.34322900      24.91310100
+O        3.95599400       1.18703300      26.14588000
+O       20.55715100      20.45753800      21.54897500
+O       19.73819000      20.46822400      22.48790900
+O       15.54716700      11.17956000      10.54123500
+O       15.52405700      11.70380700      11.67129500
+O       27.57196100      18.81697000      18.10230200
+O       28.42915400      18.07345000      18.61689400
+O       13.08594500      21.54734200      24.54405200
+O       14.11395800      21.16325000      25.13402600
+O       26.73666800      18.41630100      23.76197500
+O       26.39312500      17.34817200      24.30371700
+O        3.75150300      18.39675300       7.29152900
+O        3.53600500      18.86067200       6.15541900
+O       17.47592700      22.97389700      14.10462800
+O       17.52163900      23.02279100      15.34878500
+O       18.26366400       1.00092300      19.49677400
+O       18.47839600       1.21883800      20.70458600
+O       13.79429800      11.31045500      17.14879000
+O       14.75715400      11.10528500      16.38510100
+O       23.04461100      15.79867400      21.94317000
+O       24.21186100      16.06476100      22.28832400
+O       18.39899300      20.45438100       6.25233500
+O       18.17334100      19.38368000       5.65644400
+O       12.45708500      28.00076400      22.87959600
+O       12.89052700      26.91111300      23.30054100
+O       26.92707100       5.06905500       5.23373400
+O       25.77848800       5.33532200       4.83091500
+O       21.42684100      24.25676000      11.83089700
+O       21.61828400      25.40374200      11.38347100
+O        5.87933600      27.53672900      16.61577900
+O        4.73532800      27.38741900      17.08626700
+O       24.28470500       3.34526900       2.03178600
+O       24.76461000       2.73333800       1.05832000
+O       28.95365000       2.62596100      13.80868600
+O       28.43270400       3.61259000      13.25408200
+O        4.22171400      14.80081000      20.83480200
+O        4.21389800      14.77292800      19.58918200
+O       22.63711100       1.37775300      19.58937200
+O       23.87485500       1.47510200      19.48488400
+O       24.56301300      18.58751900      21.22319300
+O       23.67458500      19.30549400      20.72558500
+O        8.94762000      14.49730700      12.31189400
+O        9.38425200      15.14940800      13.27963600
+O       16.95516300       6.13080300      19.32331700
+O       16.15902300       5.87750900      18.39897300
+O        8.63504700      19.77961300       8.44117000
+O        9.68469700      20.38220800       8.73700800
+O       11.31282500       2.46998800       2.58334500
+O       11.07470700       1.38104200       2.02665600
+O       27.39479700      24.64841500       9.77024200
+O       26.71049700      25.67139800       9.57621600
+O        2.05963700       8.45476400      16.28023600
+O        2.01230900       8.30200000      17.51588500
+O        9.59314000      26.93822800       8.47063300
+O       10.09094600      25.91018600       7.97291300
+O       25.13932300       4.62957000       8.63843500
+O       25.91372600       5.60386200       8.57958800
+O       21.90534500       3.90912800      17.80877300
+O       21.81900900       3.92307100      19.05165600
+O       19.73746700       6.78254100       2.45368300
+O       19.39816400       6.01988000       3.37868500
+O        6.15237600      26.10484900      18.61638000
+O        4.93823500      25.98411400      18.86874900
+O        7.89451900       2.41931000      18.51611600
+O        7.61766700       2.45592400      17.30185900
+O       24.96184600       3.83646000      17.86831300
+O       24.40660700       3.89389800      16.75439300
+O       19.23185100       4.23752100       5.63971900
+O       20.27904100       4.39660700       4.98360000
diff --git a/mlip_arena/tasks/combustion/water.ipynb b/mlip_arena/tasks/combustion/water.ipynb
index 5a39e4e..e890f7c 100644
--- a/mlip_arena/tasks/combustion/water.ipynb
+++ b/mlip_arena/tasks/combustion/water.ipynb
@@ -34,7 +34,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -49,7 +49,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -60,7 +60,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -78,9 +78,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Atoms(symbols='H256O128', pbc=True, cell=[30.0, 30.0, 30.0])\n"
+     ]
+    }
+   ],
    "source": [
     "tolerance = 2.0\n",
     "input_gen = PackmolBoxGen(\n",
@@ -136,7 +144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -155,7 +163,7 @@
       "source ~/.bashrc\n",
       "module load python\n",
       "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
-      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.19:34913 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.19:39737 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
       "\n"
      ]
     },
@@ -165,8 +173,10 @@
      "text": [
       "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
       "Perhaps you already have a cluster running?\n",
-      "Hosting the HTTP server on port 38239 instead\n",
-      "  warnings.warn(\n"
+      "Hosting the HTTP server on port 38693 instead\n",
+      "  warnings.warn(\n",
+      "2024-10-08 03:30:05,696 - distributed.scheduler - ERROR - Task run-b8e0881eb1e8959f855ada9515884ada marked as failed because 4 workers died while trying to run it\n",
+      "2024-10-08 03:30:05,723 - distributed.scheduler - ERROR - Task run-3307095fc41590db704148cc499bad14 marked as failed because 4 workers died while trying to run it\n"
      ]
     }
    ],
@@ -230,7 +240,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -250,7 +260,7 @@
     "            total_time=1000_000,\n",
     "            temperature=[300, 3000, 3000, 300],\n",
     "            pressure=None,\n",
-    "            mb_velocity_seed=0,\n",
+    "            md_velocity_seed=0,\n",
     "            traj_file=Path(REGISTRY[model.name][\"family\"])\n",
     "            / f\"{model.name}_{atoms.get_chemical_formula()}.traj\",\n",
     "            traj_interval=1000,\n",
@@ -264,7 +274,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 5,
    "metadata": {
     "tags": []
    },
@@ -272,24 +282,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.449 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
-       "</pre>\n"
-      ],
-      "text/plain": [
-       "14:50:51.449 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'liberal-beetle'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.453 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4803afc0-0472-4cf3-8d76-794fc218e2c3</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:25:41.097 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'combustion'</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "14:50:51.453 | \u001b[36mINFO\u001b[0m    | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/flow-runs/flow-run/4803afc0-0472-4cf3-8d76-794fc218e2c3\u001b[0m\n"
+       "01:25:41.097 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'green-turaco'\u001b[0m for flow\u001b[1;35m 'combustion'\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -298,11 +295,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.454 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.19:34913\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:25:41.101 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/784dc0e0-9f4f-4320-8bfa-4c5e8d3b35fb</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "14:50:51.454 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to an existing Dask cluster at tcp://128.55.64.19:34913\n"
+       "01:25:41.101 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/784dc0e0-9f4f-4320-8bfa-4c5e8d3b35fb\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -311,11 +308,11 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.460 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - The Dask dashboard is available at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">http://128.55.64.19:38239/status</span>\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">01:25:41.409 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(91d5c261, 'tcp://128.55.64.19:39737', workers=0, threads=0, memory=0 B)\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "14:50:51.460 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - The Dask dashboard is available at \u001b[94mhttp://128.55.64.19:38239/status\u001b[0m\n"
+       "01:25:41.409 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(91d5c261, 'tcp://128.55.64.19:39737', workers=0, threads=0, memory=0 B)\n"
       ]
      },
      "metadata": {},
@@ -324,77 +321,63 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:51.876 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - Encountered exception during execution:\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:30:05.739 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> - Encountered exception during execution: KilledWorker('run-3307095fc41590db704148cc499bad14', &lt;WorkerState 'tcp://128.55.65.2:37007', name: SLURMCluster-1, status: closed, memory: 0, processing: 0&gt;, 3)\n",
        "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 50, in get_call_parameters\n",
-       "    bound_signature = inspect.signature(fn).bind(*call_args, **call_kwargs)\n",
-       "                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3212, in bind\n",
-       "    return self._bind(args, kwargs)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3201, in _bind\n",
-       "    raise TypeError(\n",
-       "TypeError: got an unexpected keyword argument 'mb_velocity_seed'\n",
-       "\n",
-       "During handling of the above exception, another exception occurred:\n",
-       "\n",
-       "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
-       "    result = await flow_call.aresult()\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 327, in aresult\n",
-       "    return await asyncio.wrap_future(self.future)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
-       "    result = self.fn(*self.args, **self.kwargs)\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_997512/1920202683.py\", line 6, in combustion\n",
-       "    future = MD.submit(\n",
-       "             ^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py\", line 811, in submit\n",
-       "    parameters = get_call_parameters(self.fn, args, kwargs)\n",
-       "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 52, in get_call_parameters\n",
-       "    raise ParameterBindError.from_bind_failure(fn, exc, call_args, call_kwargs)\n",
-       "prefect.exceptions.ParameterBindError: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\n",
-       "Function 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
+       "    yield self\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
+       "    engine.call_flow_fn()\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
+       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
+       "    return fn(*args, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in combustion\n",
+       "    return [future.result() for future in futures]\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in &lt;listcomp&gt;\n",
+       "    return [future.result() for future in futures]\n",
+       "            ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 132, in result\n",
+       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
+       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 328, in result\n",
+       "    return self.client.sync(self._result, callback_timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
+       "    raise exc.with_traceback(tb)\n",
+       "distributed.scheduler.KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "14:50:51.876 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - Encountered exception during execution:\n",
-       "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 50, in get_call_parameters\n",
-       "    bound_signature = inspect.signature(fn).bind(*call_args, **call_kwargs)\n",
-       "                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3212, in bind\n",
-       "    return self._bind(args, kwargs)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py\", line 3201, in _bind\n",
-       "    raise TypeError(\n",
-       "TypeError: got an unexpected keyword argument 'mb_velocity_seed'\n",
-       "\n",
-       "During handling of the above exception, another exception occurred:\n",
-       "\n",
+       "03:30:05.739 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'green-turaco'\u001b[0m - Encountered exception during execution: KilledWorker('run-3307095fc41590db704148cc499bad14', <WorkerState 'tcp://128.55.65.2:37007', name: SLURMCluster-1, status: closed, memory: 0, processing: 0>, 3)\n",
        "Traceback (most recent call last):\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py\", line 867, in orchestrate_flow_run\n",
-       "    result = await flow_call.aresult()\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 327, in aresult\n",
-       "    return await asyncio.wrap_future(self.future)\n",
-       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py\", line 352, in _run_sync\n",
-       "    result = self.fn(*self.args, **self.kwargs)\n",
-       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/tmp/ipykernel_997512/1920202683.py\", line 6, in combustion\n",
-       "    future = MD.submit(\n",
-       "             ^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py\", line 811, in submit\n",
-       "    parameters = get_call_parameters(self.fn, args, kwargs)\n",
-       "                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
-       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 52, in get_call_parameters\n",
-       "    raise ParameterBindError.from_bind_failure(fn, exc, call_args, call_kwargs)\n",
-       "prefect.exceptions.ParameterBindError: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\n",
-       "Function 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].\n"
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 652, in run_context\n",
+       "    yield self\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 696, in run_flow_sync\n",
+       "    engine.call_flow_fn()\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py\", line 675, in call_flow_fn\n",
+       "    result = call_with_parameters(self.flow.fn, self.parameters)\n",
+       "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py\", line 206, in call_with_parameters\n",
+       "    return fn(*args, **kwargs)\n",
+       "           ^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in combustion\n",
+       "    return [future.result() for future in futures]\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/tmp/ipykernel_892374/2043615938.py\", line 26, in <listcomp>\n",
+       "    return [future.result() for future in futures]\n",
+       "            ^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py\", line 132, in result\n",
+       "    future_result = self._wrapped_future.result(timeout=timeout)\n",
+       "                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 328, in result\n",
+       "    return self.client.sync(self._result, callback_timeout=timeout)\n",
+       "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+       "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py\", line 336, in _result\n",
+       "    raise exc.with_traceback(tb)\n",
+       "distributed.scheduler.KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\n"
       ]
      },
      "metadata": {},
@@ -403,47 +386,38 @@
     {
      "data": {
       "text/html": [
-       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">14:50:52.106 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'liberal-beetle'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>('Flow run encountered an exception. ParameterBindError: Error binding parameters for function \\'run\\': got an unexpected keyword argument \\'mb_velocity_seed\\'.\\nFunction \\'run\\' has signature \\'atoms: \\'Atoms\\', calculator_name: \\'str | MLIPEnum\\', calculator_kwargs: \\'dict | None\\', dispersion: \\'str | None\\' = None, dispersion_kwargs: \\'dict | None\\' = None, device: \\'str | None\\' = None, ensemble: \"Literal[\\'nve\\', \\'nvt\\', \\'npt\\']\" = \\'nvt\\', dynamics: \\'str | MolecularDynamics\\' = \\'langevin\\', time_step: \\'float | None\\' = None, total_time: \\'float\\' = 1000, temperature: \\'float | Sequence | np.ndarray | None\\' = 300.0, pressure: \\'float | Sequence | np.ndarray | None\\' = None, ase_md_kwargs: \\'dict | None\\' = None, md_velocity_seed: \\'int | None\\' = None, zero_linear_momentum: \\'bool\\' = True, zero_angular_momentum: \\'bool\\' = True, traj_file: \\'str | Path | None\\' = None, traj_interval: \\'int\\' = 1, restart: \\'bool\\' = True\\' but received args: () and kwargs: [\\'atoms\\', \\'calculator_name\\', \\'calculator_kwargs\\', \\'ensemble\\', \\'dynamics\\', \\'time_step\\', \\'ase_md_kwargs\\', \\'total_time\\', \\'temperature\\', \\'pressure\\', \\'mb_velocity_seed\\', \\'traj_file\\', \\'traj_interval\\', \\'restart\\'].')\n",
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">03:30:06.504 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'green-turaco'</span> - Finished in state <span style=\"color: #d70000; text-decoration-color: #d70000\">Failed</span>(\"Flow run encountered an exception: KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://distributed.dask.org/en/stable/killed.html.</span>\")\n",
        "</pre>\n"
       ],
       "text/plain": [
-       "14:50:52.106 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'liberal-beetle'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m('Flow run encountered an exception. ParameterBindError: Error binding parameters for function \\'run\\': got an unexpected keyword argument \\'mb_velocity_seed\\'.\\nFunction \\'run\\' has signature \\'atoms: \\'Atoms\\', calculator_name: \\'str | MLIPEnum\\', calculator_kwargs: \\'dict | None\\', dispersion: \\'str | None\\' = None, dispersion_kwargs: \\'dict | None\\' = None, device: \\'str | None\\' = None, ensemble: \"Literal[\\'nve\\', \\'nvt\\', \\'npt\\']\" = \\'nvt\\', dynamics: \\'str | MolecularDynamics\\' = \\'langevin\\', time_step: \\'float | None\\' = None, total_time: \\'float\\' = 1000, temperature: \\'float | Sequence | np.ndarray | None\\' = 300.0, pressure: \\'float | Sequence | np.ndarray | None\\' = None, ase_md_kwargs: \\'dict | None\\' = None, md_velocity_seed: \\'int | None\\' = None, zero_linear_momentum: \\'bool\\' = True, zero_angular_momentum: \\'bool\\' = True, traj_file: \\'str | Path | None\\' = None, traj_interval: \\'int\\' = 1, restart: \\'bool\\' = True\\' but received args: () and kwargs: [\\'atoms\\', \\'calculator_name\\', \\'calculator_kwargs\\', \\'ensemble\\', \\'dynamics\\', \\'time_step\\', \\'ase_md_kwargs\\', \\'total_time\\', \\'temperature\\', \\'pressure\\', \\'mb_velocity_seed\\', \\'traj_file\\', \\'traj_interval\\', \\'restart\\'].')\n"
+       "03:30:06.504 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'green-turaco'\u001b[0m - Finished in state \u001b[38;5;160mFailed\u001b[0m(\"Flow run encountered an exception: KilledWorker: Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see \u001b[94mhttps://distributed.dask.org/en/stable/killed.html.\u001b[0m\")\n"
       ]
      },
      "metadata": {},
      "output_type": "display_data"
     },
     {
-     "ename": "ParameterBindError",
-     "evalue": "Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\nFunction 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart'].",
+     "ename": "KilledWorker",
+     "evalue": "Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html.",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:50\u001b[0m, in \u001b[0;36mget_call_parameters\u001b[0;34m(fn, call_args, call_kwargs, apply_defaults)\u001b[0m\n\u001b[1;32m     49\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 50\u001b[0m     bound_signature \u001b[38;5;241m=\u001b[39m \u001b[43minspect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msignature\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcall_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcall_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     51\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py:3212\u001b[0m, in \u001b[0;36mSignature.bind\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   3208\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Get a BoundArguments object, that maps the passed `args`\u001b[39;00m\n\u001b[1;32m   3209\u001b[0m \u001b[38;5;124;03mand `kwargs` to the function's signature.  Raises `TypeError`\u001b[39;00m\n\u001b[1;32m   3210\u001b[0m \u001b[38;5;124;03mif the passed arguments can not be bound.\u001b[39;00m\n\u001b[1;32m   3211\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m-> 3212\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/inspect.py:3201\u001b[0m, in \u001b[0;36mSignature._bind\u001b[0;34m(self, args, kwargs, partial)\u001b[0m\n\u001b[1;32m   3200\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 3201\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\n\u001b[1;32m   3202\u001b[0m             \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mgot an unexpected keyword argument \u001b[39m\u001b[38;5;132;01m{arg!r}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m   3203\u001b[0m                 arg\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mnext\u001b[39m(\u001b[38;5;28miter\u001b[39m(kwargs))))\n\u001b[1;32m   3205\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_bound_arguments_cls(\u001b[38;5;28mself\u001b[39m, arguments)\n",
-      "\u001b[0;31mTypeError\u001b[0m: got an unexpected keyword argument 'mb_velocity_seed'",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[0;31mParameterBindError\u001b[0m                        Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1224\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1219\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m task_viz_tracker:\n\u001b[1;32m   1220\u001b[0m     \u001b[38;5;66;03m# this is a subflow, for now return a single task and do not go further\u001b[39;00m\n\u001b[1;32m   1221\u001b[0m     \u001b[38;5;66;03m# we can add support for exploring subflows for tasks in the future.\u001b[39;00m\n\u001b[1;32m   1222\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[0;32m-> 1224\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43menter_flow_run_engine_from_flow_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1225\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1226\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1227\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1228\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1229\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:297\u001b[0m, in \u001b[0;36menter_flow_run_engine_from_flow_call\u001b[0;34m(flow, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    290\u001b[0m     retval \u001b[38;5;241m=\u001b[39m from_async\u001b[38;5;241m.\u001b[39mwait_for_call_in_loop_thread(\n\u001b[1;32m    291\u001b[0m         begin_run,\n\u001b[1;32m    292\u001b[0m         done_callbacks\u001b[38;5;241m=\u001b[39mdone_callbacks,\n\u001b[1;32m    293\u001b[0m         contexts\u001b[38;5;241m=\u001b[39mcontexts,\n\u001b[1;32m    294\u001b[0m     )\n\u001b[1;32m    296\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 297\u001b[0m     retval \u001b[38;5;241m=\u001b[39m \u001b[43mfrom_sync\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwait_for_call_in_loop_thread\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    298\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbegin_run\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    299\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdone_callbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdone_callbacks\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    300\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcontexts\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcontexts\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    301\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    303\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m retval\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/api.py:243\u001b[0m, in \u001b[0;36mfrom_sync.wait_for_call_in_loop_thread\u001b[0;34m(_from_sync__call, timeout, done_callbacks, contexts)\u001b[0m\n\u001b[1;32m    241\u001b[0m     stack\u001b[38;5;241m.\u001b[39menter_context(context)\n\u001b[1;32m    242\u001b[0m waiter\u001b[38;5;241m.\u001b[39mwait()\n\u001b[0;32m--> 243\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:318\u001b[0m, in \u001b[0;36mCall.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    312\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mresult\u001b[39m(\u001b[38;5;28mself\u001b[39m, timeout: Optional[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    313\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    314\u001b[0m \u001b[38;5;124;03m    Wait for the result of the call.\u001b[39;00m\n\u001b[1;32m    315\u001b[0m \n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03m    Not safe for use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    317\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 318\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:179\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    177\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError()\n\u001b[1;32m    178\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;241m==\u001b[39m FINISHED:\n\u001b[0;32m--> 179\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__get_result\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_condition\u001b[38;5;241m.\u001b[39mwait(timeout)\n\u001b[1;32m    183\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_state \u001b[38;5;129;01min\u001b[39;00m [CANCELLED, CANCELLED_AND_NOTIFIED]:\n\u001b[1;32m    184\u001b[0m     \u001b[38;5;66;03m# Raise Prefect cancelled error instead of\u001b[39;00m\n\u001b[1;32m    185\u001b[0m     \u001b[38;5;66;03m# `concurrent.futures._base.CancelledError`\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/concurrent/futures/_base.py:401\u001b[0m, in \u001b[0;36mFuture.__get_result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception:\n\u001b[1;32m    400\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 401\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         \u001b[38;5;66;03m# Break a reference cycle with the exception in self._exception\u001b[39;00m\n\u001b[1;32m    404\u001b[0m         \u001b[38;5;28mself\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:389\u001b[0m, in \u001b[0;36mCall._run_async\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    387\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_async_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    388\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 389\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m coro\n\u001b[1;32m    390\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    391\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    392\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    393\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    394\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/client/utilities.py:78\u001b[0m, in \u001b[0;36minject_client.<locals>.with_injected_client\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01masync\u001b[39;00m \u001b[38;5;28;01mwith\u001b[39;00m context \u001b[38;5;28;01mas\u001b[39;00m new_client:\n\u001b[1;32m     77\u001b[0m     kwargs\u001b[38;5;241m.\u001b[39msetdefault(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mclient\u001b[39m\u001b[38;5;124m\"\u001b[39m, new_client \u001b[38;5;129;01mor\u001b[39;00m client)\n\u001b[0;32m---> 78\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m fn(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:400\u001b[0m, in \u001b[0;36mcreate_then_begin_flow_run\u001b[0;34m(flow, parameters, wait_for, return_type, client, user_thread)\u001b[0m\n\u001b[1;32m    398\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m state\n\u001b[1;32m    399\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresult\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m--> 400\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m state\u001b[38;5;241m.\u001b[39mresult(fetch\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    401\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInvalid return type for flow engine \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mreturn_type\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/states.py:91\u001b[0m, in \u001b[0;36m_get_state_result\u001b[0;34m(state, raise_on_failure)\u001b[0m\n\u001b[1;32m     84\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m UnfinishedRun(\n\u001b[1;32m     85\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRun is in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate\u001b[38;5;241m.\u001b[39mtype\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m state, its result is not available.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     86\u001b[0m     )\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m     89\u001b[0m     state\u001b[38;5;241m.\u001b[39mis_crashed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_failed() \u001b[38;5;129;01mor\u001b[39;00m state\u001b[38;5;241m.\u001b[39mis_cancelled()\n\u001b[1;32m     90\u001b[0m ):\n\u001b[0;32m---> 91\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m get_state_exception(state)\n\u001b[1;32m     93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(state\u001b[38;5;241m.\u001b[39mdata, DataDocument):\n\u001b[1;32m     94\u001b[0m     result \u001b[38;5;241m=\u001b[39m result_from_state_with_data_document(\n\u001b[1;32m     95\u001b[0m         state, raise_on_failure\u001b[38;5;241m=\u001b[39mraise_on_failure\n\u001b[1;32m     96\u001b[0m     )\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/engine.py:867\u001b[0m, in \u001b[0;36morchestrate_flow_run\u001b[0;34m(flow, flow_run, parameters, wait_for, interruptible, client, partial_flow_run_context, user_thread)\u001b[0m\n\u001b[1;32m    862\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    863\u001b[0m             from_async\u001b[38;5;241m.\u001b[39mcall_soon_in_new_thread(\n\u001b[1;32m    864\u001b[0m                 flow_call, timeout\u001b[38;5;241m=\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds\n\u001b[1;32m    865\u001b[0m             )\n\u001b[0;32m--> 867\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m flow_call\u001b[38;5;241m.\u001b[39maresult()\n\u001b[1;32m    869\u001b[0m         waited_for_task_runs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m wait_for_task_runs_and_report_crashes(\n\u001b[1;32m    870\u001b[0m             flow_run_context\u001b[38;5;241m.\u001b[39mtask_run_futures, client\u001b[38;5;241m=\u001b[39mclient\n\u001b[1;32m    871\u001b[0m         )\n\u001b[1;32m    872\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m PausedRun \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    873\u001b[0m     \u001b[38;5;66;03m# could get raised either via utility or by returning Paused from a task run\u001b[39;00m\n\u001b[1;32m    874\u001b[0m     \u001b[38;5;66;03m# if a task run pauses, we set its state as the flow's state\u001b[39;00m\n\u001b[1;32m    875\u001b[0m     \u001b[38;5;66;03m# to preserve reschedule and timeout behavior\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:327\u001b[0m, in \u001b[0;36mCall.aresult\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    321\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    322\u001b[0m \u001b[38;5;124;03mWait for the result of the call.\u001b[39;00m\n\u001b[1;32m    323\u001b[0m \n\u001b[1;32m    324\u001b[0m \u001b[38;5;124;03mFor use from asynchronous contexts.\u001b[39;00m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 327\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mwrap_future(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture)\n\u001b[1;32m    328\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m asyncio\u001b[38;5;241m.\u001b[39mCancelledError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    329\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m CancelledError() \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mexc\u001b[39;00m\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/_internal/concurrency/calls.py:352\u001b[0m, in \u001b[0;36mCall._run_sync\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    350\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39menforce_sync_deadline() \u001b[38;5;28;01mas\u001b[39;00m cancel_scope:\n\u001b[1;32m    351\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 352\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    353\u001b[0m     \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    354\u001b[0m         \u001b[38;5;66;03m# Forget this call's arguments in order to free up any memory\u001b[39;00m\n\u001b[1;32m    355\u001b[0m         \u001b[38;5;66;03m# that may be referenced by them; after a call has happened,\u001b[39;00m\n\u001b[1;32m    356\u001b[0m         \u001b[38;5;66;03m# there's no need to keep a reference to them\u001b[39;00m\n\u001b[1;32m    357\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "Cell \u001b[0;32mIn[6], line 6\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      3\u001b[0m futures \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m model \u001b[38;5;129;01min\u001b[39;00m MLIPEnum:\n\u001b[0;32m----> 6\u001b[0m     future \u001b[38;5;241m=\u001b[39m \u001b[43mMD\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m        \u001b[49m\u001b[43matoms\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43matoms\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcalculator_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcalculator_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m        \u001b[49m\u001b[43mensemble\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnvt\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdynamics\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnose-hoover\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtime_step\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m        \u001b[49m\u001b[43mase_md_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mttime\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m25\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43munits\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpfactor\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_time\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000_000\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtemperature\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m300\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m300\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpressure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmb_velocity_seed\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtraj_file\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mPath\u001b[49m\u001b[43m(\u001b[49m\u001b[43mREGISTRY\u001b[49m\u001b[43m[\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m]\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mfamily\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mname\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43matoms\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_chemical_formula\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m.traj\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtraj_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     21\u001b[0m \u001b[43m        \u001b[49m\u001b[43mrestart\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     22\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [future\u001b[38;5;241m.\u001b[39mresult() \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py:811\u001b[0m, in \u001b[0;36mTask.submit\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m    808\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mengine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m create_autonomous_task_run, enter_task_run_engine\n\u001b[1;32m    810\u001b[0m \u001b[38;5;66;03m# Convert the call args/kwargs to a parameter dict\u001b[39;00m\n\u001b[0;32m--> 811\u001b[0m parameters \u001b[38;5;241m=\u001b[39m \u001b[43mget_call_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    812\u001b[0m return_type \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m return_state \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfuture\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    814\u001b[0m task_viz_tracker \u001b[38;5;241m=\u001b[39m get_task_viz_tracker()\n",
-      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:52\u001b[0m, in \u001b[0;36mget_call_parameters\u001b[0;34m(fn, call_args, call_kwargs, apply_defaults)\u001b[0m\n\u001b[1;32m     50\u001b[0m     bound_signature \u001b[38;5;241m=\u001b[39m inspect\u001b[38;5;241m.\u001b[39msignature(fn)\u001b[38;5;241m.\u001b[39mbind(\u001b[38;5;241m*\u001b[39mcall_args, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcall_kwargs)\n\u001b[1;32m     51\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[0;32m---> 52\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m ParameterBindError\u001b[38;5;241m.\u001b[39mfrom_bind_failure(fn, exc, call_args, call_kwargs)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m apply_defaults:\n\u001b[1;32m     55\u001b[0m     bound_signature\u001b[38;5;241m.\u001b[39mapply_defaults()\n",
-      "\u001b[0;31mParameterBindError\u001b[0m: Error binding parameters for function 'run': got an unexpected keyword argument 'mb_velocity_seed'.\nFunction 'run' has signature 'atoms: 'Atoms', calculator_name: 'str | MLIPEnum', calculator_kwargs: 'dict | None', dispersion: 'str | None' = None, dispersion_kwargs: 'dict | None' = None, device: 'str | None' = None, ensemble: \"Literal['nve', 'nvt', 'npt']\" = 'nvt', dynamics: 'str | MolecularDynamics' = 'langevin', time_step: 'float | None' = None, total_time: 'float' = 1000, temperature: 'float | Sequence | np.ndarray | None' = 300.0, pressure: 'float | Sequence | np.ndarray | None' = None, ase_md_kwargs: 'dict | None' = None, md_velocity_seed: 'int | None' = None, zero_linear_momentum: 'bool' = True, zero_angular_momentum: 'bool' = True, traj_file: 'str | Path | None' = None, traj_interval: 'int' = 1, restart: 'bool' = True' but received args: () and kwargs: ['atoms', 'calculator_name', 'calculator_kwargs', 'ensemble', 'dynamics', 'time_step', 'ase_md_kwargs', 'total_time', 'temperature', 'pressure', 'mb_velocity_seed', 'traj_file', 'traj_interval', 'restart']."
+      "\u001b[0;31mKilledWorker\u001b[0m                              Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[5], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mcombustion\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1345\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1341\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[1;32m   1343\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow_engine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m run_flow\n\u001b[0;32m-> 1345\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1350\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:818\u001b[0m, in \u001b[0;36mrun_flow\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_flow_async(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 818\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:698\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[1;32m    696\u001b[0m             engine\u001b[38;5;241m.\u001b[39mcall_flow_fn()\n\u001b[0;32m--> 698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:255\u001b[0m, in \u001b[0;36mFlowRunEngine.result\u001b[0;34m(self, raise_on_failure)\u001b[0m\n\u001b[1;32m    253\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m NotSet:\n\u001b[1;32m    254\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m raise_on_failure:\n\u001b[0;32m--> 255\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    256\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_raised\n\u001b[1;32m    258\u001b[0m \u001b[38;5;66;03m# This is a fall through case which leans on the existing state result mechanics to get the\u001b[39;00m\n\u001b[1;32m    259\u001b[0m \u001b[38;5;66;03m# return value. This is necessary because we currently will return a State object if the\u001b[39;00m\n\u001b[1;32m    260\u001b[0m \u001b[38;5;66;03m# the State was Prefect-created.\u001b[39;00m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;66;03m# TODO: Remove the need to get the result from a State except in cases where the return value\u001b[39;00m\n\u001b[1;32m    262\u001b[0m \u001b[38;5;66;03m# is a State object.\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:652\u001b[0m, in \u001b[0;36mFlowRunEngine.run_context\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    645\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m timeout_context(\n\u001b[1;32m    646\u001b[0m         seconds\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mtimeout_seconds,\n\u001b[1;32m    647\u001b[0m         timeout_exc_type\u001b[38;5;241m=\u001b[39mFlowRunTimeoutError,\n\u001b[1;32m    648\u001b[0m     ):\n\u001b[1;32m    649\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\n\u001b[1;32m    650\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExecuting flow \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m for flow run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mflow_run\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m...\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    651\u001b[0m         )\n\u001b[0;32m--> 652\u001b[0m         \u001b[38;5;28;01myield\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    654\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_timeout(exc)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:696\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    694\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mis_running():\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[0;32m--> 696\u001b[0m             \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_flow_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:675\u001b[0m, in \u001b[0;36mFlowRunEngine.call_flow_fn\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    673\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _call_flow_fn()\n\u001b[1;32m    674\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 675\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mcall_with_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    676\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m(fn, parameters)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36mcombustion\u001b[0;34m(atoms)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m[\u001b[49m\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfuture\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfutures\u001b[49m\u001b[43m]\u001b[49m\n",
+      "Cell \u001b[0;32mIn[4], line 26\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      6\u001b[0m     future \u001b[38;5;241m=\u001b[39m MD\u001b[38;5;241m.\u001b[39msubmit(\n\u001b[1;32m      7\u001b[0m         atoms\u001b[38;5;241m=\u001b[39matoms,\n\u001b[1;32m      8\u001b[0m         calculator_name\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m         restart\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     22\u001b[0m     )\n\u001b[1;32m     24\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[0;32m---> 26\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m [\u001b[43mfuture\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m future \u001b[38;5;129;01min\u001b[39;00m futures]\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:132\u001b[0m, in \u001b[0;36mPrefectDaskFuture.result\u001b[0;34m(self, timeout, raise_on_failure)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_final_state:\n\u001b[1;32m    131\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 132\u001b[0m         future_result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_wrapped_future\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtimeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    133\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m distributed\u001b[38;5;241m.\u001b[39mTimeoutError \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m    134\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTimeoutError\u001b[39;00m(\n\u001b[1;32m    135\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTask run \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtask_run_id\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m did not complete within \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtimeout\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m seconds\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    136\u001b[0m         ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mexc\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:328\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self, timeout)\u001b[0m\n\u001b[1;32m    326\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_verify_initialized()\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m shorten_traceback():\n\u001b[0;32m--> 328\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclient\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_result\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtimeout\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:336\u001b[0m, in \u001b[0;36mFuture._result\u001b[0;34m(self, raiseit)\u001b[0m\n\u001b[1;32m    334\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m raiseit:\n\u001b[1;32m    335\u001b[0m     typ, exc, tb \u001b[38;5;241m=\u001b[39m exc\n\u001b[0;32m--> 336\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mwith_traceback(tb)\n\u001b[1;32m    337\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    338\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m exc\n",
+      "\u001b[0;31mKilledWorker\u001b[0m: Attempted to run task 'run-3307095fc41590db704148cc499bad14' on 4 different workers, but all those workers died while running it. The last worker that attempt to run the task was tcp://128.55.65.2:37007. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html."
      ]
     }
    ],
diff --git a/mlip_arena/tasks/eos/run.ipynb b/mlip_arena/tasks/eos/run.ipynb
new file mode 100644
index 0000000..a5ae894
--- /dev/null
+++ b/mlip_arena/tasks/eos/run.ipynb
@@ -0,0 +1,645 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from mp_api.client import MPRester\n",
+    "from dask.distributed import Client\n",
+    "from dask_jobqueue import SLURMCluster\n",
+    "from prefect import task, flow\n",
+    "from prefect_dask import DaskTaskRunner\n",
+    "from pymatgen.core.structure import Structure\n",
+    "from dotenv import load_dotenv\n",
+    "from ase import Atoms\n",
+    "from ase.io import write, read\n",
+    "from pathlib import Path\n",
+    "import pandas as pd\n",
+    "from prefect.futures import wait\n",
+    "\n",
+    "from mlip_arena.tasks.eos.run import fit as EOS\n",
+    "from mlip_arena.models.utils import REGISTRY, MLIPEnum\n",
+    "\n",
+    "load_dotenv()\n",
+    "\n",
+    "MP_API_KEY = os.environ.get(\"MP_API_KEY\", None)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "MP Database version: 2023.11.1\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bf1668eac663453aa9a20149ecf726c4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Retrieving SummaryDoc documents:   0%|          | 0/5135 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "\n",
+    "with MPRester(MP_API_KEY) as mpr:\n",
+    "    print(\"MP Database version:\", mpr.get_database_version())\n",
+    "\n",
+    "    summary_docs = mpr.materials.summary.search(\n",
+    "        num_elements=(1, 2),\n",
+    "        is_stable=True,\n",
+    "        fields=[\"material_id\", \"structure\", \"formula_pretty\"]\n",
+    "    )\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "atoms_list = []\n",
+    "\n",
+    "for doc in summary_docs:\n",
+    "\n",
+    "    structure = doc.structure\n",
+    "    assert isinstance(structure, Structure)\n",
+    "\n",
+    "    atoms = structure.to_ase_atoms()\n",
+    "\n",
+    "    atoms_list.append(atoms)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "write(\"all.extxyz\", atoms_list)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "atoms_list = read(\"all.extxyz\", index=':')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "#!/bin/bash\n",
+      "\n",
+      "#SBATCH -A matgen\n",
+      "#SBATCH --mem=0\n",
+      "#SBATCH -t 00:30:00\n",
+      "#SBATCH -N 1\n",
+      "#SBATCH -G 4\n",
+      "#SBATCH -q debug\n",
+      "#SBATCH -C gpu\n",
+      "#SBATCH -J eos\n",
+      "source ~/.bashrc\n",
+      "module load python\n",
+      "source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\n",
+      "/pscratch/sd/c/cyrusyc/.conda/mlip-arena/bin/python -m distributed.cli.dask_worker tcp://128.55.64.41:45323 --name dummy-name --nthreads 1 --memory-limit 59.60GiB --nanny --death-timeout 60\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Exception in thread Profile:\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/threading.py\", line 1045, in _bootstrap_inner\n",
+      "    self.run()\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/ipykernel/ipkernel.py\", line 766, in run_closure\n",
+      "    _threading_Thread_run(self)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/threading.py\", line 982, in run\n",
+      "    self._target(*self._args, **self._kwargs)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 366, in _watch\n",
+      "    process(frame, None, recent, omit=omit)\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 183, in process\n",
+      "    if any(frame.f_code.co_filename.endswith(o) for o in omit):\n",
+      "       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "  File \"/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/profile.py\", line 183, in <genexpr>\n",
+      "    if any(frame.f_code.co_filename.endswith(o) for o in omit):\n",
+      "           ^^^^^^^^^^^^^^^^^^^^^^^^\n",
+      "AttributeError: 'dict' object has no attribute 'co_filename'\n"
+     ]
+    }
+   ],
+   "source": [
+    "nodes_per_alloc = 1\n",
+    "gpus_per_alloc = 4\n",
+    "ntasks = 1\n",
+    "\n",
+    "cluster_kwargs = {\n",
+    "    \"cores\": 1,\n",
+    "    \"memory\": \"64 GB\",\n",
+    "    \"shebang\": \"#!/bin/bash\",\n",
+    "    \"account\": \"matgen\",\n",
+    "    \"walltime\": \"00:30:00\",\n",
+    "    \"job_mem\": \"0\",\n",
+    "    \"job_script_prologue\": [\n",
+    "        \"source ~/.bashrc\",\n",
+    "        \"module load python\",\n",
+    "        \"source activate /pscratch/sd/c/cyrusyc/.conda/mlip-arena\",\n",
+    "    ],\n",
+    "    \"job_directives_skip\": [\"-n\", \"--cpus-per-task\", \"-J\"],\n",
+    "    \"job_extra_directives\": [f\"-N {nodes_per_alloc}\", f\"-G {gpus_per_alloc}\", \"-q debug\", \"-C gpu\", \"-J eos\"],\n",
+    "}\n",
+    "cluster = SLURMCluster(**cluster_kwargs)\n",
+    "\n",
+    "print(cluster.job_script())\n",
+    "cluster.adapt(minimum_jobs=2, maximum_jobs=2)\n",
+    "client = Client(cluster)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "def postprocess(output, model: str, formula: str):\n",
+    "    row = {\n",
+    "        \"formula\": formula,\n",
+    "        \"method\": model,\n",
+    "        \"volumes\": output[\"eos\"][\"volumes\"],\n",
+    "        \"energies\": output[\"eos\"][\"energies\"],\n",
+    "        \"K\": output[\"K\"],\n",
+    "    }\n",
+    "\n",
+    "    fpath = Path(REGISTRY[model][\"family\"]) / f\"{model}.parquet\"\n",
+    "\n",
+    "    if not fpath.exists():\n",
+    "        fpath.parent.mkdir(parents=True, exist_ok=True)\n",
+    "        df = pd.DataFrame([row])  # Convert the dictionary to a DataFrame with a list\n",
+    "    else:\n",
+    "        df = pd.read_parquet(fpath)\n",
+    "        new_row = pd.DataFrame([row])  # Convert dictionary to DataFrame with a list\n",
+    "        df = pd.concat([df, new_row], ignore_index=True)\n",
+    "\n",
+    "    df.drop_duplicates(subset=[\"formula\", \"method\"], keep='last', inplace=True)\n",
+    "    df.to_parquet(fpath)\n",
+    "\n",
+    "from prefect.runtime import flow_run, task_run\n",
+    "EOS = EOS.with_options(timeout_seconds=240, result_storage=None)\n",
+    "\n",
+    "def generate_task_run_name():\n",
+    "    task_name = task_run.task_name\n",
+    "\n",
+    "    parameters = task_run.parameters\n",
+    "\n",
+    "    atoms = parameters[\"atoms\"]\n",
+    "    \n",
+    "    return f\"{task_name}: {atoms.get_chemical_formula()}\"\n",
+    "\n",
+    "\n",
+    "@task(task_run_name=generate_task_run_name, result_storage=None)\n",
+    "def fit_one(atoms: Atoms):\n",
+    "    \n",
+    "    outputs = []\n",
+    "    for model in MLIPEnum:\n",
+    "        try:\n",
+    "            eos = EOS(\n",
+    "                atoms=atoms,\n",
+    "                calculator_name=model.name,\n",
+    "                calculator_kwargs={},\n",
+    "                device=None,\n",
+    "                optimizer=\"QuasiNewton\",\n",
+    "                optimizer_kwargs=None,\n",
+    "                filter=\"FrechetCell\",\n",
+    "                filter_kwargs=None,\n",
+    "                criterion=dict(\n",
+    "                    fmax=0.1,\n",
+    "                ),\n",
+    "                max_abs_strain=0.1,\n",
+    "                npoints=7,\n",
+    "            )\n",
+    "            if isinstance(eos, dict):\n",
+    "                postprocess(output=eos, model=model.name, formula=atoms.get_chemical_formula())\n",
+    "                eos[\"method\"] = model.name\n",
+    "                outputs.append(eos)\n",
+    "        except:\n",
+    "            continue\n",
+    "    \n",
+    "    return outputs\n",
+    "\n",
+    "\n",
+    "@flow(task_runner=DaskTaskRunner(address=client.scheduler.address), log_prints=True, result_storage=None)\n",
+    "def fit_all(atoms_list: list[Atoms]):\n",
+    "    \n",
+    "    futures = []\n",
+    "    for atoms in atoms_list:\n",
+    "        future = fit_one.submit(atoms)\n",
+    "        futures.append(future)\n",
+    "            \n",
+    "    wait(futures)\n",
+    "    \n",
+    "    return [f.result(raise_on_failure=False) for f in futures]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:51:45.507 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - Created flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> for flow<span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\"> 'fit-all'</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "00:51:45.507 | \u001b[36mINFO\u001b[0m    | prefect.engine - Created flow run\u001b[35m 'thistle-pronghorn'\u001b[0m for flow\u001b[1;35m 'fit-all'\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:51:45.509 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.engine - View at <span style=\"color: #0000ff; text-decoration-color: #0000ff\">https://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/9b293dae-4c2d-43b7-b8bb-9a4749eee74e</span>\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "00:51:45.509 | \u001b[36mINFO\u001b[0m    | prefect.engine - View at \u001b[94mhttps://app.prefect.cloud/account/f7d40474-9362-4bfa-8950-ee6a43ec00f3/workspace/d4bb0913-5f5e-49f7-bfc5-06509088baeb/runs/flow-run/9b293dae-4c2d-43b7-b8bb-9a4749eee74e\u001b[0m\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:51:45.818 | <span style=\"color: #008080; text-decoration-color: #008080\">INFO</span>    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(5d2876ed, 'tcp://128.55.64.41:45323', workers=0, threads=0, memory=0 B)\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "00:51:45.818 | \u001b[36mINFO\u001b[0m    | prefect.task_runner.dask - Connecting to existing Dask cluster SLURMCluster(5d2876ed, 'tcp://128.55.64.41:45323', workers=0, threads=0, memory=0 B)\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:58:09.352 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> - Crash detected! Execution was aborted by an interrupt signal.\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "00:58:09.352 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'thistle-pronghorn'\u001b[0m - Crash detected! Execution was aborted by an interrupt signal.\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">00:58:09.642 | <span style=\"color: #d70000; text-decoration-color: #d70000\">ERROR</span>   | Flow run<span style=\"color: #800080; text-decoration-color: #800080\"> 'thistle-pronghorn'</span> - Finished in state <span style=\"color: #ff0000; text-decoration-color: #ff0000\">Crashed</span>('Execution was aborted by an interrupt signal.')\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "00:58:09.642 | \u001b[38;5;160mERROR\u001b[0m   | Flow run\u001b[35m 'thistle-pronghorn'\u001b[0m - Finished in state \u001b[91mCrashed\u001b[0m('Execution was aborted by an interrupt signal.')\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[6], line 21\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# import os\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# import tempfile\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m# import shutil\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     18\u001b[0m \n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# with twd():\u001b[39;00m\n\u001b[0;32m---> 21\u001b[0m \u001b[43mfit_all\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms_list\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flows.py:1345\u001b[0m, in \u001b[0;36mFlow.__call__\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1341\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m track_viz_task(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39misasync, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname, parameters)\n\u001b[1;32m   1343\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mprefect\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow_engine\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m run_flow\n\u001b[0;32m-> 1345\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1346\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1347\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_type\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1350\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:818\u001b[0m, in \u001b[0;36mrun_flow\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    816\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m run_flow_async(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    817\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 818\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrun_flow_sync\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:696\u001b[0m, in \u001b[0;36mrun_flow_sync\u001b[0;34m(flow, flow_run, parameters, wait_for, return_type)\u001b[0m\n\u001b[1;32m    694\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mis_running():\n\u001b[1;32m    695\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mrun_context():\n\u001b[0;32m--> 696\u001b[0m             \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_flow_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    698\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mstate \u001b[38;5;28;01mif\u001b[39;00m return_type \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m engine\u001b[38;5;241m.\u001b[39mresult()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/flow_engine.py:675\u001b[0m, in \u001b[0;36mFlowRunEngine.call_flow_fn\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    673\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _call_flow_fn()\n\u001b[1;32m    674\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 675\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[43mcall_with_parameters\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mflow\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    676\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandle_success(result)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/utilities/callables.py:206\u001b[0m, in \u001b[0;36mcall_with_parameters\u001b[0;34m(fn, parameters)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;124;03mCall a function with parameters extracted with `get_call_parameters`\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    203\u001b[0m \u001b[38;5;124;03mthe args/kwargs using `parameters_to_positional_and_keyword` directly\u001b[39;00m\n\u001b[1;32m    204\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    205\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m parameters_to_args_kwargs(fn, parameters)\n\u001b[0;32m--> 206\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[5], line 72\u001b[0m, in \u001b[0;36mfit_all\u001b[0;34m(atoms_list)\u001b[0m\n\u001b[1;32m     70\u001b[0m futures \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m     71\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m atoms \u001b[38;5;129;01min\u001b[39;00m atoms_list:\n\u001b[0;32m---> 72\u001b[0m     future \u001b[38;5;241m=\u001b[39m \u001b[43mfit_one\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\u001b[43matoms\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     73\u001b[0m     futures\u001b[38;5;241m.\u001b[39mappend(future)\n\u001b[1;32m     75\u001b[0m wait(futures)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect/tasks.py:1163\u001b[0m, in \u001b[0;36mTask.submit\u001b[0;34m(self, return_state, wait_for, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1158\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m VisualizationUnsupportedError(\n\u001b[1;32m   1159\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`task.submit()` is not currently supported by `flow.visualize()`\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1160\u001b[0m     )\n\u001b[1;32m   1162\u001b[0m task_runner \u001b[38;5;241m=\u001b[39m flow_run_context\u001b[38;5;241m.\u001b[39mtask_runner\n\u001b[0;32m-> 1163\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[43mtask_runner\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_state:\n\u001b[1;32m   1165\u001b[0m     future\u001b[38;5;241m.\u001b[39mwait()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/task_runners.py:353\u001b[0m, in \u001b[0;36mDaskTaskRunner.submit\u001b[0;34m(self, task, parameters, wait_for, dependencies)\u001b[0m\n\u001b[1;32m    349\u001b[0m \u001b[38;5;66;03m# unpack the upstream call in order to cast Prefect futures to Dask futures\u001b[39;00m\n\u001b[1;32m    350\u001b[0m \u001b[38;5;66;03m# where possible to optimize Dask task scheduling\u001b[39;00m\n\u001b[1;32m    351\u001b[0m parameters \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_optimize_futures(parameters)\n\u001b[0;32m--> 353\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_client\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    354\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    355\u001b[0m \u001b[43m    \u001b[49m\u001b[43mparameters\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    356\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwait_for\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mwait_for\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    357\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdependencies\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdependencies\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    358\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstate\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    359\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    360\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m PrefectDaskFuture(wrapped_future\u001b[38;5;241m=\u001b[39mfuture, task_run_id\u001b[38;5;241m=\u001b[39mfuture\u001b[38;5;241m.\u001b[39mtask_run_id)\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/prefect_dask/client.py:64\u001b[0m, in \u001b[0;36mPrefectDaskClient.submit\u001b[0;34m(self, func, key, workers, resources, retries, priority, fifo_timeout, allow_other_workers, actor, actors, pure, *args, **kwargs)\u001b[0m\n\u001b[1;32m     61\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     62\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m run_task_sync(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m---> 64\u001b[0m future \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msubmit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     65\u001b[0m \u001b[43m    \u001b[49m\u001b[43mwrapper_func\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     66\u001b[0m \u001b[43m    \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     67\u001b[0m \u001b[43m    \u001b[49m\u001b[43mworkers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mworkers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     68\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresources\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresources\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretries\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretries\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpriority\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpriority\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfifo_timeout\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfifo_timeout\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_other_workers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mallow_other_workers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m    \u001b[49m\u001b[43mactor\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mactor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[43m    \u001b[49m\u001b[43mactors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mactors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     75\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpure\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     76\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mrun_task_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     77\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m future\u001b[38;5;241m.\u001b[39mtask_run_id \u001b[38;5;241m=\u001b[39m run_task_kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtask_run_id\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m future\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/distributed/client.py:1931\u001b[0m, in \u001b[0;36mClient.submit\u001b[0;34m(self, func, key, workers, resources, retries, priority, fifo_timeout, allow_other_workers, actor, actors, pure, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1929\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m key \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1930\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m pure:\n\u001b[0;32m-> 1931\u001b[0m         key \u001b[38;5;241m=\u001b[39m funcname(func) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m-\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[43mtokenize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1932\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1933\u001b[0m         key \u001b[38;5;241m=\u001b[39m funcname(func) \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m-\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(uuid\u001b[38;5;241m.\u001b[39muuid4())\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1033\u001b[0m, in \u001b[0;36mtokenize\u001b[0;34m(ensure_deterministic, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1015\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Deterministic token\u001b[39;00m\n\u001b[1;32m   1016\u001b[0m \n\u001b[1;32m   1017\u001b[0m \u001b[38;5;124;03m>>> tokenize([1, 2, '3'])\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1030\u001b[0m \u001b[38;5;124;03m    Defaults to the `tokenize.ensure-deterministic` configuration parameter.\u001b[39;00m\n\u001b[1;32m   1031\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1032\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m _seen_ctx(reset\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m), _ensure_deterministic_ctx(ensure_deterministic):\n\u001b[0;32m-> 1033\u001b[0m     token: \u001b[38;5;28mobject\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1034\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m kwargs:\n\u001b[1;32m   1035\u001b[0m         token \u001b[38;5;241m=\u001b[39m token, _normalize_seq_func(\u001b[38;5;28msorted\u001b[39m(kwargs\u001b[38;5;241m.\u001b[39mitems()))\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
+      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (1 times), _normalize_seq_func at line 1159 (1 times)]\u001b[0m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
+      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (1 times), _normalize_seq_func at line 1159 (1 times)]\u001b[0m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
+      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (3 times), _normalize_seq_func at line 1159 (3 times), normalize_dict at line 1129 (1 times), normalize_seq at line 1166 (1 times)]\u001b[0m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1129\u001b[0m, in \u001b[0;36mnormalize_dict\u001b[0;34m(d)\u001b[0m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((types\u001b[38;5;241m.\u001b[39mMappingProxyType, \u001b[38;5;28mdict\u001b[39m))\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_dict\u001b[39m(d):\n\u001b[0;32m-> 1129\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdict\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43msorted\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43md\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitems\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mkv\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mstr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mkv\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
+      "    \u001b[0;31m[... skipping similar frames: Dispatch.__call__ at line 773 (2 times), _normalize_seq_func at line 1159 (2 times), normalize_seq at line 1166 (1 times)]\u001b[0m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1166\u001b[0m, in \u001b[0;36mnormalize_seq\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;129m@normalize_token\u001b[39m\u001b[38;5;241m.\u001b[39mregister((\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m))\n\u001b[1;32m   1165\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mnormalize_seq\u001b[39m(seq):\n\u001b[0;32m-> 1166\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(seq)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, \u001b[43m_normalize_seq_func\u001b[49m\u001b[43m(\u001b[49m\u001b[43mseq\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1159\u001b[0m, in \u001b[0;36m_normalize_seq_func\u001b[0;34m(seq)\u001b[0m\n\u001b[1;32m   1157\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1158\u001b[0m         seen[\u001b[38;5;28mid\u001b[39m(item)] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(seen), item\n\u001b[0;32m-> 1159\u001b[0m         item \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_token\u001b[49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1160\u001b[0m     out\u001b[38;5;241m.\u001b[39mappend(item)\n\u001b[1;32m   1161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/utils.py:773\u001b[0m, in \u001b[0;36mDispatch.__call__\u001b[0;34m(self, arg, *args, **kwargs)\u001b[0m\n\u001b[1;32m    769\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    770\u001b[0m \u001b[38;5;124;03mCall the corresponding method based on type of argument.\u001b[39;00m\n\u001b[1;32m    771\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    772\u001b[0m meth \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch(\u001b[38;5;28mtype\u001b[39m(arg))\n\u001b[0;32m--> 773\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1212\u001b[0m, in \u001b[0;36mnormalize_object\u001b[0;34m(o)\u001b[0m\n\u001b[1;32m   1209\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _normalize_dataclass(o)\n\u001b[1;32m   1211\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 1212\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_normalize_pickle\u001b[49m\u001b[43m(\u001b[49m\u001b[43mo\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1213\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n\u001b[1;32m   1214\u001b[0m     _maybe_raise_nondeterministic(\n\u001b[1;32m   1215\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mObject \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mo\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m cannot be deterministically hashed. See \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1216\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhttps://docs.dask.org/en/latest/custom-collections.html#implementing-deterministic-hashing \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1217\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfor more information.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   1218\u001b[0m     )\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/base.py:1250\u001b[0m, in \u001b[0;36m_normalize_pickle\u001b[0;34m(o)\u001b[0m\n\u001b[1;32m   1247\u001b[0m     buffers\u001b[38;5;241m.\u001b[39mclear()\n\u001b[1;32m   1248\u001b[0m     pik \u001b[38;5;241m=\u001b[39m cloudpickle\u001b[38;5;241m.\u001b[39mdumps(o, protocol\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m5\u001b[39m, buffer_callback\u001b[38;5;241m=\u001b[39mbuffers\u001b[38;5;241m.\u001b[39mappend)\n\u001b[0;32m-> 1250\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhash_buffer_hex\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpik\u001b[49m\u001b[43m)\u001b[49m, [hash_buffer_hex(buf) \u001b[38;5;28;01mfor\u001b[39;00m buf \u001b[38;5;129;01min\u001b[39;00m buffers]\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:98\u001b[0m, in \u001b[0;36mhash_buffer_hex\u001b[0;34m(buf, hasher)\u001b[0m\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mhash_buffer_hex\u001b[39m(buf, hasher\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m     95\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     96\u001b[0m \u001b[38;5;124;03m    Same as hash_buffer, but returns its result in hex-encoded form.\u001b[39;00m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 98\u001b[0m     h \u001b[38;5;241m=\u001b[39m \u001b[43mhash_buffer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhasher\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     99\u001b[0m     s \u001b[38;5;241m=\u001b[39m binascii\u001b[38;5;241m.\u001b[39mb2a_hex(h)\n\u001b[1;32m    100\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m s\u001b[38;5;241m.\u001b[39mdecode()\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:88\u001b[0m, in \u001b[0;36mhash_buffer\u001b[0;34m(buf, hasher)\u001b[0m\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m hasher \u001b[38;5;129;01min\u001b[39;00m hashers:\n\u001b[1;32m     87\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 88\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhasher\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mTypeError\u001b[39;00m, \u001b[38;5;167;01mOverflowError\u001b[39;00m):\n\u001b[1;32m     90\u001b[0m         \u001b[38;5;28;01mpass\u001b[39;00m\n",
+      "File \u001b[0;32m/pscratch/sd/c/cyrusyc/.conda/mlip-arena/lib/python3.11/site-packages/dask/hashing.py:44\u001b[0m, in \u001b[0;36m_hash_xxhash\u001b[0;34m(buf)\u001b[0m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_hash_xxhash\u001b[39m(buf):\n\u001b[1;32m     41\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;124;03m    Produce a 8-bytes hash of *buf* using xxHash.\u001b[39;00m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 44\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mxxhash\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxxh64\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbuf\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mdigest()\n",
+      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "# import os\n",
+    "# import tempfile\n",
+    "# import shutil\n",
+    "# from contextlib import contextmanager\n",
+    "\n",
+    "# @contextmanager\n",
+    "# def twd():\n",
+    "    \n",
+    "#     pwd = os.getcwd()\n",
+    "#     temp_dir = tempfile.mkdtemp()\n",
+    "    \n",
+    "#     try:\n",
+    "#         os.chdir(temp_dir)\n",
+    "#         yield\n",
+    "#     finally:\n",
+    "#         os.chdir(pwd)\n",
+    "#         shutil.rmtree(temp_dir)\n",
+    "\n",
+    "# with twd():\n",
+    "\n",
+    "fit_all(atoms_list)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "\n",
+    "df = pd.read_parquet('mace-mp/MACE-MP(M).parquet')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>formula</th>\n",
+       "      <th>method</th>\n",
+       "      <th>volumes</th>\n",
+       "      <th>energies</th>\n",
+       "      <th>K</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>Ac2O3</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[82.36010147441682, 85.41047560309894, 88.4608...</td>\n",
+       "      <td>[-39.47665786743164, -39.65583419799805, -39.7...</td>\n",
+       "      <td>95.869141</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>Ac4</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[166.09086069175856, 172.2423740507126, 178.39...</td>\n",
+       "      <td>[-16.326059341430664, -16.406923294067383, -16...</td>\n",
+       "      <td>25.409891</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Ac16S24</td>\n",
+       "      <td>MACE-MP(M)</td>\n",
+       "      <td>[1006.5670668063424, 1043.84732853991, 1081.12...</td>\n",
+       "      <td>[-249.42129516601562, -250.79556274414062, -25...</td>\n",
+       "      <td>61.372858</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   formula      method                                            volumes  \\\n",
+       "0    Ac2O3  MACE-MP(M)  [82.36010147441682, 85.41047560309894, 88.4608...   \n",
+       "1      Ac4  MACE-MP(M)  [166.09086069175856, 172.2423740507126, 178.39...   \n",
+       "3  Ac16S24  MACE-MP(M)  [1006.5670668063424, 1043.84732853991, 1081.12...   \n",
+       "\n",
+       "                                            energies          K  \n",
+       "0  [-39.47665786743164, -39.65583419799805, -39.7...  95.869141  \n",
+       "1  [-16.326059341430664, -16.406923294067383, -16...  25.409891  \n",
+       "3  [-249.42129516601562, -250.79556274414062, -25...  61.372858  "
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mlip-arena",
+   "language": "python",
+   "name": "mlip-arena"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "state": {
+     "1e2c24979c4e4b3b91e8f8a8e214d56f": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_e03d78a760b8402cad0827279f4075c1",
+       "style": "IPY_MODEL_743568418ec24e369c2ac0f9a9573a1c",
+       "value": "Retrieving SummaryDoc documents: 100%"
+      }
+     },
+     "41e363c9f8514e6a993426e30d45ec98": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "4dde6375de954082a87b1483182e40a3": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "607687cda1ce4276a772b626ef3e6179": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLModel",
+      "state": {
+       "layout": "IPY_MODEL_93a95097ee984d95b1dfedcb2d913435",
+       "style": "IPY_MODEL_41e363c9f8514e6a993426e30d45ec98",
+       "value": " 5135/5135 [00:04&lt;00:00, 958.06it/s]"
+      }
+     },
+     "743568418ec24e369c2ac0f9a9573a1c": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HTMLStyleModel",
+      "state": {
+       "description_width": "",
+       "font_size": null,
+       "text_color": null
+      }
+     },
+     "8d602e7a8e734c50b0fad54bf76cc7d5": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "ProgressStyleModel",
+      "state": {
+       "description_width": ""
+      }
+     },
+     "93a95097ee984d95b1dfedcb2d913435": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "b09359f88f954062aff30df3bf2ba46b": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     },
+     "bf1668eac663453aa9a20149ecf726c4": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "HBoxModel",
+      "state": {
+       "children": [
+        "IPY_MODEL_1e2c24979c4e4b3b91e8f8a8e214d56f",
+        "IPY_MODEL_cf1131b45bf04178ac93962259718700",
+        "IPY_MODEL_607687cda1ce4276a772b626ef3e6179"
+       ],
+       "layout": "IPY_MODEL_4dde6375de954082a87b1483182e40a3"
+      }
+     },
+     "cf1131b45bf04178ac93962259718700": {
+      "model_module": "@jupyter-widgets/controls",
+      "model_module_version": "2.0.0",
+      "model_name": "FloatProgressModel",
+      "state": {
+       "bar_style": "success",
+       "layout": "IPY_MODEL_b09359f88f954062aff30df3bf2ba46b",
+       "max": 5135,
+       "style": "IPY_MODEL_8d602e7a8e734c50b0fad54bf76cc7d5",
+       "value": 5135
+      }
+     },
+     "e03d78a760b8402cad0827279f4075c1": {
+      "model_module": "@jupyter-widgets/base",
+      "model_module_version": "2.0.0",
+      "model_name": "LayoutModel",
+      "state": {}
+     }
+    },
+    "version_major": 2,
+    "version_minor": 0
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/mlip_arena/tasks/eos/run.py b/mlip_arena/tasks/eos/run.py
index 5f5279b..9efecd9 100644
--- a/mlip_arena/tasks/eos/run.py
+++ b/mlip_arena/tasks/eos/run.py
@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-import asyncio
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -15,9 +14,9 @@
 from ase.optimize import *  # type: ignore
 from ase.optimize.optimize import Optimizer
 from prefect import flow
-from pymatgen.analysis.eos import BirchMurnaghan
-
 from prefect.futures import wait
+from prefect.runtime import flow_run, task_run
+from pymatgen.analysis.eos import BirchMurnaghan
 
 from mlip_arena.models.utils import MLIPEnum
 from mlip_arena.tasks.optimize import run as OPT
@@ -26,13 +25,36 @@
     from ase.filters import Filter
 
 
-@flow
+def generate_flow_run_name():
+    flow_name = flow_run.flow_name
+
+    parameters = flow_run.parameters
+
+    atoms = parameters["atoms"]
+    calculator_name = parameters["calculator_name"]
+
+    return f"{flow_name}: {atoms.get_chemical_formula()} - {calculator_name}"
+
+
+def generate_task_run_name():
+    task_name = task_run.task_name
+
+    parameters = task_run.parameters
+
+    atoms = parameters["atoms"]
+    calculator_name = parameters["calculator_name"]
+
+    return f"{task_name}: {atoms.get_chemical_formula()} - {calculator_name}"
+
+
+# @task(task_run_name=generate_task_run_name)
+@flow(flow_run_name=generate_flow_run_name, validate_parameters=False)
 def fit(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,
     calculator_kwargs: dict | None,
     device: str | None = None,
-    optimizer: Optimizer | str = BFGSLineSearch, # type: ignore
+    optimizer: Optimizer | str = "BFGSLineSearch",  # type: ignore
     optimizer_kwargs: dict | None = None,
     filter: Filter | str | None = None,
     filter_kwargs: dict | None = None,
@@ -59,7 +81,7 @@ def fit(
     Returns:
         A dictionary containing the EOS data and the bulk modulus.
     """
-    result = OPT(
+    first_relax = OPT(
         atoms=atoms,
         calculator_name=calculator_name,
         calculator_kwargs=calculator_kwargs,
@@ -71,7 +93,7 @@ def fit(
         criterion=criterion,
     )
 
-    relaxed = result["atoms"]
+    relaxed = first_relax["atoms"]
 
     # p0 = relaxed.get_positions()
     c0 = relaxed.get_cell()
@@ -99,14 +121,20 @@ def fit(
 
     wait(futures)
 
-    volumes = [f.result()["atoms"].get_volume() for f in futures]
-    energies = [f.result()["atoms"].get_potential_energy() for f in futures]
+    volumes = [f.result()["atoms"].get_volume() for f in futures if isinstance(f.result(), dict)]
+    energies = [f.result()["atoms"].get_potential_energy() for f in futures if isinstance(f.result(), dict)]
+
+    volumes, energies = map(
+        list,
+        zip(
+            *sorted(zip(volumes, energies, strict=True), key=lambda i: i[0]),
+            strict=True,
+        ),
+    )
 
     bm = BirchMurnaghan(volumes=volumes, energies=energies)
     bm.fit()
 
-    volumes, energies = map(list, zip(*sorted(zip(volumes, energies, strict=False), key=lambda i: i[0]), strict=False))
-
     return {
         "eos": {"volumes": volumes, "energies": energies},
         "K": bm.b0_GPa,
diff --git a/mlip_arena/tasks/optimize.py b/mlip_arena/tasks/optimize.py
index 5a21a11..29a61df 100644
--- a/mlip_arena/tasks/optimize.py
+++ b/mlip_arena/tasks/optimize.py
@@ -41,7 +41,11 @@
 } # type: ignore
 
 
-@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(days=1))
+# @task(
+#     cache_key_fn=task_input_hash, 
+#     cache_expiration=timedelta(days=1),
+#     timeout_seconds=120)
+@task(timeout_seconds=120, result_storage=None)
 def run(
     atoms: Atoms,
     calculator_name: str | MLIPEnum,


Test patch:
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 0e0495b..88b2744 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -30,5 +30,8 @@ jobs:
         pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
 
     - name: Run tests
+      env:
+        PREFECT_API_KEY: ${{ secrets.PREFECT_API_KEY }}
+        PREFECT_API_URL: ${{ secrets.PREFECT_API_URL }}
       run: |
         pytest -v tests
\ No newline at end of file
diff --git a/tests/test_eos.py b/tests/test_eos.py
index ffc9e66..58d7a75 100644
--- a/tests/test_eos.py
+++ b/tests/test_eos.py
@@ -12,7 +12,7 @@
 
 @pytest.mark.parametrize("model", [MLIPEnum["MACE-MP(M)"]])
 @pytest.mark.skipif(
-    sys.version_info != (3, 11), reason="requires Python 3.11 to use prefect"
+    tuple(sys.version_info)[:2] != (3, 11), reason="requires Python 3.11 to use prefect"
 )
 def test_eos(model: MLIPEnum):
     """