Request for assistance: Interacting directly with a jitted function

[jdebacker]

I'm working with @andrewjs00 to develop more units tests for the functions in calcfunctions.py.

We are running into difficulty understanding how to interact with the @iterate_jit-decorated functions directly and would appreciate some guidance on this.

I suggested @andrewjs00 start with calcfunctions.DependentCare because its logic is straightforward and it takes a relatively small set of arguments. The function appears to take a tuple with 9 arguments and so the following test was written:

def test_DependentCare():
    """
    Tests the DependentCare function in calcfunctions.py
    
    Note: The tests variable is only 1 of 5 tests I have created. This works
    when @iterate_jit is removed
    """
    test_tuple = (3, 2, 100000, 1, [250000, 500000, 250000, 500000, 250000], 
                  .2, 7165, 5000, 0)
    test_value = calcfunctions.DependentCare(*test_tuple)
    expected_value = 25196

    assert np.allclose(test_value, expected_value)

If the @iterate_jit is removed, this tests passed. However, with the decorator, the test fails with the following traceback:

(taxcalc-dev) jason.debacker@x86_64-apple-darwin13 taxcalc % pytest tests/test_calcfunctions.py::test_DependentCare
/Users/jason.debacker/anaconda3/envs/taxcalc-dev/lib/python3.8/site-packages/pep8.py:110: FutureWarning: Possible nested set at position 1
  EXTRANEOUS_WHITESPACE_REGEX = re.compile(r'[[({] | []}),;:]')
============================= test session starts ==============================
platform darwin -- Python 3.8.5, pytest-6.0.1, py-1.9.0, pluggy-0.13.1
rootdir: /Users/jason.debacker/repos/tax-calculator, configfile: pytest.ini
plugins: forked-1.2.0, xdist-2.1.0, pep8-1.0.6
collected 1 item

tests/test_calcfunctions.py F                                            [100%]

=================================== FAILURES ===================================
______________________________ test_DependentCare ______________________________

    def test_DependentCare():
        """
        Tests the DependentCare function in calcfunctions.py

        Note: The tests variable is only 1 of 5 tests I have created. This works
        when @iterate_jit is removed
        """
        test_tuple = (3, 2, 100000, 1, [250000, 500000, 250000, 500000, 250000],
                      .2, 7165, 5000, 0)
>       test_value = calcfunctions.DependentCare(*test_tuple)

tests/test_calcfunctions.py:176:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

args = (3, 2, 100000, 1, [250000, 500000, 250000, 500000, 250000], 0.2, ...)
kwargs = {}, in_arrays = [], pm_or_pf = [], farg = 'care_deduction'
high_level_func = "def hl_func(pm, pf):\n    from pandas import DataFrame\n    import numpy as np\n    import pandas as pd\n    def get_...      () = \\\n        applied_f()\n    header = ['care_deduction']\n    return DataFrame(data=outputs,columns=header)"
func_code = <code object <module> at 0x7fdcf27cd5b0, file "<string>", line 1>
fakeglobals = {'hl_func': <function hl_func at 0x7fdcf28f99d0>}
high_level_fn = <function hl_func at 0x7fdcf28f99d0>

    def wrapper(*args, **kwargs):
        """
        wrapper function nested in make_wrapper function nested
        in iterate_jit decorator.
        """
        in_arrays = []
        pm_or_pf = []
        print('ARGS = ', args)
        print('KWARGS = ', kwargs)
        for farg in all_out_args + in_args:
            if hasattr(args[0], farg):
                in_arrays.append(getattr(args[0], farg))
                pm_or_pf.append("pm")
            elif hasattr(args[1], farg):
                in_arrays.append(getattr(args[1], farg))
                pm_or_pf.append("pf")
        # Create the high level function
        high_level_func = create_toplevel_function_string(all_out_args,
                                                          list(in_args),
                                                          pm_or_pf)
        func_code = compile(high_level_func, "<string>", "exec")
        fakeglobals = {}
        eval(func_code,  # pylint: disable=eval-used
             {"applied_f": applied_jitted_f}, fakeglobals)
        high_level_fn = fakeglobals['hl_func']
>       ans = high_level_fn(*args, **kwargs)
E       TypeError: hl_func() takes 2 positional arguments but 9 were given

decorators.py:322: TypeError
----------------------------- Captured stdout call -----------------------------
ARGS =  (3, 2, 100000, 1, [250000, 500000, 250000, 500000, 250000], 0.2, 7165, 5000, 0)
KWARGS =  {}
=========================== short test summary info ============================
FAILED tests/test_calcfunctions.py::test_DependentCare - TypeError: hl_func()...
============================== 1 failed in 0.18s ===============================

We've tried some modifications in how to pass arguments to the DependentCare() function, but none has worked with the decorator in place. The docstrings and inline comments in the function or in the decorators.py module aren't helpful in figuring out our issue.

cc @MattHJensen @Peter-Metz @andersonfrailey

[MattHJensen]

Turning off jit for the tests may help. E.g., #2144

I can play around with this in the next few days if you're still blocked.

Thanks for working on this!

[Peter-Metz]

Hi @jdebacker and @andrewjs00. Seems like a great project. I'm not sure of the best way to interact with a jitted function, but an alternative approach to this sort of test is to create a Records object with the tax filer you'd like to test, create a Calculator object with that Records object, and run calc_all() to generate the care_deduction value. This approach is used in test_qbid_calculation(). Hope that helps.

[jdebacker]

@MattHJensen Thanks for linking to those issues - I guess this has been a problem in the past.

I tried adding:

importlib.reload(taxcalc.decorators)
taxcalc.decorators.DO_JIT = False

To my test, but while taxcalc.decorators.DO_JIT was now false, the issues were the same. I also tried to implement the functools.wraps function to help strip away decorators (as noted here), but couldn't get that to work successfully.

@Peter-Metz Thanks for this suggestion. I'd like to be able to deal with functions directly - one of the advantages being that you can more easily find the source of bugs where a particular value of a final output variable may be a result of several functions (which perhaps isn't the case in our application to DependentCare). But if that seems unworkable, this is a useful option.

[MattHJensen]

Can we close this issue now thanks to @chusloj's contribution in #2515?

[jdebacker]

@MattHJensen Yes, @chusloj has addressed this issue. Thanks!