Dynamic python code

[rburghol]

Goal: have a set of functions defined by local code that are executable in @njit. see respec#113

• Auto load files with name [riverseg].py - here: lines 342-359
• Currently testing with branch state on deq1
  • test seg: PL3_5250_0001
  • '/home/rob/working/modeling/hsp2/river_wd/PL3_5250_0001.py'
  • If theres a PL3_5250_0001.py file in that directory, it'll load the file and the functioned defined within
    • Functions defined in that file
      • That way you can have custom equations for individual riversegs
      • And it allows faster development without making pull requests
  • Within PL3_5250_0001.py is the function om_test_fn()
• @jdkleiner this is now working nicely.
• Performance Notes: Performance lags a bit since it recompiles the dynamic code every time (adds 7-8 seconds), so maybe a necessary evil, or maybe only to be used for testing or when we simply can't make it another way. Lag quantification and potential ways of unerstanding this lag.
  • Performance comparison:
    • branch fixjup (develop with install fix for bad jupyter line) takes 7 seconds (:15 on first compile)
    • branch spects loads 20-30 OM components and takes 12 seconds (:22 on 1st compile)
    • branch state loads 0 OM, and takes 17 seconds each time (no diff w/1st compile), even if there is no dynamic code loaded, or empty dynamic code function.
    • branch statex used to disable, then sequentially enable things (run from 1984-2001):
      1. 9 seconds with state loaded and added as arg to hydr(), and does setting state_ix[o1_ix' ... each timestep, but no computations
      2. 9 seconds with state loaded, dynamic module loaded in main, and all state loading stuff in HYDR, but no njoi module loading, nor any code in hydr executed.
      3. 9.1 seconds with b + load dynamic module in HYDR
      4. 9.1 seconds d + pass state_info to _hydr_()
      5. 9.0 seconds e + pass state_paths
      6. 15.3 f + pass all other Dicts and state_step_hydr
      7. 9.0 seconds g + do not pass state_step_hydr
      8. 9.0 seconds h + added and set all state data in _hydr_()
  • Known and potential overhead:
    • Passing in the function accounts for over 6.x seconds
    • Executing the dynamic code accounts for 1.x seconds
    • 80% of performance hit involved with dynamic code loading is simply passing the dynamic function into _hydr_
    • Performance hit does not scale with increased time span, so it is a 1-time incursion each model run.
• main.py
  • calls load_dynamics() (need to move this from main.py since the dynamic code must be loaded in the same module that it is called in, however, maybe we can also pass it in as part of a standard python dictionary, but then split it out as an individual function when passing t
• hsp2_local_py = dynamic_module_import(fbase, local_path + "/" + fbase + ".py", "hsp2_local_py")
  • This imports the hsp2_local_py module
• hsp2_local_py.state_step_hydr(model_exec_list, op_tokens, state_ix, dict_ix, ts_ix, hydr_ix, step)

Loading and Executing @njit code,

Since all functions called by njit code must be njit also, the dynamic functions must be defined before the first call of the larger routine, such as _hydr_().

• njit function must be loaded prior to calling hydr()`

Example Withdrawal

• Using test h5 file PL3_5250_0001.h5
• Create a file with the following py code PL3_5250_0001.py:

import sys
from numba import int8, float32, njit, types, typed # import the types
print("Loaded a set of  HSP2 code!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")

@njit
def state_step_hydr(state_ix, dict_ix, ts_ix, hydr_ix, step):
    fn_defined = False
    state_ix[hydr_ix['O1']] = 0.0 * 1.547
    if (step <= 1):
        print("Custom state_step_hydr() called")
        print("state at start", state_ix)
    if (step > 24*40*364):
        print("state at end", state_ix)
    return

• Then run hsp2 and export hydr:

hsp2 hsp2 run PL3_5250_0001.h5
Rscript /opt/model/meta_model/scripts/h5/export_hsp_h5.R PL3_5250_0001.h5 hydr_10mgd.csv "//RESULTS/RCHRES_R001/HYDR/table"

• Then edit PL3_5250_0001.py and set the withdrawa (O1)l to 0.0
• Then run and export again:

hsp2 hsp2 run PL3_5250_0001.h5
Rscript /opt/model/meta_model/scripts/h5/export_hsp_h5.R PL3_5250_0001.h5 hydr.csv "//RESULTS/RCHRES_R001/HYDR/table"

[jdkleiner]

Notes 6/21/23:

These notes have been incorporated into the body of this issue, and in a separate issue #67

[rburghol]

This allows you to test manually at the command line. More illuminating is the code in the body of the issue.

• cd to a directory where you have a [rivername].h5 file and [rivername].py file
• Enter python shell python3
• Paste the code for function dynamic_module_import() from https://github.com/HARPgroup/HSPsquared/blob/state/HSP2/state.py
• Do the code below

from numba.typed import Dict
from numba import types
import numpy as np
os.chdir("/opt/model/HSPsquared")
from HSP2.utilities_specl import *
from HSP2.SPECL import specl, _specl_
from HSP2.om_model_object import *
from HSP2.om_equation import *
from HSP2.om_data_matrix import *
from HSP2.om_sim_timer import *
from HSP2.om_model_linkage import ModelLinkage, step_model_link
from HSP2.om_model_broadcast import *

hydr_ix = Dict.empty(key_type=types.unicode_type, value_type=types.int64)
op_tokens = Dict.empty(key_type=types.int64, value_type=types.i8[:])
state_paths = Dict.empty(key_type=types.unicode_type, value_type=types.int64)
state_ix = Dict.empty(key_type=types.int64, value_type=types.float64)
dict_ix = Dict.empty(key_type=types.int64, value_type=types.float64[:,:])
ts_ix = Dict.empty(key_type=types.int64, value_type=types.float64[:])

local_path = os.getcwd()
fbase="PL3_5250_0001"
numdict = numdict=np.asarray([],dtype="float32")
io_manager, siminfo, op_tokens, state_paths, state_ix, dict_ix, ts_ix, model_object_cache = numdict,numdict,numdict,numdict,numdict,numdict,numdict,numdict

hsp2_local_py = dynamic_module_import(fbase, local_path + "/" + fbase + ".py", "hsp2_local_py")

# just set some test values
hydr_ix['O1'] = 1
from hsp2_local_py import state_step_hydr
state_step_hydr(numdict, op_tokens, state_ix, dict_ix, ts_ix, hydr_ix, 1)
# now show the values
state_ix[hydr_ix['O1']]

Sample code to put into a model [riverseg].py file.

import sys

print("Loaded a dynamic HSP2 code file!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")

def om_init_model(io_manager, siminfo, op_tokens, state_paths, state_ix, dict_ix, ts_ix, model_object_cache):
    print("Initialize om_init_model from code")

@njit
def state_step_hydr(model_exec_list, op_tokens, state_ix, dict_ix, ts_ix, hydr_ix, step):
    print("Do something")

print("Loaded state_step_hydr()")

[jdkleiner]

@rburghol Couple notes when running the test above:

• Need to import os before the line os.chdir("/opt/model/HSPsquared") i.e. just needs a import os
• When I try running I get NameError: name 'dynamic_module_import' is not defined, where does the function dynamic_module_import get loaded from?

[rburghol]

Hey @jdkleiner -- sorry, the test code does suffer from the lack you describe :) -- dynamic_module_import() is in the HSP2/STATE.py file. However, you may find the actual embedded code in hsp2 more interesting. That is the code that is in the body of this issue at the top. Also, in the state branch now there is a sample file in tests/test10/test10.py that will function fully, and should run out of the box in that branch on your install, provided of course that there is no incompatibility that I did not foresee in a windows versus linux install?

[rburghol]

• here: lines 342-359

[rburghol]

Tested on:

• Linux, Python 3.8.10
• Windows Python 3.8.9
  • powershell works
  • git bash.exe works also