`simulation_app` does not work as documented

[theo-brown]

Hi! I might be being dumb, but I couldn't get the simulation app to work from within a script.

My example script:

from absl import app
from torax.examples import basic_config
from torax import simulation_app
from torax.config.build_sim import build_sim_from_config

sim_builder = lambda: build_sim_from_config(basic_config.CONFIG)
app.run(simulation_app.main(sim_builder))

raises the error

Traceback (most recent call last):
  File ".../torax-dev/torax/tests/test_data/test_prescribed_current_source.py", line 80, in <module>
    app.run(simulation_app.main(sim_builder))
  File ".../python3.12/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File ".../python3.12/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
TypeError: 'tuple' object is not callable

Note that there is a typo in the simulation_app.py docstring; where it says app.run(simulation_app.main(get_sim())), I believe it should say app.run(simulation_app.main(get_sim)), based on the signature of simulation_app.main. In my example above I have corrected for this.

[jcitrin]

Thanks for noting the simulation_app typo, indeed a callable needs to be passed. I will remove the () in a quick PR.

[jcitrin]

About running the script, the issue seems to be on exit, and I think it's an absl issue related to the scripting pattern. You can see here, when doing the same thing but with logging enabled, that it crashes on exit

from absl import app
from absl import logging
from torax.examples import basic_config
from torax import simulation_app
from torax.config import build_sim

logging.set_verbosity(logging.INFO)
sim_builder = lambda: build_sim.build_sim_from_config(basic_config.CONFIG)
app.run(simulation_app.main(
        sim_builder,
        log_sim_progress=True,
        ))

[jcitrin]

I recommend just using the if __name__ == '__main__': pattern

The following should work, and you can see two different patterns to run TORAX. One using the Sim.run() method directly, and the other with simulation_app.main

from absl import app
from absl import logging
from torax.examples import basic_config
from torax import simulation_app
from torax.config import build_sim

def run(_):
  sim = build_sim.build_sim_from_config(basic_config.CONFIG)
  torax_outputs = sim.run(log_timestep_info=True)

def run2(_):
  sim_builder = lambda: build_sim.build_sim_from_config(basic_config.CONFIG)
  simulation_app.main(
    sim_builder,
    log_sim_progress=True,
    )

if __name__ == '__main__':
  logging.set_verbosity(logging.INFO)
  app.run(run)
  #app.run(run2)

[jcitrin]

Looking through it deeper, it seems that app.run calls sys.exit on the output of the function it runs.

sys.exit can take as input an exitcode of type str, int, or None. The idea here is to provide functionality for error handling of code outputs.

However, simulation_app.main outputs a tuple [xr.Dataset, str], so sys.exit crashes. In the examples I showed above, run and run2 have None outputs, so they run fine.

I will fix the docstring to make this all clearer. Thanks for raising this.

@araju FYI

[theo-brown]

Yes, I think that's the source of the error I was experiencing! Sorry for not debugging it myself, thanks for having a look at it.

[jcitrin]

Fixed with #237