allenai / tango

Organize your experiments into discrete steps that can be cached and reused throughout the lifetime of your research project.
https://ai2-tango.readthedocs.io/
Apache License 2.0
533 stars 47 forks source link

CLI-style execution for Python-defined experiments #600

Closed apoorvkh closed 1 year ago

apoorvkh commented 1 year ago

Hey!

I think the Tango CLI is really nice and would really like if we could support Python-defined experiments with CLI-style execution. So I modularized the CLI run functions and decoupled them from __main__.py.

Now we can also run scripts like:

from tango import StepGraph
from tango.cli import tango_cli, execute_step_graph, prepare_workspace, prepare_executor

from examples.euler.complex_arithmetic import *
from tango.steps import PrintStep

i = [0.0, 1.0]
pi = [3.1415926535, 0.0]

i_times_pi = MultiplyStep(a=i, b=pi)
pow_e = ExponentiateStep(x=i_times_pi)
plus_one = AdditionStep(a=pow_e, b=[1,0])
print_step = PrintStep(input=plus_one)

step_graph = StepGraph({
    "i_times_pi": i_times_pi,
    "pow_e": pow_e,
    "plus_one": plus_one,
    "print": print_step
})

if __name__ == '__main__':
    with tango_cli():
        workspace = prepare_workspace()
        executor = prepare_executor(workspace)
        execute_step_graph(step_graph, workspace, executor)

This should be the same as tango run examples/euler/euler.jsonnet!

If we want to use the default TangoGlobalSettings, we can also shorten the block to

    with tango_cli():
        execute_step_graph(step_graph)

This includes the changes from (@BigRedT) Tanmay's PR #491. CLI functionality should be completely the same. I ran the tests in tests/main_test.py successfully. Formatted with isort and black and checked with ruff and mypy. Docs needed.

What do you think? Happy to make adjustments as you see fit!

apoorvkh commented 1 year ago

Seems like the lint and type check failed, but for reasons independent of this PR. Shall I resolve those anyway?

apoorvkh commented 1 year ago

That's great -- thanks a lot @AkshitaB!