simple gpu launcher

omx_launch.py

import subprocess
subprocess.run("getgist miraculixx progress.py".split(' '))

import omegaml as om_
from progress import progress

om = om_

def match_active_queues(queue=None, nodes=None):
    queues = om.runtime.celeryapp.control.inspect().active_queues()
    for node, node_qs in queues.items():
        is_system_q = lambda user_q: user_q['name'].startswith('amq.')
        for q in (user_q['name'] for user_q in node_qs if not is_system_q(user_q)):
            matched = (q == queue and (not nodes or node in nodes))
            yield matched, node, q
            
@progress            
def wait_for(queue, nodes=None, pbar=None):
    matched, a_nodes, a_queues = list(zip(*match_active_queues(queue, nodes)))
    pbar.set_description_str(desc=f'waiting for {queue}')
    pbar.set_postfix_str('active={}'.format(sorted(a_queues)))                
    return any(matched), a_queues
    
def launch(service, specs, queue, timeout=-1, check=None, using=None):
    from omegaml.client import cli
    global om
    om = using if using else om
    check = check or (lambda t : wait_for(queue, timeout=t)[0])
    if not check(0):
        print(f"{queue} not found, requesting {service} with {specs}")
        cli.main(argv=f'cloud update {service} --specs="{specs}"'.split(' '))
    msg = "ok, {queue} is available" if check(timeout) else f'{queue} did not appear so far'
    return msg