File size: 2,209 Bytes
3cdaa7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import traceback
import os
from inspect import signature
import pickle
import datetime
def pickle_function_call_wrapper(func, output_dir='pickled_inputs'):
i = 0
os.makedirs(output_dir)
# pickle.dump({'args': args, 'kwargs': kwargs}, fh)
def wrapper(*args, **kwargs):
"""
Wrap the original function call to print the arguments before
calling the intended function
"""
nonlocal i
i += 1
func_sig = signature(func)
# Create the argument binding so we can determine what
# parameters are given what values
argument_binding = func_sig.bind(*args, **kwargs)
argument_map = argument_binding.arguments
# Perform the print so that it shows the function name
# and arguments as a dictionary
path = os.path.join(output_dir, f'{i:05d}.pkl')
print(f"logging {func.__name__} arguments: {[k for k in argument_map]} to {path}")
argument_map['stack'] = traceback.format_stack()
for k, v in argument_map.items():
if hasattr(v, 'detach'):
argument_map[k] = v.cpu().detach()
with open(path, 'wb') as fh:
pickle.dump(argument_map, fh)
return func(*args, **kwargs)
return wrapper
def wrap_it(wrapper, instance, method, **kwargs):
class_method = getattr(instance, method)
wrapped_method = wrapper(class_method, **kwargs)
setattr(instance, method, wrapped_method)
def pickle_function_call(instance, method, subdir):
output_dir = os.path.join(os.getcwd(), 'pickled_inputs', subdir, datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S"))
wrap_it(pickle_function_call_wrapper, instance, method, output_dir=output_dir)
return output_dir
# For testing
if __name__=='__main__':
import glob
class Dog:
def __init__(self, name):
self.name = name
def bark(self, arg, kwarg=None):
print(f'{self.name}:{arg}:{kwarg}')
dog = Dog('fido')
dog.bark('ruff')
output_dir = pickle_function_call(dog, 'bark', 'debugging')
dog.bark('ruff', kwarg='wooof')
for p in glob.glob(os.path.join(output_dir, '*')):
print(p)
with open(p, 'rb') as fh:
print(pickle.load(fh))
|