build_wf.py

#!/usr/bin/env python
"""
Generate many fireworks workflows.
"""
import argparse
import os
import sys

def main():
    p = argparse.ArgumentParser(description="Generate fireworks")
    p.add_argument("--output", help="Output directory (.)", dest="out_dir", default=".")
    p.add_argument("--tasks", help="Number of fireworks in a workflow (10)",
                   dest="tasks", type=int, default=10)
    p.add_argument("--wf", help="Number of workflows (1)", dest="wf", type=int, default=1)
    p.add_argument("--type", help="Type of workflow (complex)", dest="wf_type",
                   choices=["sequence", "complex"], default="complex")
    args = p.parse_args()
    command = eval("build_" + args.wf_type)
    path = args.out_dir
    if not os.path.exists(path):
        os.mkdir(path)
    multi_wf = args.wf > 1
    for i in xrange(args.wf):
        spec = command(tasks=args.tasks)
        wf_num_sfx = "_{:d}".format(i) if multi_wf else ""
        outfile = os.path.join(path, "fw_{}_{:d}{}.yaml"
                                     .format(args.wf_type, args.tasks, wf_num_sfx))
        with open(outfile, "w") as f:
            f.write(spec)
    return 0

def build_complex(tasks=1):
    n = max(tasks / 2, 1)
    fws = ["""
    - fw_id: {{:d}}
      spec:
        _tasks:
        - _fw_name: ScriptTask
          script: echo '{}'
    """.format(pith) for pith in 'To be, or not to be,', 'that is the question:']
    fwlist = [fws[0].format(i + 1) for i in xrange(n)] + \
             [fws[1].format(i + n + 1) for i in xrange(n + 1)]
    linklist = ["""
      {:d}:
      - {:d}
      - {:d}
    """.format(i + 1, i + n + 1, i + n + 2) for i in xrange(n)]

    return build_spec(fwlist, linklist)

def build_sequence(tasks=1):
    n = max(tasks / 2, 1)
    fws = ["""
    - fw_id: {{:d}}
      spec:
        _tasks:
        - _fw_name: ScriptTask
          script: echo '{}'
    """.format(pith) for pith in 'To be, or not to be,', 'that is the question:']
    fwlist = [fws[0].format(i) for i in xrange(1, n + 1)] + \
             [fws[1].format(i + n) for i in xrange(1, n + 1)]
    linklist = ["""
      {:d}:
      - {:d}\n""".format(i + n, i) for i in xrange(1, n + 1)]

    return build_spec(fwlist, linklist)

def build_spec(fwlist, linklist):
    spec ="""fws:
    {fireworks}
links:
    {links}

metadata: {{}}
    """.format(fireworks=''.join(fwlist), links=''.join(linklist))
    return spec


if __name__ == "__main__":
    sys.exit(main())