JIT Cleanup

[aaronzedwick]

Closes #969

Overview

Cleans up jit usage, standardizes jit cache, and removes redundant code.

PR Checklist

General

• An issue is linked created and linked
• Add appropriate labels
• Filled out Overview and Expected Usage (if applicable) sections

[philipc2]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

[aaronzedwick]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

[rajeeja]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

+1 we should get rid of ENABLE_JIT_CACHE as you mention. For disable_numba what if we disable it at first - use case "codecov", in our current implementatation we don't have disable_numba.

[philipc2]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

Yeah, I agee and don't really see much purpose in having a constant. Some functions can benefit from caching, while others may not. This should be something that we handle on a per-function basis.

[aaronzedwick]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

+1 we should get rid of ENABLE_JIT_CACHE as you mention. For disable_numba what if we disable it at first - use case "codecov", in our current implementatation we don't have disable_numba.

Disabling numba doesn't seem practical, or even really possible. We would have to somehow disable it before all the tests ran (no idea if that is possible, I don't know that it is but I am not well informed on the specifics of that), and it would slow down all the test cases. The CI would take forever to run since numba would be off for everything.

[rajeeja]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

+1 we should get rid of ENABLE_JIT_CACHE as you mention. For disable_numba what if we disable it at first - use case "codecov", in our current implementatation we don't have disable_numba.

Disabling numba doesn't seem practical, or even really possible. We would have to somehow disable it before all the tests ran (no idea if that is possible, I don't know that it is but I am not well informed on the specifics of that), and it would slow down all the test cases. The CI would take forever to run since numba would be off for everything.

I'll try some test to see, but disabling should be possible, even if it slows down one use case would be to get correct codecov and another to just check if sometime Numba has some issues going on with conversion to byte code etc.

[aaronzedwick]

This is great. I assume these global variables in Numba need to be set before importing UXarray right?

Well, they aren't really global variables in Numba. There isn't much point in having the disable_numba function, as we have already shown once the function is loaded it won't change anything. That would be a global variable in Numba, but it doesn't help much once the function is loaded. And ENABLE_JIT_CACHE is just our own constant we have in the constants file. Honestly, I would be in favor of removing that as well and just rewriting all the functions with jit to have @njit(cache=True). Since that is what we are doing anyway, I don't really see a reason to have an extra constant variable. We have never really even used it. It is basically the same as just saying cache=True.

+1 we should get rid of ENABLE_JIT_CACHE as you mention. For disable_numba what if we disable it at first - use case "codecov", in our current implementatation we don't have disable_numba.

Disabling numba doesn't seem practical, or even really possible. We would have to somehow disable it before all the tests ran (no idea if that is possible, I don't know that it is but I am not well informed on the specifics of that), and it would slow down all the test cases. The CI would take forever to run since numba would be off for everything.

I'll try some test to see, but disabling should be possible, even if it slows down one use case would be to get correct codecov and another to just check if sometime Numba has some issues going on with conversion to byte code etc.

We did try already. Once the function is loaded, it doesn't matter if numba is disabled after the function is loaded with the wrapper. Numba is either off or on it seems, for the whole package, and you'd have to disable it before the tests ran somehow. At least based on our testing. You would have to rewrap the functions somehow if you wanted to turn them on and off mid-function. Just seems like a lot of work, just to increase code coverage. As long as we make sure we are covering with tests the functions that numba is on, I don't see it being that useful to turn numba on/off.

[github-actions]

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change	Before [691999b] <v2024.10.0>	After [5f77bbf]	Ratio	Benchmark (Parameter)
-	442M	375M	0.85	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
-	442M	375M	0.85	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
-	445M	379M	0.85	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
-	444M	375M	0.84	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
-	6.77±0.09ms	5.99±0.09ms	0.89	mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
-	467M	372M	0.8	mpas_ocean.Integrate.peakmem_integrate('480km')

Benchmarks that have stayed the same:

