Skip to content

Commit 5e0c850

Browse files
alexmalyshevfacebook-github-bot
authored andcommitted
Speed up tests in test_cinderjit
Summary: Some of these tests take way too long to run, making it hard to iterate on the test suite. The biggest issue is the use of `-X jit-all` in subprocesses, they end up compiling many functions that are ultimately unused. Reviewed By: grantlouisherman Differential Revision: D79673399 fbshipit-source-id: fded418988e17e66fdd1b07e513dac515baedc67
1 parent a92edcd commit 5e0c850

3 files changed

Lines changed: 84 additions & 112 deletions

File tree

cinderx/PythonLib/test_cinderx/cinder_preload_helper_main.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@ def main_func() -> str:
1818

1919

2020
print("disabling jit")
21-
cinderx.jit.precompile_all()
21+
# Force main_func to be compiled, which will trigger preloading.
22+
cinderx.jit.force_compile(main_func)
2223
cinderx.jit.disable()
2324

2425
print("jit disabled")

cinderx/PythonLib/test_cinderx/perf_fork_helper.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import os
66
import sys
77

8+
import cinderx.jit
9+
810

911
def compute():
1012
n = 0
@@ -47,5 +49,20 @@ def main():
4749
os.waitpid(pid2, 0)
4850

4951

52+
@cinderx.jit.jit_suppress
53+
def schedule_compilation() -> None:
54+
"""
55+
Set up functions to be compiled when they're first called. Processes that don't
56+
call them will leave them uncompiled.
57+
"""
58+
59+
cinderx.jit.lazy_compile(main)
60+
cinderx.jit.lazy_compile(parent)
61+
cinderx.jit.lazy_compile(child1)
62+
cinderx.jit.lazy_compile(child2)
63+
cinderx.jit.lazy_compile(compute)
64+
65+
5066
if __name__ == "__main__":
67+
schedule_compilation()
5168
sys.exit(main())

cinderx/PythonLib/test_cinderx/test_cinderjit.py

Lines changed: 65 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import weakref
1616

1717
from pathlib import Path
18+
from typing import Callable
1819

1920
AT_LEAST_312 = sys.version_info[:2] >= (3, 12)
2021

@@ -1609,8 +1610,10 @@ def junk{i}(j):
16091610
codepath = dirpath / "mod.py"
16101611
codepath.write_text(code)
16111612

