Last active
July 17, 2020 13:08
-
-
Save davidlange6/3b0cb365aac669a714d9f288b0bf420d to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
# usage | |
# edm_pset_tweak.py --input_pkl RunPromptRecoCfg.pkl --output_pkl ouput.pkl --json tweaks.json | |
# | |
# tweaks.json can be | |
# [{ "process.maxEvents.input" : 100}] | |
# or | |
# { "process.maxEvents.input" : 100} | |
# or | |
# [{ "maxEvents.input" : 100}] | |
# or | |
# { "maxEvents.input" : 100} | |
# | |
# properly handles multiple pkls and mulitple jsons (all jsons are applied to all pkls) | |
# | |
import pickle | |
import argparse | |
import sys | |
import json | |
def apply_tweak(process, key, value): | |
key_split = key.split('.') | |
param = process | |
if key_split[0] == "process": | |
key_split = key_split[1:] | |
for p in key_split: | |
param = getattr(param, p) | |
if param is None: | |
return 1 | |
param = value | |
return 0 | |
def init_argparse(): | |
parser = argparse.ArgumentParser( | |
usage="%(prog)s [OPTION] [FILE]...", | |
description="Tweak pset(s) based on json list of tweaks to apply and write out resulting pset(s)" | |
) | |
parser.add_argument('--json', nargs='+', required=True) | |
parser.add_argument('--input_pkl', nargs='+', required=True) | |
parser.add_argument('--output_pkl', nargs='+', required=False) | |
parser.add_argument('--allow_failed_tweaks', action='store_true') | |
return parser | |
def main(): | |
parser = init_argparse() | |
args = parser.parse_args() | |
if len(args.output_pkl) != 0: | |
output_file_names = args.output_pkl | |
if len(args.input_pkl) != len(args.output_pkl): | |
print( | |
"Incorrect arguments. --input_pkl and --output_pkl must be the same length") | |
sys.exit(1) | |
else: | |
output_file_names=args.input_pkl | |
tweaks = [] | |
for json_file_name in args.json: | |
try: | |
with open(json_file_name) as json_file: | |
json_data = json.load(json_file) | |
except Exception as e: | |
print("Error opening file "+json_file_name) | |
sys.exit(1) | |
if not isinstance(json_data,(dict,list)): | |
print("Error loading json "+json_file_name) | |
sys.exit(1) | |
if isinstance(json_data, list): | |
for d in json_data: | |
for key in d: | |
tweaks.append((key, d[key])) | |
elif isinstance(json_data, dict): | |
for key in json_data: | |
tweaks.append((key, json_data[key])) | |
else: | |
print("Error loading json "+json_file_name) | |
sys.exit(1) | |
for i in range(len(args.input_pkl)): | |
with open(args.input_pkl[i], "rb") as pkl_file: | |
try: | |
process = pickle.load(pkl_file) | |
except (IOError, OSError, pickle.PickleError, pickle.UnpicklingError): | |
print("Not a valid pickle file " + | |
args.input_pkl[i]) + "- stopping" | |
sys.exit(1) | |
for tweak in tweaks: | |
err_val = apply_tweak(process, tweak[0], tweak[1]) | |
if err_val != 0: | |
print("Tweak not applied "+tweak[0]+" "+str(tweak[1])) | |
if not args.allow_failed_tweaks: | |
sys.exit(1) | |
else: | |
print("Tweak applied "+tweak[0]+" "+str(tweak[1])) | |
with open(output_file_names[i], "wb") as output_file: | |
if output_file.closed: | |
print("Error loading pickle input "+args.output_pkl[i]) | |
sys.exit(1) | |
pickle.dump(process, output_file, protocol=0) | |
main() |
I believe there is a little mistake here:
if isinstance(json_data,list) and isinstance(json_data,dict):
print("Error loading json "+json_file_name)
sys.exit(1)
It should probably be (the negation):
if not isinstance(json_data, (list, dict)):
print("Error loading json " + json_file_name)
sys.exit(1)
Or even if you bring back elif isinstance(json_data, dict)
, error can be moved to else
:
if isinstance(json_data, list):
for d in json_data:
for key in d:
tweaks.append((key, d[key]))
elif isinstance(json_data, dict):
for key in json_data:
tweaks.append((key, json_data[key]))
else:
print("Error loading json "+json_file_name)
sys.exit(1)
But this is not necessary, just might save one isinstance
.
definitely a bug.. I'll take your version
now to write some tests for all the different cases..
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
updated for --output_pkl and isinstance
(still mostly untested)