Support Parsl special options: `parsl_resource_specification`, `stdout`, `stderr`, and `walltime`

[Andrew-S-Rosen]

What new feature would you like to see?

As shown here, we need to be able to support Parsl "special options" with @job decorators: namely, parsl_resource_specification, stdout, stderr, and walltime.

Originally mentioned in https://github.com/Quantum-Accelerators/quacc/discussions/1881.

[Andrew-S-Rosen]

Comment from Parsl team:

YB: Here's how I see it from the HTEX perspective. An app behaves mostly like a python function, but with some additional parsl behavior (like walltime limits, parsl_resource_spec... ). If the issue is that you do not want to/can't modify a user supplied function, you can always wrap your function with a wrapper that takes the parsl_resource_specification and internally calls the user supplied function.

AR: Yeah, the issue is that it would not be practical or feasible to modify the underlying functions. Interesting idea about the wrapper. That seems like a potentially viable path forward, but my concern would be that this would then mess up things with monitoring (e.g. the function name would be the wrapped function name and such). I guess that might be solvable with a @functools.wraps kind of thing?

[tomdemeyere]

Following yesterday's deleted comment, if you do the following change in the job decorator you get what you want. There is no additional assumptions about the parameter that _func takes. The only problem is what you already reported above: what is being launched is function job.<locals>.wrapper which can be problematic for some situations (monitoring...), doing wrapper.__name__ = _func.__name__ does not not seem to have any impact.

elif SETTINGS.WORKFLOW_ENGINE == "parsl":
    from parsl import python_app
    stdout = kwargs.pop("stdout", None)
    stderr = kwargs.pop("stderr", None)
    walltime = kwargs.pop("walltime", None)
    parsl_resource_specification = kwargs.pop("parsl_resource_specification", None)
    def wrapper(
        *f_args,
        stdout=stdout,
        stderr=stderr,
        walltime=walltime,
        parsl_resource_specification=parsl_resource_specification,
        **f_kwargs,
    ):
        return _func(*f_args, **f_kwargs)
    return python_app(wrapper, **kwargs)

config = Config(
    executors=[
        HighThroughputExecutor(
            label="htex_local",
            cores_per_worker=1,
            max_workers=1,
            provider=LocalProvider(
                channel=LocalChannel(),
                init_blocks=1,
                max_blocks=1,
            ),
        )
    ],
)

resource_spec = {
    "num_nodes": 2,
    "ranks_per_node": 2,
    "num_ranks": 4,
}

parsl.load(config)

@job(parsl_resource_specification=resource_spec)
def test_mpi_apps():
    from os import environ

    return environ.get("PARSL_RANKS_PER_NODE")

future = test_mpi_apps()

print(future.result())

This should print 2

[Andrew-S-Rosen]

Yup, that makes perfect sense to me. It's possible that @functools.wraps around the wrapper could fix the function metadata (should probably be used either way).

Feel free to open a PR if you would like. Otherwise, I appreciate your comment here and will get to it one day. :)