1612-
def run_test(asserts_func, params):
1613-
args = [sys.executable, "-X", "jit-all"]
1613+
def run_test(
1614+
asserts_func: Callable[[list[str]], None], params: list[str]
1615+
) -> None:
1616+
args = [sys.executable]
16141617
args.extend(params)
16151618
args.append("mod.py")
16161619
proc = subprocess.run(
@@ -1624,14 +1627,14 @@ def run_test(asserts_func, params):
16241627
actual_stdout = [x.strip() for x in proc.stdout.split("\n")]
16251628
asserts_func(actual_stdout)
16261629

1627-
def zero_asserts(actual_stdout):
1630+
def zero_asserts(actual_stdout: list[str]) -> None:
16281631
expected_stdout = "max_size: 0"
16291632
self.assertEqual(actual_stdout[0], expected_stdout)
16301633
self.assertIn("used_size", actual_stdout[1])
16311634
used_size = int(actual_stdout[1].split(" ")[1])
16321635
self.assertGreater(used_size, 0)
16331636

1634-
def onek_asserts(actual_stdout):
1637+
def onek_asserts(actual_stdout: list[str]) -> None:
16351638
expected_stdout = "max_size: 1024"
16361639
self.assertEqual(actual_stdout[0], expected_stdout)
16371640
self.assertIn("used_size", actual_stdout[1])
@@ -1641,14 +1644,25 @@ def onek_asserts(actual_stdout):
16411644
# allocation is; we assume < 200K.
16421645
self.assertLess(used_size, 1024 * 200)
16431646

1644-
run_test(zero_asserts, ["-X", "jit-max-code-size=0"])
1647+
# Run the zero-assert tests with JitAuto=1000 to test "normal" behavior
1648+
# where we compile some code but don't have any limits to trip.
1649+
run_test(zero_asserts, ["-X", "jit-auto=1000", "-X", "jit-max-code-size=0"])
16451650
run_test(
16461651
zero_asserts,
1647-
["-X", "jit-max-code-size=0", "-X", "jit-huge-pages=0"],
1652+
[
1653+
"-X",
1654+
"jit-auto=1000",
1655+
"-X",
1656+
"jit-max-code-size=0",
1657+
"-X",
1658+
"jit-huge-pages=0",
1659+
],
16481660
)
16491661
run_test(
16501662
zero_asserts,
16511663
[
1664+
"-X",
1665+
"jit-auto=1000",
16521666
"-X",
16531667
"jit-max-code-size=0",
16541668
"-X",
@@ -1659,14 +1673,26 @@ def onek_asserts(actual_stdout):
16591673
"jit-cold-code-section-size=1048576",
16601674
],
16611675
)
1662-
run_test(onek_asserts, ["-X", "jit-max-code-size=1024"])
1676+
1677+
# Run the onek-assert tests with JitAll so that we quickly trip the limit
1678+
# and stop compiling.
1679+
run_test(onek_asserts, ["-X", "jit-all", "-X", "jit-max-code-size=1024"])
16631680
run_test(
16641681
onek_asserts,
1665-
["-X", "jit-max-code-size=1024", "-X", "jit-huge-pages=0"],
1682+
[
1683+
"-X",
1684+
"jit-all",
1685+
"-X",
1686+
"jit-max-code-size=1024",
1687+
"-X",
1688+
"jit-huge-pages=0",
1689+
],
16661690
)
16671691
run_test(
16681692
onek_asserts,
16691693
[
1694+
"-X",
1695+
"jit-all",
16701696
"-X",
16711697
"jit-max-code-size=1024",
16721698
"-X",
@@ -1691,7 +1717,12 @@ def test_max_code_size_fast(self) -> None:
16911717
codepath = dirpath / "mod.py"
16921718
codepath.write_text(code)
16931719

1694-
def run_proc():
1720+
def run_proc(size: str | None = None) -> str:
1721+
args = [sys.executable, "-X", "jit"]
1722+
if size:
1723+
args.extend(["-X", f"jit-max-code-size={size}"])
1724+
args.append("mod.py")
1725+
16951726
proc = subprocess.run(
16961727
args,
16971728
cwd=tmp,
@@ -1703,73 +1734,24 @@ def run_proc():
17031734
actual_stdout = [x.strip() for x in proc.stdout.split("\n")]
17041735
return actual_stdout[0]
17051736

1706-
args = [sys.executable, "-X", "jit-all", "mod.py"]
17071737
self.assertEqual(run_proc(), "max_size: 0")
1708-
args = [
1709-
sys.executable,
1710-
"-X",
1711-
"jit-all",
1712-
"-X",
1713-
"jit-max-code-size=1234567",
1714-
"mod.py",
1715-
]
1716-
self.assertEqual(run_proc(), "max_size: 1234567")
1717-
args = [
1718-
sys.executable,
1719-
"-X",
1720-
"jit-all",
1721-
"-X",
1722-
"jit-max-code-size=1k",
1723-
"mod.py",
1724-
]
1725-
self.assertEqual(run_proc(), "max_size: 1024")
1726-
args = [
1727-
sys.executable,
1728-
"-X",
1729-
"jit-all",
1730-
"-X",
1731-
"jit-max-code-size=1K",
1732-
"mod.py",
1733-
]
1734-
self.assertEqual(run_proc(), "max_size: 1024")
1735-
args = [
1736-
sys.executable,
1737-
"-X",
1738-
"jit-all",
1739-
"-X",
1740-
"jit-max-code-size=1m",
1741-
"mod.py",
1742-
]
1743-
self.assertEqual(run_proc(), "max_size: 1048576")
1744-
args = [
1745-
sys.executable,
1746-
"-X",
1747-
"jit-all",
1748-
"-X",
1749-
"jit-max-code-size=1M",
1750-
"mod.py",
1751-
]
1752-
self.assertEqual(run_proc(), "max_size: 1048576")
1753-
args = [
1754-
sys.executable,
1755-
"-X",
1756-
"jit-all",
1757-
"-X",
1758-
"jit-max-code-size=1g",
1759-
"mod.py",
1760-
]
1761-
self.assertEqual(run_proc(), "max_size: 1073741824")
1762-
args = [
1763-
sys.executable,
1764-
"-X",
1765-
"jit-all",
1766-
"-X",
1767-
"jit-max-code-size=1G",
1768-
"mod.py",
1769-
]
1770-
self.assertEqual(run_proc(), "max_size: 1073741824")
1771-
1772-
def run_proc():
1738+
self.assertEqual(run_proc("1234567"), "max_size: 1234567")
1739+
self.assertEqual(run_proc("1k"), "max_size: 1024")
1740+
self.assertEqual(run_proc("1K"), "max_size: 1024")
1741+
self.assertEqual(run_proc("1m"), "max_size: 1048576")
1742+
self.assertEqual(run_proc("1M"), "max_size: 1048576")
1743+
self.assertEqual(run_proc("1g"), "max_size: 1073741824")
1744+
self.assertEqual(run_proc("1G"), "max_size: 1073741824")
1745+
1746+
def run_proc(size: str) -> str:
1747+
args = [
1748+
sys.executable,
1749+
"-X",
1750+
"jit",
1751+
"-X",
1752+
f"jit-max-code-size={size}",
1753+
"mod.py",
1754+
]
17731755
proc = subprocess.run(
17741756
args,
17751757
cwd=tmp,
@@ -1780,44 +1762,16 @@ def run_proc():
17801762
self.assertEqual(proc.returncode, -6, proc)
17811763
return proc.stderr
17821764

1783-
args = [
1784-
sys.executable,
1785-
"-X",
1786-
"jit-all",
1787-
"-X",
1788-
"jit-max-code-size=-1",
1789-
"mod.py",
1790-
]
1791-
self.assertIn("Invalid unsigned integer in input string: '-1'", run_proc())
1792-
args = [
1793-
sys.executable,
1794-
"-X",
1795-
"jit-all",
1796-
"-X",
1797-
"jit-max-code-size=1.1",
1798-
"mod.py",
1799-
]
1800-
self.assertIn("Invalid unsigned integer in input string: '1.1'", run_proc())
1801-
args = [
1802-
sys.executable,
1803-
"-X",
1804-
"jit-all",
1805-
"-X",
1806-
"jit-max-code-size=dogs",
1807-
"mod.py",
1808-
]
1809-
self.assertIn("Invalid character in input string", run_proc())
1810-
args = [
1811-
sys.executable,
1812-
"-X",
1813-
"jit-all",
1814-
"-X",
1815-
"jit-max-code-size=1152921504606846976g",
1816-
"mod.py",
1817-
]
1765+
self.assertIn(
1766+
"Invalid unsigned integer in input string: '-1'", run_proc("-1")
1767+
)
1768+
self.assertIn(
1769+
"Invalid unsigned integer in input string: '1.1'", run_proc("1.1")
1770+
)
1771+
self.assertIn("Invalid character in input string", run_proc("dogs"))
18181772
self.assertIn(
18191773
"Unsigned Integer overflow in input string: '1152921504606846976g'",
1820-
run_proc(),
1774+
run_proc("1152921504606846976g"),
18211775
)
18221776

18231777

0 commit comments

Comments
 (0)