queue_main.py

importdatetime
importos
fromsubprocessimportcall
fromtimeimportsleep
fromtimeimporttime

importheroku3

fromappimportHEROKU_APP_NAME
fromappimportdb
fromappimportlogger
fromutilimportelapsed
fromutilimportget_sql_answer
fromutilimportget_sql_answers
fromutilimportrun_sql
fromutilimportsafe_commit


classDbQueue(object):

def__init__(self, **kwargs):
self.parsed_vars= {}
super(DbQueue, self).__init__(**kwargs)

defmonitor_till_done(self, job_type):
logger.info("collecting data. will have some stats soon...")
logger.info("\n\n")

num_total=self.number_total_on_queue(job_type)
print("num_total", num_total)
num_unfinished=self.number_unfinished(job_type)
print("num_unfinished", num_unfinished)

loop_thresholds= {"short": 30, "long": 10*60, "medium": 60}
loop_unfinished= {"short": num_unfinished, "long": num_unfinished}
loop_start_time= {"short": time(), "long": time()}

# print_idle_dynos(job_type)

whileall(loop_unfinished.values()):
forloopin ["short", "long"]:
ifelapsed(loop_start_time[loop]) >loop_thresholds[loop]:
ifloopin ["short", "long"]:
num_unfinished_now=self.number_unfinished(job_type)
num_finished_this_loop=loop_unfinished[loop] -num_unfinished_now
loop_unfinished[loop] =num_unfinished_now
ifloop=="long":
logger.info("\n****"),
logger.info(" {} finished in the last {} seconds, {} of {} are now finished ({}%). ".format(
num_finished_this_loop, loop_thresholds[loop],
num_total-num_unfinished_now,
num_total,
int(100*float(num_total-num_unfinished_now)/num_total)
 )), # comma so the next part will stay on the same line
ifnum_finished_this_loop:
minutes_left=float(num_unfinished_now) /num_finished_this_loop*loop_thresholds[loop] /60
logger.info("{} estimate: done in {} mins, which is {} hours".format(
loop, round(minutes_left, 1), round(minutes_left/60, 1)))
else:
print()
loop_start_time[loop] =time()
# print_idle_dynos(job_type)
print(".")
sleep(3)
logger.info("everything is done. turning off all the dynos")
self.scale_dyno(0, job_type)


defnumber_total_on_queue(self, job_type):
num=get_sql_answer(db, "select count(*) from {}".format(self.table_name(job_type)))
returnnum

defnumber_waiting_on_queue(self, job_type):
num=get_sql_answer(db, "select count(*) from {} where started is null".format(self.table_name(job_type)))
returnnum

defnumber_unfinished(self, job_type):
num=get_sql_answer(db, "select count(*) from {} where finished is null".format(self.table_name(job_type)))
returnnum

defprint_status(self, job_type):
num_dois=self.number_total_on_queue(job_type)
num_waiting=self.number_waiting_on_queue(job_type)
ifnum_dois:
logger.info("There are {} dois in the queue, of which {} ({}%) are waiting to run".format(
num_dois, num_waiting, int(100*float(num_waiting)/num_dois)))

defkick(self, job_type):
q="""update {table_name} set started=null, finished=null
 where finished is null""".format(
table_name=self.table_name(job_type))
run_sql(db, q)
self.print_status(job_type)

defreset_enqueued(self, job_type):
q="update {} set started=null, finished=null".format(self.table_name(job_type))
run_sql(db, q)

defnum_dynos(self, job_type):
heroku_conn=heroku3.from_key(os.getenv("HEROKU_API_KEY"))
num_dynos=0
try:
dynos=heroku_conn.apps()[HEROKU_APP_NAME].dynos()[self.process_name(job_type)]
num_dynos=len(dynos)
except (KeyError, TypeError) ase:
pass
returnnum_dynos

defprint_idle_dynos(self, job_type):
heroku_conn=heroku3.from_key(os.getenv("HEROKU_API_KEY"))
app=heroku_conn.apps()[HEROKU_APP_NAME]
running_dynos= []
try:
running_dynos= [dynofordynoinapp.dynos() ifdyno.name.startswith(self.process_name(job_type))]
except (KeyError, TypeError) ase:
pass

dynos_still_working=get_sql_answers(db, "select dyno from {} where started is not null and finished is null".format(self.table_name(job_type)))
dynos_still_working_names= [nfornindynos_still_working]

logger.info("dynos still running: {}".format([d.namefordinrunning_dynosifd.nameindynos_still_working_names]))
# logger.info(u"dynos stopped:", [d.name for d in running_dynos if d.name not in dynos_still_working_names])
# kill_list = [d.kill() for d in running_dynos if d.name not in dynos_still_working_names]

