Spaces:
Sleeping
Sleeping
from google_sheet import conf_repo, secret_repo | |
import json | |
from kaggle_service import KernelRerunService, NbJob | |
from logger import sheet_logger, get_now | |
from run_stt import run_stt_service | |
def get_secret_dict(): | |
# load all account and secret into a dict | |
secret_dict = {} | |
try: | |
for i in range(2, 50): # note: read row 2 - 50 | |
rs = secret_repo.read(i) | |
if not rs: | |
break | |
assert 'username' in rs | |
assert 'secret' in rs | |
username = rs['username'].rstrip() | |
secret = rs['secret'].rstrip() | |
secret_dict[username] = secret | |
except Exception as e: | |
sheet_logger.log(log="Get secret table failed!!" + str(e)) | |
return secret_dict | |
def get_config_list(secret_dict): | |
configs = [] | |
ids = [] | |
try: | |
for i in range(2, 50): # note: read row 2 - 50 | |
rs = conf_repo.read(i) | |
if not rs: | |
break | |
if not rs['config']: | |
reset_keys = [] | |
for key in rs: | |
if rs[key]: | |
reset_keys.append(key) | |
if len(reset_keys) > 0: | |
conf_repo.update(row_index=i, data={k: '' for k in reset_keys}) | |
break | |
print(i, rs) | |
try: | |
# validation | |
# assert 'config' in rs, 'require config column' | |
assert 'usernames' in rs, 'require usernames column' | |
assert 'enable' in rs, 'require enable column' | |
assert rs['enable'].rstrip() in ['enable', 'disable'], rs['enable'] | |
assert 'slug' in rs['config'], 'require config.slug' | |
cfg = json.loads(rs['config']) | |
usernames = set(rs['usernames'] | |
.rstrip().replace('\n', ' ').replace(',', ' ').replace(';', ' ').replace('|', ' ') | |
.split(' ')) | |
usernames = [u for u in usernames if u] | |
is_enabled = rs['enable'].rstrip() == 'enable' | |
if not is_enabled: | |
sheet_logger.log(log="Disabled, skip!", nb=cfg['slug']) | |
print("skip ", cfg['slug']) | |
continue | |
assert len(usernames) > 0, 'len usernames == 0' | |
# process config | |
accounts = {u: secret_dict[u] for u in secret_dict if u in usernames} | |
assert not set(usernames).difference(set(accounts.keys())), set(usernames).difference(set(accounts.keys())) | |
cfg = {**cfg, "accounts": accounts} | |
# save | |
configs.append(cfg) | |
ids.append(i) | |
print(cfg) | |
except AssertionError: | |
import traceback | |
sheet_logger.update_job_status(i, validate_status=str(traceback.format_exc())) | |
except Exception: | |
import traceback | |
sheet_logger.log(log="Get config failed!!" + str(traceback.format_exc())) | |
return configs, ids | |
def create_service(configs, ids): | |
service = KernelRerunService() | |
for config, idx in zip(configs, ids): | |
try: | |
service.add_job(NbJob.from_dict(config, id=idx)) | |
except Exception as e: | |
sheet_logger.update_job_status(idx, validate_status=str(e)) | |
return service | |
def main(): | |
sheet_logger.log("========= start ==========") | |
obj = run_stt_service.get_obj() | |
if run_stt_service.is_need_to_run(obj): | |
# start | |
run_stt_service.set_is_running({}) | |
secret_dict = get_secret_dict() | |
# load config from gg sheet file | |
configs, ids = get_config_list(secret_dict) | |
# add config to service | |
service = create_service(configs, ids) | |
# run service | |
try: | |
service.validate_all() | |
# service.status_all() | |
service.run_all() | |
except Exception as e: | |
sheet_logger.log(log=str(e)) | |
# done | |
run_stt_service.set_run_done() | |
else: | |
sheet_logger.log(f"Current time is {get_now()} is not after {obj.get('last_run')} {obj.get('auto_run_after_last_run')} or not require new run") | |
sheet_logger.log("========= end ==========") | |
if __name__ == "__main__": | |
main() | |