Passing SBATCH Args via FirecREST

[AryanAhadinia]

We have a python script, that uses Jinja to render the a shell script and submit it as a job either using SLURM or FirecREST. So, technically, we want to get rid of Jinja since it deteriorated the readability of script. We are using Jinja template for providing some SBATCH args and some env vars as below:

#!/bin/bash
#SBATCH --time={{ time }}
...

FRAMEWORK="{{ framework }}"
...

So, we are completely fine with the env vars. Fortunately, pyfirecrest provides an API to pass the env vars as await self.client.submit(job_script, ..., env_vars=...) and also sbatch/srun commands provide an CLI argument to do so as srun --export=VAR1=value1 job.sh. What is problematic for us, is the sbatch args at the header of the script, e.g., #SBATCH --time={{ time }}. We can move these args from the header of the script to the command as srun --time=02:00:00 ..., but the only remaining piece here is to pass these attributes with firecrest. I wonder if there is a way to do so. It would be more than great for us!

[ekouts]

Hi @AryanAhadinia , thanks for opening the issue! Are you using FirecREST v2? The API of the submit method is slightly different. You can use script_str, script_local_path or script_remote_path depending if your script is locally or on the filesystem of the remote machine. In your case the first two would be the best. So you are free to dynamically create locally (in your python code) the job script and pass it directly to the client.
I think you may need some of the args as SBATCH directives (for example if you want the slurm allocation to have the time limit and not just the job step).

[AryanAhadinia]

@ekouts Thanks for your reply.

Actually, our code is public and you can take a look to that. Our script is a Jinja template which is rendered and submitted using this python script.

Are you using FirecREST v2?

I just realized that we were mistakenly using V1, but I think it does not affect the issue in hand since neither of them have support for such functionality.

You can use script_str, script_local_path or script_remote_path depending if your script is locally or on the filesystem of the remote machine. In your case the first two would be the best. So you are free to dynamically create locally (in your python code) the job script and pass it directly to the client.

Exactly! this is what we are doing now, in a glance, we (1) read the script Jinja template from the assets, (2) render it using Jinja, and (3) submit it as client.submit(script_str=rendered_script, ...). The SBATCH directives are also injected during the rendering process. But indeed, we were thinking it would be nicer if we can pass the SBATCH directives as an argument to the submit method.

So, with simplifications, what we have now, is like below.

template = Template(_TEMPLATE_PATH.read_text())
rendered_script = template.render(
    time=time,                 # SBATCH directive
    partition=partition,       # SBATCH directive
    ...,                       # and so on ...
    framework=framework,       # env var in the script
    router_port=router_port,   # env var in the script
    ...,                       # and so on ...
)
client.submit(
    script_str=rendered_script,
)

But, what we are interested to have is something like below.

script_str = open("script.sh").read()
client.submit(
    script_str=script_str,
    time=time,                     # SBATCH directive
    partition=partition,           # SBATCH directive
    ...,                           # and so on ...
    env_vars={
        framework=framework,       # env var in the script
        router_port=router_port,   # env var in the script
        ...,                       # and so on ...
    },
)

In this way, the code would be much cleaner. As of for now, we don't miss a functionality that blocks our work, but adding this API will make the code much cleaner and we can get rid of Jinja (or any other formatting/rendering tools). Also, this will make the script much readable since reading a template script with Jinja place holders is a bit hard.

There is already support for the env_vars but there is no support for the directives. I think this would be a very nice addition and will help a lot of people for cleaner and simpler codes.

This is also aligned with srun commands. You can have submit a job with SBATCH directives at the top of the shell script, or adding the directives to the command as srun time=1:00:00 partition=normal ....