Non-blocking training.

examples/Model definition

@@ -1,19 +1,5 @@
 {
   "edges": [
-    {
-      "id": "Linear 2 Activation 1",
-      "source": "Linear 2",
-      "sourceHandle": "output",
-      "target": "Activation 1",
-      "targetHandle": "x"
-    },
-    {
-      "id": "Input: tensor 1 Linear 2",
-      "source": "Input: tensor 1",
-      "sourceHandle": "x",
-      "target": "Linear 2",
-      "targetHandle": "x"
-    },
     {
       "id": "MSE loss 2 Optimizer 2",
       "source": "MSE loss 2",
@@ -43,10 +29,24 @@
       "targetHandle": "input"
     },
     {
-      "id": "Repeat 1 Linear 2",
+      "id": "Input: tensor 1 Linear 1",
+      "source": "Input: tensor 1",
+      "sourceHandle": "x",
+      "target": "Linear 1",
+      "targetHandle": "x"
+    },
+    {
+      "id": "Linear 1 Activation 1",
+      "source": "Linear 1",
+      "sourceHandle": "output",
+      "target": "Activation 1",
+      "targetHandle": "x"
+    },
+    {
+      "id": "Repeat 1 Linear 1",
       "source": "Repeat 1",
       "sourceHandle": "output",
-      "target": "Linear 2",
+      "target": "Linear 1",
       "targetHandle": "x"
     }
   ],
@@ -114,60 +114,6 @@
       "type": "basic",
       "width": 232.0
     },
