Skip to content

Commit

Permalink
Promote variables to cluster=True if their dependencies are
Browse files Browse the repository at this point in the history
  • Loading branch information
takluyver committed Jul 10, 2024
1 parent 0e4537b commit c39d836
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 16 deletions.
22 changes: 9 additions & 13 deletions damnit/ctxsupport/ctxrunner.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,24 +127,20 @@ def __init__(self, vars, code):

# Check for cycles
try:
self.ordered_vars()
ordered_names = self.ordered_vars()
except CycleError as e:
# Tweak the error message to make it clearer
raise CycleError(f"These Variables have cyclical dependencies, which is not allowed: {e.args[1]}") from e

# Check for raw-data variables that depend on proc-data variables
for name, var in self.vars.items():
if var.data != RunData.RAW:
continue
proc_dependencies = [dep for dep in self.all_dependencies(var)
if self.vars[dep].data == RunData.PROC]
# 'Promote' variables to match characters of their dependencies
for name in ordered_names:
var = self.vars[name]
deps = [self.vars[dep] for dep in self.all_dependencies(var)]
if var._data is None and any(v.data == RunData.PROC for v in deps):
var._data = RunData.PROC.value

if len(proc_dependencies) > 0:
# If we have a variable that depends on proc data but didn't
# explicitly set `data`, then promote this variable to use proc
# data.
if var._data == None:
var._data = RunData.PROC.value
if any(v.cluster for v in deps):
var.cluster = True

def check(self):
problems = []
Expand Down
14 changes: 11 additions & 3 deletions tests/test_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,11 @@ def bar(run, foo: "var#foo"):
var_promotion_code = """
from damnit.context import Variable
@Variable(title="foo", data="proc")
@Variable()
def baz(run, bar: "var#bar"):
return bar * 2
@Variable(title="foo", data="proc", cluster=True)
def foo(run):
return 42
Expand All @@ -144,9 +148,13 @@ def bar(run, foo: "var#foo"):
# This should not raise an exception
var_promotion_ctx = mkcontext(var_promotion_code)

# `bar` should automatically be promoted to use proc data because it depends
# on a proc variable.
# `bar` & `baz` should be promoted to use proc data & run on a dedicated
# node because they depend on foo.
assert var_promotion_ctx.vars["bar"].data == RunData.PROC
assert var_promotion_ctx.vars["baz"].data == RunData.PROC

assert var_promotion_ctx.vars["bar"].cluster is True
assert var_promotion_ctx.vars["baz"].cluster is True

# Test depending on mymdc fields
good_mymdc_code = """
Expand Down

0 comments on commit c39d836

Please sign in to comment.