[batch] Fix job_group cancellation logic in SQL procedures

[ehigham]

The following sql triggers and procedures were not updated to reflect
changes to job cancellation when job groups merged:

• jobs_after_update
• cancel_batch
• schedule_job
• mark_job_creating
• mark_job_started

These error if there's more than 1 cancelled job group in the batch,
preventing jobs being marked complete even if they had been executed
sucessfully.

This change drops and re-creates these entities with updated job
cancellation logic.

Fixes #14864

Security Assessment

This change impacts the Hail Batch instance as deployed by Broad Institute in GCP.
This change has a medium security impact, affecting

• the stored SQL procedures in the batch database, and
• error handling in the batch worker.

This change does not expose new functionality that could be exploited.

Appsec Review

• Required: The impact has been assessed and approved by appsec

[ehigham]

Compare with
https://github.com/hail-is/hail/blob/62efb32d9544b0c964019345281e89b235f02f92/batch/sql/finalize-job-groups.sql
https://github.com/hail-is/hail/blob/62efb32d9544b0c964019345281e89b235f02f92/batch/sql/rename-job-groups-tables.sql

[grohli]

Is this the original cancel_job_group? (unless I missed it in those two places you linked, which is very likely)

[ehigham]

I believe so yes

[ehigham]

Tested by submitting stress tests and monitoring batch-driver logs
for errors to batch services running in

• the default k8s namespace to reproduce the errors
• a dev deployment patched with this branch.

The stress test is below and requires #14817.

The default-namespace deployment logged errors matching that of #14864
whereas the patched instance did not. I claim that this change fixes that issue.

from random import choice, random
from typing import Callable, Optional, TypeVar, Union

import hailtop.batch as hb


def flip(p):
    return random() <= p


A = TypeVar('A')


def prob(p: float, f: Callable[[], A]) -> Optional[A]:
    return f() if flip(p) else None


def job_group_name(a: Union[hb.Batch, hb.JobGroup]) -> str:
    return a.attributes['name'] if isinstance(a, hb.JobGroup) else ''


def job_group_tree(root: Union[hb.Batch, hb.JobGroup], depth: int, arity: int) -> hb.JobGroup:
    if depth <= 1:
        return

    failure_threashold = prob(0.5, lambda: 1)
    group_name = f'{job_group_name(root)}/{depth}'
    jg = root.create_job_group(attributes={'name': group_name}, cancel_after_n_failures=failure_threashold)
    j = jg.new_bash_job(name=f'({group_name}, 0)')
    j.command(f'sleep {choice(range(4))}; exit {1 if flip(0.1) else 0}')

    for m in range(arity):
        jm = jg.new_bash_job(name=f'({group_name}, {m + 1})')
        jm = jm.depends_on(j)
        jm.command(f'sleep {choice(range(4))}; exit {1 if flip(0.1) else 0}')
        jm._always_run = flip(0.1)

        if flip(0.01):
            jm._machine_type = 'n1-standard-1'
            jm.spot(flip(0.5))

    for n in range(arity):
        job_group_tree(jg, depth - 1, arity)


def stress():
    b = hb.Batch(name='stress', backend=hb.ServiceBackend(billing_project='ehigham-trial'))
    job_group_tree(b, depth=3, arity=100)
    b.run(open=False, wait=False)


if __name__ == "__main__":
    stress()

[grohli]

Looks great, thanks for the improvement Ed

[sarahgibs]

approved by security.