-    {
-      "data": {
-        "__execution_delay": 0.0,
-        "collapsed": null,
-        "display": null,
-        "error": null,
-        "input_metadata": null,
-        "meta": {
-          "inputs": {
-            "x": {
-              "name": "x",
-              "position": "bottom",
-              "type": {
-                "type": "<class 'inspect._empty'>"
-              }
-            }
-          },
-          "name": "Linear",
-          "outputs": {
-            "output": {
-              "name": "output",
-              "position": "top",
-              "type": {
-                "type": "None"
-              }
-            }
-          },
-          "params": {
-            "output_dim": {
-              "default": "",
-              "name": "output_dim",
-              "type": {
-                "type": "<class 'int'>"
-              }
-            }
-          },
-          "type": "basic"
-        },
-        "params": {
-          "output_dim": "4"
-        },
-        "status": "planned",
-        "title": "Linear"
-      },
-      "dragHandle": ".bg-primary",
-      "height": 200.0,
-      "id": "Linear 2",
-      "position": {
-        "x": 92.32755761444682,
-        "y": 20.626371289630676
-      },
-      "type": "basic",
-      "width": 200.0
-    },
     {
       "data": {
         "__execution_delay": 0.0,
@@ -242,8 +188,8 @@
           "inputs": {},
           "name": "Input: tensor",
           "outputs": {
-            "output": {
-              "name": "output",
+            "x": {
+              "name": "x",
               "position": "top",
               "type": {
                 "type": "tensor"
@@ -292,8 +238,8 @@
           "inputs": {},
           "name": "Input: tensor",
           "outputs": {
-            "output": {
-              "name": "output",
+            "x": {
+              "name": "x",
               "position": "top",
               "type": {
                 "type": "tensor"
@@ -383,6 +329,130 @@
       },
       "type": "basic",
       "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {
+            "input": {
+              "name": "input",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Repeat",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {
+            "same_weights": {
+              "default": false,
+              "name": "same_weights",
+              "type": {
+                "type": "<class 'bool'>"
+              }
+            },
+            "times": {
+              "default": 1.0,
+              "name": "times",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            }
+          },
+          "position": {
+            "x": 487.0,
+            "y": 443.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "same_weights": false,
+          "times": "2"
+        },
+        "status": "planned",
+        "title": "Repeat"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Repeat 1",
+      "position": {
+        "x": -210.0,
+        "y": -135.0
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            }
+          },
+          "name": "Linear",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "top",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "output_dim": {
+              "default": 1024.0,
+              "name": "output_dim",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            }
+          },
+          "position": {
+            "x": 359.0,
+            "y": 310.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "output_dim": "4"
+        },
+        "status": "planned",
+        "title": "Linear"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Linear 1",
+      "position": {
+        "x": 88.83370222907377,
+        "y": 48.642890099180136
+      },
+      "type": "basic",
+      "width": 200.0
     }
   ]
 }

examples/Model use

@@ -575,58 +575,58 @@
               "columns": [
                 "x",
                 "y",
-                "predicted"
+                "pred"
               ],
               "data": [
                 [
-                  "[0.94559073 0.65736622 0.25761551 0.48553199]",
-                  "[1.94559073 1.65736628 1.25761557 1.48553205]",
-                  "[1.5948047637939453, 1.619612693786621, 1.5269112586975098, -0.008817584253847599]"
+                  "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
+                  "[1.19908059 1.175704   1.51475513 1.18939424]",
+                  "[1.560641884803772, 1.5941988229751587, 1.5775359869003296, 1.4935821294784546]"
                 ],
                 [
-                  "[0.47856545 0.46267092 0.6376707  0.84747767]",
-                  "[1.47856545 1.46267092 1.63767076 1.84747767]",
-                  "[1.5928349494934082, 1.6176562309265137, 1.52553391456604, -0.008808750659227371]"
+                  "[0.43681622 0.74680805 0.83598751 0.12414402]",
+                  "[1.43681622 1.74680805 1.83598757 1.12414408]",
+                  "[1.5766589641571045, 1.7117265462875366, 1.7645087242126465, 1.3384637832641602]"
                 ],
                 [
-                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
-                  "[1.59492421 1.90274489 1.38069057 1.46101224]",
-                  "[1.5592631101608276, 1.5841729640960693, 1.5020174980163574, -0.008660631254315376]"
+                  "[0.9829582  0.59269661 0.40120947 0.95487177]",
+                  "[1.9829582  1.59269667 1.40120947 1.95487177]",
+                  "[1.5375217199325562, 1.4159281253814697, 1.2972962856292725, 1.7269455194473267]"
                 ],
                 [
-                  "[0.12858278 0.09930819 0.83222693 0.72485673]",
-                  "[1.12858272 1.09930825 1.83222699 1.72485673]",
-                  "[1.469609260559082, 1.4953691959381104, 1.4395854473114014, -0.00825517252087593]"
+                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
+                  "[1.32565451 1.90939188 1.07488036 1.13730896]",
+                  "[1.562728762626648, 1.6061222553253174, 1.597141146659851, 1.4772177934646606]"
                 ],
                 [
-                  "[0.94516498 0.08422136 0.5608117  0.07652664]",
-                  "[1.94516492 1.08422136 1.56081176 1.07652664]",
-                  "[1.5648787021636963, 1.5899150371551514, 1.5060429573059082, -0.008683123625814915]"
+                  "[0.31518555 0.49643308 0.11509258 0.95458382]",
+                  "[1.31518555 1.49643302 1.11509252 1.95458388]",
+                  "[1.528311848640442, 1.3380011320114136, 1.171952247619629, 1.8305948972702026]"
                 ],
                 [
-                  "[0.54914117 0.03810108 0.87531954 0.73044223]",
-                  "[1.54914117 1.03810108 1.87531948 1.73044229]",
-                  "[1.6262837648391724, 1.650843620300293, 1.548863410949707, -0.00895910244435072]"
+                  "[0.79905868 0.89367443 0.75429088 0.3190186 ]",
+                  "[1.79905868 1.89367437 1.75429082 1.3190186 ]",
+                  "[1.5757312774658203, 1.7105278968811035, 1.7636661529541016, 1.3394038677215576]"
                 ],
                 [
-                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
-                  "[1.94221354 1.57740951 1.98649526 1.40934443]",
-                  "[1.8493703603744507, 1.8721930980682373, 1.704444169998169, -0.009961890056729317]"
+                  "[0.80893755 0.92237449 0.88346356 0.93164903]",
+                  "[1.80893755 1.92237449 1.88346362 1.93164897]",
+                  "[1.562132716178894, 1.6031286716461182, 1.593322992324829, 1.4810831546783447]"
                 ],
                 [
-                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
-                  "[1.8065424  1.08253479 1.74478531 1.71257162]",
-                  "[1.672502040863037, 1.6967051029205322, 1.5810983180999756, -0.009166811592876911]"
+                  "[0.26661873 0.45946234 0.13510543 0.81294441]",
+                  "[1.26661873 1.4594624  1.13510537 1.81294441]",
+                  "[1.533058762550354, 1.3753284215927124, 1.230975866317749, 1.7815138101577759]"
                 ],
                 [
-                  "[0.50272274 0.54912758 0.17663097 0.79070699]",
-                  "[1.50272274 1.54912758 1.17663097 1.79070699]",
-                  "[1.4309396743774414, 1.4570224285125732, 1.4126286506652832, -0.008081027306616306]"
+                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
+                  "[1.39147139 1.29854035 1.84663737 1.58175623]",
+                  "[1.5607244968414307, 1.5942375659942627, 1.5779708623886108, 1.4935153722763062]"
                 ],
                 [
                   "[0.34084332 0.73018837 0.54168713 0.91440833]",
                   "[1.34084332 1.73018837 1.54168713 1.91440833]",
-                  "[1.5581963062286377, 1.5832865238189697, 1.5013742446899414, -0.008653069846332073]"
+                  "[1.5488454103469849, 1.4963982105255127, 1.422922968864441, 1.622254490852356]"
                 ]
               ]
             },
@@ -693,12 +693,12 @@
                   "[1.48959708 1.48549271 1.32688856 1.35667706]"
                 ],
                 [
-                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
-                  "[1.04508126 1.76880038 1.80721331 1.62542391]"
+                  "[0.50272274 0.54912758 0.17663097 0.79070699]",
+                  "[1.50272274 1.54912758 1.17663097 1.79070699]"
                 ],
                 [
-                  "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
-                  "[1.19908059 1.175704   1.51475513 1.18939424]"
+                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
+                  "[1.04508126 1.76880038 1.80721331 1.62542391]"
                 ],
                 [
                   "[0.40167677 0.25953674 0.9407078  0.76308483]",
@@ -736,10 +736,6 @@
                   "[0.68062544 0.98093534 0.14778823 0.53244978]",
                   "[1.68062544 1.98093534 1.14778829 1.53244972]"
                 ],
-                [
-                  "[0.31518555 0.49643308 0.11509258 0.95458382]",
-                  "[1.31518555 1.49643302 1.11509252 1.95458388]"
-                ],
                 [
                   "[0.79121011 0.54161114 0.69369799 0.1520769 ]",
                   "[1.79121017 1.54161119 1.69369793 1.15207696]"
@@ -753,8 +749,8 @@
                   "[1.23942459 1.90487361 1.69337189 1.65089428]"
                 ],
                 [
-                  "[0.26661873 0.45946234 0.13510543 0.81294441]",
-                  "[1.26661873 1.4594624  1.13510537 1.81294441]"
+                  "[0.94516498 0.08422136 0.5608117  0.07652664]",
+                  "[1.94516492 1.08422136 1.56081176 1.07652664]"
                 ],
                 [
                   "[0.30754459 0.77694583 0.09278506 0.38326019]",
@@ -780,6 +776,10 @@
                   "[0.78956431 0.87284744 0.06880784 0.03455889]",
                   "[1.78956437 1.87284744 1.06880784 1.03455889]"
                 ],
+                [
+                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
+                  "[1.94221354 1.57740951 1.98649526 1.40934443]"
+                ],
                 [
                   "[0.00497234 0.39319336 0.57054168 0.75150961]",
                   "[1.00497234 1.39319336 1.57054162 1.75150967]"
@@ -845,16 +845,8 @@
                   "[1.95928192 1.84273899 1.7151463  1.38619852]"
                 ],
                 [
-                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
-                  "[1.32565451 1.90939188 1.07488036 1.13730896]"
-                ],
-                [
-                  "[0.9829582  0.59269661 0.40120947 0.95487177]",
-                  "[1.9829582  1.59269667 1.40120947 1.95487177]"
-                ],
-                [
-                  "[0.79905868 0.89367443 0.75429088 0.3190186 ]",
-                  "[1.79905868 1.89367437 1.75429082 1.3190186 ]"
+                  "[0.54914117 0.03810108 0.87531954 0.73044223]",
+                  "[1.54914117 1.03810108 1.87531948 1.73044229]"
                 ],
                 [
                   "[0.67418337 0.79634351 0.23229051 0.71345252]",
@@ -869,12 +861,12 @@
                   "[1.81788456 1.58174157 1.29376316 1.79712534]"
                 ],
                 [
-                  "[0.60075855 0.12234765 0.00614399 0.30560958]",
-                  "[1.60075855 1.12234759 1.00614405 1.30560958]"
+                  "[0.94559073 0.65736622 0.25761551 0.48553199]",
+                  "[1.94559073 1.65736628 1.25761557 1.48553205]"
                 ],
                 [
-                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
-                  "[1.39147139 1.29854035 1.84663737 1.58175623]"
+                  "[0.60075855 0.12234765 0.00614399 0.30560958]",
+                  "[1.60075855 1.12234759 1.00614405 1.30560958]"
                 ],
                 [
                   "[0.02162331 0.81861657 0.92468154 0.07808572]",
@@ -904,6 +896,10 @@
                   "[0.60609657 0.96257663 0.19292736 0.95702219]",
                   "[1.60609651 1.96257663 1.19292736 1.95702219]"
                 ],
+                [
+                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
+                  "[1.8065424  1.08253479 1.74478531 1.71257162]"
+                ],
                 [
                   "[0.70167565 0.26930219 0.5660674  0.61194974]",
                   "[1.70167565 1.26930213 1.56606746 1.61194968]"
@@ -912,6 +908,10 @@
                   "[0.76933283 0.86241865 0.44114518 0.65644735]",
                   "[1.76933289 1.86241865 1.44114518 1.65644741]"
                 ],
+                [
+                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
+                  "[1.59492421 1.90274489 1.38069057 1.46101224]"
+                ],
                 [
                   "[0.15064228 0.03198934 0.25754827 0.51484001]",
                   "[1.15064228 1.03198934 1.25754833 1.51484001]"
@@ -949,13 +949,17 @@
                   "[1.40234613 1.54987347 1.49542785 1.5415318 ]"
                 ],
                 [
-                  "[0.80893755 0.92237449 0.88346356 0.93164903]",
-                  "[1.80893755 1.92237449 1.88346362 1.93164897]"
+                  "[0.12858278 0.09930819 0.83222693 0.72485673]",
+                  "[1.12858272 1.09930825 1.83222699 1.72485673]"
                 ],
                 [
                   "[0.72470158 0.4940322  0.41027349 0.89364016]",
                   "[1.72470164 1.49403214 1.41027355 1.89364016]"
                 ],
+                [
+                  "[0.47856545 0.46267092 0.6376707  0.84747767]",
+                  "[1.47856545 1.46267092 1.63767076 1.84747767]"
+                ],
                 [
                   "[0.49584109 0.80599248 0.07096875 0.75872749]",
                   "[1.49584103 1.80599248 1.07096875 1.75872755]"
@@ -976,10 +980,6 @@
                   "[0.68094063 0.45189077 0.22661722 0.37354094]",
                   "[1.68094063 1.45189071 1.22661722 1.37354088]"
                 ],
-                [
-                  "[0.43681622 0.74680805 0.83598751 0.12414402]",
-                  "[1.43681622 1.74680805 1.83598757 1.12414408]"
-                ],
                 [
                   "[0.47870928 0.17129105 0.27300501 0.20634609]",
                   "[1.47870922 1.17129111 1.27300501 1.20634604]"
@@ -1000,7 +1000,7 @@
             }
           },
           "other": {
-            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (2) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (4) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (5) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (6) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (7) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (8) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (9) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (10) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__tensor_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['Input__tensor_3_x', 'END_Repeat_1_output'], loss=Sequential(\n  (0) - <function mse_loss at 0x75e660939d00>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> loss\n), optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace=None, trained=True)"
+            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (2) - <function leaky_relu at 0x762d1f82c680>: Linear_1_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (4) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (5) - <function leaky_relu at 0x762d1f82c680>: Linear_1_output -> Activation_1_output\n  (6) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (7) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__tensor_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['END_Repeat_1_output', 'Input__tensor_3_x'], loss=Sequential(\n  (0) - <function mse_loss at 0x762d1f82e160>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> loss\n), optimizer_parameters={'lr': 0.1, 'type': <OptionsFor_type.SGD: 4>}, optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace='Model definition', trained=True)"
           },
           "relations": []
         },
