FluMapModel/upload_models.py

72 строки
2.9 KiB
Python

import argparse
import csv
from logging import getLogger
import requests
import os
from tqdm import tqdm
logger = getLogger()
def upload_model(model, api_url, models_path, api_key, continue_on_error=False):
headers = dict(Authorization=f'Bearer {api_key}')
model_path = os.path.join(models_path, f"{model['filename']}.csv")
if len(model['rds']) > 0:
rds_path = os.path.join(models_path, os.path.basename(model['rds']))
else:
rds_path = None
model_data = {
"id": model['filename'],
"name": model['name'],
"query_str": model['queryJSON'],
"model_type": model['type'],
"created": model['created'],
}
try:
files = {
'model': (os.path.basename(model_path), open(model_path, 'rb'), 'text/plain')
}
if rds_path and os.path.exists(rds_path):
files['rds'] = (os.path.basename(rds_path), open(rds_path, 'rb'), 'application/octet-stream')
r = requests.post(api_url, data=model_data, headers=headers, files=files)
if r.status_code != 201 and not continue_on_error:
logger.error(f"Failed to upload: {model['filename']}")
raise Exception(f"upload failed of {model['filename']}: {r.status_code}: {r.content}")
elif r.status_code != 201:
logger.warning(f"Failed to upload: {model['filename']}: {r.status_code}: {r.content}")
except FileNotFoundError as e:
if not continue_on_error:
raise e
else:
logger.warning(f"Could not find file from modelDB.tsv for {model['name']}")
def get_models(filename):
with open(filename) as tsvfile:
reader = csv.DictReader(tsvfile, dialect='excel-tab')
rows = [row for row in reader]
return rows
if __name__ == "__main__":
default_model_store_path = os.path.abspath(os.path.join( os.path.os.getcwd(), 'test_model_store'))
parser = argparse.ArgumentParser(description='Uploads trained SF Models to production')
parser.add_argument("--db-file", default=os.path.join(default_model_store_path, "modelDB.tsv"),
help="Where the modelDB.tsv produced during training is stored")
parser.add_argument("--model-store", default=default_model_store_path)
parser.add_argument("--fail-on-error",dest='onerror', default=True, action='store_false',
help='Allows the procressing of files to continue even in the case of an error')
parser.add_argument("--api-url", default="http://40.112.165.255/v1/generic_models",
help="URL for Seattle FLU API Incidence Mapper Model Server API")
parser.add_argument("--api-key", help="API-KEY Allowing uploads")
args = parser.parse_args()
models = get_models(args.db_file)
pbar = tqdm(models)
for model in pbar:
pbar.set_description(f"Processing {model['name']}")
upload_model(model, args.api_url, args.model_store, args.api_key, args.onerror)