Merge pull request #204 from biggraph/darabos-slow

Make it simpler to declare operations "slow"

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

@@ -255,7 +255,7 @@ def op(
             func = matplotlib_to_image(func)
         if slow:
             func = mem.cache(func)
-            func = _global_slow(func)
+            func = make_async(func)
         # Positional arguments are inputs.
         inputs = [
             Input(name=name, type=param.annotation)
@@ -385,9 +385,13 @@ def passive_op_registration(env: str):
     return functools.partial(register_passive_op, env)
 
 
-def slow(func):
+def make_async(func):
     """Decorator for slow, blocking operations. Turns them into separate threads."""
 
+    if asyncio.iscoroutinefunction(func):
+        # If the function is already a coroutine, return it as is.
+        return func
+
     @functools.wraps(func)
     async def wrapper(*args, **kwargs):
         return await asyncio.to_thread(func, *args, **kwargs)
@@ -395,7 +399,6 @@ def slow(func):
     return wrapper
 
 
-_global_slow = slow  # For access inside op().
 CATALOGS_SNAPSHOTS: dict[str, Catalogs] = {}
 
 

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

@@ -8,12 +8,10 @@ from lynxkite.core import workspace
 from .pytorch import pytorch_core
 from lynxkite.core import ops
 from tqdm import tqdm
-import joblib
 import pandas as pd
 import pathlib
 
 
-mem = joblib.Memory(".joblib-cache")
 op = ops.op_registration(core.ENV)
 
 
@@ -57,8 +55,7 @@ class ModelOutputMapping(pytorch_core.ModelMapping):
     pass
 
 
-@op("Train model")
-@ops.slow
+@op("Train model", slow=True)
 def train_model(
     bundle: core.Bundle,
     *,
@@ -82,8 +79,7 @@ def train_model(
     return bundle
 
 
-@op("Model inference")
-@ops.slow
+@op("Model inference", slow=True)
 def model_inference(
     bundle: core.Bundle,
     *,

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

@@ -156,7 +156,7 @@ def wrapped(name: str, func):
         for k, v in kwargs.items():
             if v == "None":
                 kwargs[k] = None
-        res = await ops.slow(func)(*args, **kwargs)
+        res = await ops.run_in_thread(func)(*args, **kwargs)
         # Figure out what the returned value is.
         if isinstance(res, nx.Graph):
             return res