@@ -1016,7 +1016,7 @@
               },
               "df_test": {
                 "columns": [
-                  "predicted",
+                  "pred",
                   "x",
                   "y"
                 ]
@@ -1035,8 +1035,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__tensor_3_x",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1210,8 +1210,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__tensor_3_x",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1270,7 +1270,7 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "110",
+          "epochs": "1003",
           "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_x\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
           "model_name": "model"
         },
@@ -1322,8 +1322,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__tensor_3_x",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1384,13 +1384,13 @@
         "params": {
           "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
           "model_name": "model",
-          "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"predicted\"}}}"
+          "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"pred\"}}}"
         },
         "status": "done",
         "title": "Model inference"
       },
       "dragHandle": ".bg-primary",
-      "height": 893.0,
+      "height": 650.0,
       "id": "Model inference 1",
       "position": {
         "x": 2181.718373860645,

lynxkite-app/web/src/workspace/nodes/NodeParameter.tsx

@@ -91,6 +91,7 @@ function ModelMapping({ value, onChange, data, variant }: any) {
   const dfs: { [df: string]: string[] } = {};
   const inputs = data?.input_metadata?.value ?? data?.input_metadata ?? [];
   for (const input of inputs) {
+    if (!input.dataframes) continue;
     const dataframes = input.dataframes as {
       [df: string]: { columns: string[] };
     };

lynxkite-core/src/lynxkite/core/ops.py

@@ -1,6 +1,7 @@
 """API for implementing LynxKite operations."""
 
 from __future__ import annotations
+import asyncio
 import enum
 import functools
 import inspect
@@ -297,3 +298,13 @@ def op_registration(env: str):
 def passive_op_registration(env: str):
     """Returns a function that can be used to register operations without associated code."""
     return functools.partial(register_passive_op, env)
+
+
+def slow(func):
+    """Decorator for slow, blocking operations. Turns them into separate threads."""
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        return await asyncio.to_thread(func, *args, **kwargs)
+
+    return wrapper

lynxkite-graph-analytics/src/lynxkite_graph_analytics/core.py

@@ -1,5 +1,6 @@
 """Graph analytics executor and data types."""
 
+import inspect
 import os
 from lynxkite.core import ops, workspace
 import dataclasses
@@ -177,10 +178,16 @@ async def execute(ws: workspace.Workspace):
                 # All inputs for this node are ready, we can compute the output.
                 todo.remove(id)
                 progress = True
-                _execute_node(node, ws, catalog, outputs)
+                await _execute_node(node, ws, catalog, outputs)
 
 
-def _execute_node(node, ws, catalog, outputs):
+async def await_if_needed(obj):
+    if inspect.isawaitable(obj):
+        obj = await obj
+    return obj
+
+
+async def _execute_node(node, ws, catalog, outputs):
     params = {**node.data.params}
     op = catalog.get(node.data.title)
     if not op:
@@ -214,6 +221,7 @@ def _execute_node(node, ws, catalog, outputs):
     # Execute op.
     try:
         result = op(*inputs, **params)
+        result.output = await await_if_needed(result.output)
     except Exception as e:
         if os.environ.get("LYNXKITE_LOG_OP_ERRORS"):
             traceback.print_exc()

lynxkite-graph-analytics/src/lynxkite_graph_analytics/lynxkite_ops.py

@@ -369,6 +369,7 @@ class ModelOutputMapping(pytorch_model_ops.ModelMapping):
 
 
 @op("Train model")
+@ops.slow
 def train_model(
     bundle: core.Bundle,
     *,
@@ -380,9 +381,11 @@ def train_model(
     m = bundle.other[model_name].copy()
     inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
     t = tqdm(range(epochs), desc="Training model")
+    losses = []
     for _ in t:
         loss = m.train(inputs)
         t.set_postfix({"loss": loss})
+        losses.append(loss)
     m.trained = True
     bundle = bundle.copy()
     bundle.other[model_name] = m
@@ -390,6 +393,7 @@ def train_model(
 
 
 @op("Model inference")
+@ops.slow
 def model_inference(
     bundle: core.Bundle,
     *,