Change	Before [691999b] <v2024.10.0>	After [5f77bbf]	Ratio	Benchmark (Parameter)
	1.58±0s	1.56±0.01s	0.99	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
	225±1ms	219±3ms	0.97	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
	1.99±0.01s	1.98±0.05s	0.99	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
	7.68±0.1ms	7.80±0.2ms	1.02	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
	2.98±0.01s	2.99±0.02s	1.00	import.Imports.timeraw_import_uxarray
	667±10μs	650±3μs	0.97	mpas_ocean.CheckNorm.time_check_norm('120km')
	436±3μs	431±3μs	0.99	mpas_ocean.CheckNorm.time_check_norm('480km')
	641±7ms	640±9ms	1.00	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
	42.2±1ms	40.5±0.2ms	0.96	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
	1.58±0.2ms	1.62±0.09ms	1.03	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
	492±9μs	537±30μs	1.09	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
	3.50±0.02ms	3.58±0.02ms	1.02	mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
	3.59±0.01s	3.62±0.05s	1.01	mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
	227±2ms	226±0.2ms	0.99	mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
	1.20±0.01μs	1.30±0μs	1.09	mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
	306±2ns	299±0.9ns	0.98	mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
	815±1ns	791±7ns	0.97	mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
	293±6ns	285±4ns	0.97	mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 1)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 2)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 4)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 8)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 1)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 2)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 4)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 8)
	124±0.5ms	122±0.8ms	0.98	mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
	8.82±0.3ms	8.28±0.09ms	0.94	mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
	1.04±0.01s	1.05±0.01s	1.01	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
	54.9±0.3ms	52.0±0.4ms	0.95	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
	77.5±0.5ms	79.8±0.8ms	1.03	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
	5.23±0.1ms	5.08±0.07ms	0.97	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
	320M	318M	0.99	mpas_ocean.Gradient.peakmem_gradient('120km')
	296M	296M	1.00	mpas_ocean.Gradient.peakmem_gradient('480km')
	2.72±0.03ms	2.71±0.02ms	1.00	mpas_ocean.Gradient.time_gradient('120km')
	288±2μs	285±3μs	0.99	mpas_ocean.Gradient.time_gradient('480km')
	233±8μs	227±10μs	0.97	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
	123±2μs	122±0.8μs	0.99	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
	425M	389M	0.92	mpas_ocean.Integrate.peakmem_integrate('120km')
	175±0.5ms	175±0.7ms	1.00	mpas_ocean.Integrate.time_integrate('120km')
	12.0±0.04ms	11.9±0.04ms	0.99	mpas_ocean.Integrate.time_integrate('480km')
	344±3ms	345±2ms	1.01	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
	341±2ms	345±2ms	1.01	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
	347±2ms	348±4ms	1.00	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
	23.1±0.3ms	22.5±0.07ms	0.97	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
	23.0±0.3ms	22.5±0.1ms	0.98	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
	23.1±0.3ms	22.5±0.1ms	0.97	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
	54.5±0.1ms	54.4±0.05ms	1.00	mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
	44.4±0.2ms	44.3±0.2ms	1.00	mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
	359±0.9ms	361±2ms	1.01	mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
	263±0.7ms	264±0.6ms	1.00	mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
	294M	292M	0.99	quad_hexagon.QuadHexagon.peakmem_open_dataset
	291M	291M	1.00	quad_hexagon.QuadHexagon.peakmem_open_grid
	6.15±0.02ms	6.27±0.03ms	1.02	quad_hexagon.QuadHexagon.time_open_dataset
	5.25±0.02ms	5.40±0.08ms	1.03	quad_hexagon.QuadHexagon.time_open_grid

[philipc2]

This was much needed. Thanks @aaronzedwick

[aaronzedwick]

This was much needed. Thanks @aaronzedwick

Do you know what’s up with the precommit? It works locally.

[philipc2]

This was much needed. Thanks @aaronzedwick

Do you know what’s up with the precommit? It works locally.

Nothing wrong on your end. See #991