support different queue labels

paraffin/cli.py

@@ -3,7 +3,7 @@
 import typer
 
 from paraffin.submit import submit_node_graph
-from paraffin.utils import get_stage_graph
+from paraffin.utils import get_stage_graph, get_custom_queue
 
 app = typer.Typer()
 
@@ -25,7 +25,8 @@ def main(
     from paraffin.worker import app as celery_app
 
     subgraph = get_stage_graph(names=names, glob=glob)
-    submit_node_graph(subgraph, shutdown_after_finished=shutdown_after_finished)
+    custom_queues = get_custom_queue()
+    submit_node_graph(subgraph, shutdown_after_finished=shutdown_after_finished, custom_queues=custom_queues)
 
     typer.echo(f"Submitted all (n = {len(subgraph)})  tasks.")
     if concurrency > 0:

paraffin/submit.py

@@ -1,13 +1,18 @@
 import networkx as nx
 from celery import chord, group
-
+import typing as t
 from paraffin.worker import repro, shutdown_worker
+import fnmatch
 
 
-def submit_node_graph(subgraph: nx.DiGraph, shutdown_after_finished: bool = False):
+def submit_node_graph(subgraph: nx.DiGraph, shutdown_after_finished: bool = False, custom_queues: t.Optional[dict] = None):
     task_dict = {}
+    custom_queues = custom_queues or {}
     for node in subgraph.nodes:
-        task_dict[node.name] = repro.s(name=node.name)
+        if (matched_pattern := next((pattern for pattern in custom_queues if fnmatch.fnmatch(node.name, pattern)), None)):
+            task_dict[node.name] = repro.s(name=node.name).set(queue=custom_queues[matched_pattern])
+        else:
+            task_dict[node.name] = repro.s(name=node.name)
 
     endpoints = []
     chords = {}

paraffin/utils.py

@@ -2,6 +2,8 @@
 
 import dvc.api
 import networkx as nx
+import pathlib
+import yaml
 
 
 def get_subgraph_with_predecessors(G, X, reverse=False):
@@ -38,3 +40,10 @@ def get_stage_graph(names, glob=False):
     subgraph = nx.subgraph_view(subgraph, filter_node=lambda x: hasattr(x, "name"))
 
     return subgraph
+
+def get_custom_queue():
+    with pathlib.Path("paraffin.yaml").open() as f:
+        config = yaml.safe_load(f)
+
+    return config.get("queue", {})
+# TODO: what about lists, shutdown, ?