defscale_dyno(self, n, job_type):
logger.info("starting with {} dynos".format(self.num_dynos(job_type)))
logger.info("setting to {} dynos".format(n))
heroku_conn=heroku3.from_key(os.getenv("HEROKU_API_KEY"))
app=heroku_conn.apps()[HEROKU_APP_NAME]
app.process_formation()[self.process_name(job_type)].scale(n)

logger.info("sleeping for 2 seconds while it kicks in")
sleep(2)
logger.info("verifying: now at {} dynos".format(self.num_dynos(job_type)))



defprint_logs(self, job_type):
command="heroku logs -t --dyno={}".format(self.process_name(job_type))
call(command, shell=True)



defupdate_fn(self, cls, method_name, objects, index=1, kwargs_map=None):

# we are in a fork! dispose of our engine.
# will get a new one automatically
# if is pooling, need to do .dispose() instead
db.engine.dispose()

start=time()
num_obj_rows=len(objects)

# logger.info(u"{pid} {repr}.{method_name}() got {num_obj_rows} objects in {elapsed} seconds".format(
# pid=os.getpid(),
# repr=cls.__name__,
# method_name=method_name,
# num_obj_rows=num_obj_rows,
# elapsed=elapsed(start)
# ))

forcount, objinenumerate(objects):
start_time=time()

ifobjisNone:
returnNone

method_to_run=getattr(obj, method_name)

# logger.info(u"***")
logger.info("*** #{count} starting {repr}.{method_name}() method".format(
count=count+ (num_obj_rows*index),
repr=obj,
method_name=method_name
 ))

method_kwargs=kwargs_map.get(obj.id, {}) ifkwargs_mapelse {}
method_to_run(**method_kwargs)

logger.info("finished {repr}.{method_name}(). took {elapsed} seconds".format(
repr=obj,
method_name=method_name,
elapsed=elapsed(start_time, 4)
 ))

# for handling the queue
ifnot (method_name=="update"andobj.__class__.__name__=="Pub"):
obj.finished=datetime.datetime.utcnow().isoformat()
# db.session.merge(obj)


start_time=time()
commit_success=safe_commit(db)
ifnotcommit_success:
logger.info("COMMIT fail")
logger.info("commit took {} seconds".format(elapsed(start_time, 2)))
db.session.remove() # close connection nicely
returnNone# important for if we use this on RQ

defrun(self, parsed_args, job_type):
start=time()

try:
self.worker_run(**vars(parsed_args))
exceptExceptionase:
logger.fatal('worker_run died with error:', exc_info=True)


logger.info("finished update in {} seconds".format(elapsed(start)))
# resp = None
# if job_type in ["normal"]:
# my_location = Page.query.get(parsed_args.id)
# resp = my_location.__dict__
# pprint(resp)

print("done")
return


defprint_update(self, new_loop_start_time, chunk, limit, start_time, index):
num_items=limit#let's say have to do the full limit
num_jobs_remaining=num_items- (index*chunk)
try:

jobs_per_hour_this_chunk=chunk/float(elapsed(new_loop_start_time) /3600)
predicted_mins_to_finish=round(
 (num_jobs_remaining/float(jobs_per_hour_this_chunk)) *60,
1
 )
logger.info("\n\nWe're doing {} jobs per hour. At this rate, if we had to do everything up to limit, done in {}min".format(
int(jobs_per_hour_this_chunk),
predicted_mins_to_finish
 ))
logger.info("\t{} seconds this loop, {} chunks in {} seconds, {} seconds/chunk average\n".format(
elapsed(new_loop_start_time),
index,
elapsed(start_time),
round(elapsed(start_time)/float(index), 1)
 ))
exceptZeroDivisionError:
# logger.info(u"not printing status because divide by zero")
logger.info(".")


defrun_right_thing(self, parsed_args, job_type):
ifparsed_args.dynos!=None: # to tell the difference from setting to 0
self.scale_dyno(parsed_args.dynos, job_type)

ifparsed_args.reset:
self.reset_enqueued(job_type)

ifparsed_args.status:
self.print_status(job_type)

ifparsed_args.monitor:
self.monitor_till_done(job_type)
self.scale_dyno(0, job_type)

ifparsed_args.logs:
self.print_logs(job_type)

ifparsed_args.kick:
self.kick(job_type)

ifparsed_args.idorparsed_args.doiorparsed_args.run:
self.run(parsed_args, job_type)


## these are overwritten by main class

deftable_name(self, job_type):
pass

defprocess_name(self, job_type):
pass

defworker_run(self, **kwargs):
pass