Refactor experimental fast init mode (#73)

* refactor experimental fast init mode

* add ci test

* bump version "0.4.3"

.github/workflows/test.yml

@@ -15,21 +15,25 @@ on:
 
 # A workflow run is made up of one or more jobs that can run sequentially or in parallel
 jobs:
-  # This workflow contains a single job called "build"
   test:
-    name: Julia ${{ matrix.version }} - Python ${{ matrix.python-version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
-    runs-on: ${{ matrix.os }}
+    name: ${{ matrix.OS }} - Julia ${{ matrix.julia-version }} - Python ${{ matrix.python-version }} - ${{ github.event_name }}
+    runs-on: ubuntu-latest
     strategy:
-      fail-fast: false
       matrix:
-        version:
-          - '1.6'
-          - '1.7'
-        os:
-          - ubuntu-latest
-        arch:
-          - x64
-        python-version: ['3.7']
+        julia-version: ["1.6", "1.7"]
+        OS: [Ubuntu, macOS, Windows]
+        python-version: ["3.7"]
+        include:
+          - os: Ubuntu
+            image: ubuntu-22.04
+          - os: Windows
+            image: windows-2022
+          - os: macOS
+            image: macos-12
+      fail-fast: false
+    defaults:
+      run:
+        shell: bash
 
     # Steps represent a sequence of tasks that will be executed as part of the job
     steps:
@@ -40,7 +44,6 @@ jobs:
         uses: actions/setup-python@v4.2.0
         with:
           python-version: ${{ matrix.python-version }}
-          architecture: ${{ matrix.arch }}
 
       # Runs a single command using the runners shell
       - name: insatll jnumpy
@@ -58,8 +61,7 @@ jobs:
       - name: Setup Julia environment
         uses: julia-actions/setup-julia@v1
         with:
-          version: ${{ matrix.version }}
-          arch: ${{ matrix.arch }}
+          version: ${{ matrix.julia-version }}
 
       - uses: actions/cache@v1
         env:
@@ -72,8 +74,11 @@ jobs:
             ${{ runner.os }}-test-
             ${{ runner.os }}-
 
-      - name: run pytest
-        run: python -m pytest ./jnumpy/tests/test_core.py
+      - name: run pytest test_core.py
+        run: python -m pytest jnumpy/tests/test_core.py
+
+      - name: run pytest test_fast_init.py
+        run: python -m pytest jnumpy/tests/test_fast_init.py
 
       - name: Run Julia package tests
         uses: julia-actions/julia-runtest@v1

jnumpy/apis.py

@@ -85,7 +85,9 @@ def init_project(package_entry_filepath):
         with activate_project_nocheck(project_dir):
             jl_module_name = get_project_name_checked(project_dir)
             exec_julia(
-                "TyPython.CPython.@suppress_error begin import {0};TyPython.CPython.init();{0}.init() end".format(jl_module_name)
+                "TyPython.CPython.@suppress_error begin import {0};TyPython.CPython.init();{0}.init() end".format(
+                    jl_module_name
+                )
             )
     except JuliaError:
         with activate_project_checked(project_dir):

jnumpy/defaults.py

@@ -4,6 +4,7 @@
 import shutil
 import jnumpy.envars as envars
 import warnings
+import pathlib
 
 
 def get_jnumpy_dir():
@@ -24,7 +25,9 @@ def setup_julia_exe_():
     julia_exepath = os.getenv("TYPY_JL_EXE")
     if not julia_exepath:
         julia_exepath = get_default_julia_exe()
-    envars.SessionCtx.JULIA_EXE = julia_exepath
+    envars.SessionCtx.JULIA_EXE = (
+        pathlib.Path(julia_exepath).resolve().as_posix()
+    )  # resolve symlink
     return
 
 

jnumpy/envars.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 import pathlib
-import argparse
 
 CF_TYPY_MODE = "TYPY_MODE"
 CF_TYPY_PY_APIPTR = "TYPY_PY_APIPTR"
@@ -29,6 +28,3 @@ class SessionCtx:
     JULIA_EXE: str
     DEFAULT_PROJECT_DIR: str
     JULIA_START_OPTIONS: list[str]
-
-jl_opts_parse = argparse.ArgumentParser()
-jl_opts_parse.add_argument("-J", "--sysimage", type=str, default=None)

jnumpy/init.py

@@ -2,12 +2,15 @@
 import io
 import os
 import pathlib
-import sys
 import subprocess
 import ctypes
 import shlex
 import contextlib
-from .utils import escape_to_julia_rawstr
+from .utils import (
+    escape_to_julia_rawstr,
+    guess_julia_init_params,
+    sysimage_parser,
+)
 from .defaults import setup_julia_exe_, get_project_args
 from .envars import (
     CF_TYPY_JL_OPTS,
@@ -19,7 +22,6 @@
     TyPython_directory,
     InitTools_path,
     SessionCtx,
-    jl_opts_parse,
 )
 
 # XXX: adding an environment variable for fast debugging:
@@ -97,9 +99,8 @@ class JuliaError(Exception):
     pass
 
 
-def init_jl(experimental_fast_init=False):
-    # if experimental_fast_init is True, assume TyPython is included in sysimage
-    global _eval_jl
+def init_libjulia(init, experimental_fast_init=False):
+    # if experimental_fast_init is True, search libjulia by the releative path of julia executable
     if os.getenv(CF_TYPY_MODE) == CF_TYPY_MODE_JULIA:
         return
     elif os.getenv(CF_TYPY_MODE) == CF_TYPY_MODE_PYTHON:
@@ -110,7 +111,7 @@ def init_jl(experimental_fast_init=False):
             del os.environ[CF_TYPY_MODE]
             del os.environ[CF_TYPY_PY_APIPTR]
             del os.environ[CF_TYPY_PID]
-            init_jl()
+            init_libjulia(init, experimental_fast_init)
         return
     elif not os.getenv(CF_TYPY_MODE):
         os.environ[CF_TYPY_MODE] = CF_TYPY_MODE_PYTHON
@@ -123,37 +124,51 @@ def init_jl(experimental_fast_init=False):
     setup_julia_exe_()
     jl_opts = shlex.split(os.getenv(CF_TYPY_JL_OPTS, ""))
     jl_opts_proj = get_project_args()
-    opts, unkown_opts = jl_opts_parse.parse_known_args([jl_opts_proj, *jl_opts]) # parse arg --sysimage or -J
-    if experimental_fast_init:
+    SessionCtx.JULIA_START_OPTIONS = [jl_opts_proj, *jl_opts]
+    opts, unkown_opts = sysimage_parser.parse_known_args(
+        [jl_opts_proj, *jl_opts]
+    )  # parse arg --sysimage or -J
+
+    def fetch_julia_init_params(extra_cmd):
         cmd = [
             SessionCtx.JULIA_EXE,
-            *unkown_opts,
+            *extra_cmd,
             "--startup-file=no",
             "-O0",
             "--compile=min",
             "-e",
             julia_info_query,
         ]
+        bindir, libpath, sysimage, default_project_dir = subprocess.run(
+            cmd, check=True, capture_output=True, encoding="utf8"
+        ).stdout.splitlines()
+        return bindir, libpath, sysimage, default_project_dir
+
+    guess_result = guess_julia_init_params([jl_opts_proj, *jl_opts])
+    if not experimental_fast_init and opts.sysimage:
+        sysimage = pathlib.Path(opts.sysimage).absolute().as_posix()
+        bindir, libpath, _, default_project_dir = fetch_julia_init_params(unkown_opts)
+    elif not experimental_fast_init and not opts.sysimage:
+        bindir, libpath, sysimage, default_project_dir = fetch_julia_init_params(
+            unkown_opts
+        )
+    elif experimental_fast_init and opts.sysimage and guess_result:
+        sysimage = pathlib.Path(opts.sysimage).absolute().as_posix()
+        bindir, libpath, _, default_project_dir = guess_result
+    elif experimental_fast_init and opts.sysimage and not guess_result:
+        sysimage = pathlib.Path(opts.sysimage).absolute().as_posix()
+        bindir, libpath, _, default_project_dir = fetch_julia_init_params(unkown_opts)
+    elif experimental_fast_init and not opts.sysimage and guess_result:
+        bindir, libpath, sysimage, default_project_dir = guess_result
+    elif experimental_fast_init and not opts.sysimage and not guess_result:
+        bindir, libpath, sysimage, default_project_dir = fetch_julia_init_params(
+            unkown_opts
+        )
     else:
-        cmd = [
-            SessionCtx.JULIA_EXE,
-            jl_opts_proj,
-            *jl_opts,
-            "--startup-file=no",
-            "-O0",
-            "--compile=min",
-            "-e",
-            julia_info_query,
-        ]
-    bindir, libpath, sysimage, default_project_dir = subprocess.run(
-        cmd, check=True, capture_output=True, encoding="utf8"
-    ).stdout.splitlines()
-    SessionCtx.JULIA_START_OPTIONS = unkown_opts
+        raise Exception("Unknown error during fetch julia init params")
+
     SessionCtx.DEFAULT_PROJECT_DIR = default_project_dir
 
-    if experimental_fast_init and opts.sysimage:
-        sysimage_abs_paths = pathlib.Path(opts.sysimage).absolute().as_posix()
-        sysimage = sysimage_abs_paths
     old_cwd = os.getcwd()
     try:
         os.chdir(os.path.dirname(os.path.abspath(libpath)))
@@ -163,7 +178,9 @@ def init_jl(experimental_fast_init=False):
         except AttributeError:
             init_func = lib.jl_init_with_image__threading
 
-        argc, argv = args_from_config(SessionCtx.JULIA_EXE, SessionCtx.JULIA_START_OPTIONS)
+        argc, argv = args_from_config(
+            SessionCtx.JULIA_EXE, SessionCtx.JULIA_START_OPTIONS
+        )
         lib.jl_parse_opts(ctypes.pointer(argc), ctypes.pointer(argv))
 
         init_func.argtypes = [ctypes.c_char_p, ctypes.c_char_p]
@@ -173,54 +190,63 @@ def init_jl(experimental_fast_init=False):
         lib.jl_eval_string.restype = ctypes.c_void_p
         lib.jl_exception_clear.restype = None
 
-        def _eval_jl(x: str, use_gil: bool):
-            with contextlib.redirect_stderr(io.StringIO()) as ef:
-                if use_gil:
-                    source_code = gil_template.format(x)
-                else:
-                    source_code = no_gil_template.format(x)
-                source_code_bytes = source_code.encode("utf8")
-                if (
-                    not lib.jl_eval_string(source_code_bytes)
-                ) and lib.jl_exception_occurred():
-                    lib.jl_exception_clear()
-                    raise JuliaError(ef.getvalue())
-                return None
+        init(lib)
 
-        exec_julia(
-            f"""
-            import Pkg
-            Pkg.activate({escape_to_julia_rawstr(default_project_dir)}, io=devnull)
-            include({escape_to_julia_rawstr(InitTools_path)})
-        """,
-            use_gil=False,
-        )
+    finally:
+        os.chdir(old_cwd)
 
-        try:
-            exec_julia("import TyPython", use_gil=False)
-        except JuliaError:
-            try:
-                exec_julia(
-                    f"InitTools.setup_environment({escape_to_julia_rawstr(TyPython_directory)})",
-                    use_gil=False,
-                )
-            except JuliaError:
-                pass
-            exec_julia(
-                f"InitTools.force_resolve({escape_to_julia_rawstr(TyPython_directory)})",
-                use_gil=False,
-            )
+
+def init_jl_from_lib(lib):
+    global _eval_jl
+
+    def _eval_jl(x: str, use_gil: bool):
+        with contextlib.redirect_stderr(io.StringIO()) as ef:
+            if use_gil:
+                source_code = gil_template.format(x)
+            else:
+                source_code = no_gil_template.format(x)
+            source_code_bytes = source_code.encode("utf8")
+            if (
+                not lib.jl_eval_string(source_code_bytes)
+            ) and lib.jl_exception_occurred():
+                lib.jl_exception_clear()
+                raise JuliaError(ef.getvalue())
+            return None
+
+    exec_julia(
+        f"""
+        import Pkg
+        include({escape_to_julia_rawstr(InitTools_path)})
+    """,
+        use_gil=False,
+    )
+
+    try:
+        exec_julia("import TyPython", use_gil=False)
+    except JuliaError:
         try:
             exec_julia(
-                rf"""
-                import TyPython
-                import TyPython.CPython
-                TyPython.CPython.init()
-            """,
+                f"InitTools.setup_environment({escape_to_julia_rawstr(TyPython_directory)})",
                 use_gil=False,
             )
         except JuliaError:
-            raise RuntimeError("invalid julia initialization")
+            pass
+        exec_julia(
+            f"InitTools.force_resolve({escape_to_julia_rawstr(TyPython_directory)})",
+            use_gil=False,
+        )
+    try:
+        exec_julia(
+            rf"""
+            import TyPython
+            import TyPython.CPython
+            TyPython.CPython.init()
+        """,
+            use_gil=False,
+        )
+    except JuliaError:
+        raise RuntimeError("invalid julia initialization")
 
-    finally:
-        os.chdir(old_cwd)
+
+def init_jl(experimental_fast_init=False):
+    init_libjulia(init_jl_from_lib, experimental_fast_init)

jnumpy/tests/test_core.py

@@ -62,13 +62,15 @@ def test_mat_mul(dtype):
     assert actual.dtype == desired.dtype
     np.testing.assert_array_almost_equal(actual, desired, decimal=5)
 
+
 def test_bool_array():
     x = np.asarray([[True, False], [True, False]])
-    x_int = x.astype('i')
+    x_int = x.astype("i")
     actual = mat_mul(x, x)
     desired = x_int @ x_int
     assert np.all(actual == desired)
 
+
 def test_set_zero():
     x = np.random.rand(2)
     set_zero(x)
@@ -81,14 +83,13 @@ def test_fft():
     desired = np.fft.fft(x)
     np.testing.assert_array_almost_equal(actual, desired)
 
+
 def test_subprocess():
     assert os.getenv("TYPY_PY_APIPTR")
     cmd = [
         "python",
         "-c",
-        "import jnumpy as np; np.init_jl(); np.exec_julia(\"print(1)\")"
+        'import jnumpy as np; np.init_jl(); np.exec_julia("print(1)")',
     ]
-    out = subprocess.run(
-        cmd, check=True, capture_output=True, encoding="utf8"
-    )
+    out = subprocess.run(cmd, check=True, capture_output=True, encoding="utf8")
     assert out.stdout == "1"

jnumpy/tests/test_fast_init.py

@@ -0,0 +1,5 @@
+def test_fast_init():
+    import jnumpy as np
+    np.init_jl(experimental_fast_init=True)
+    np.exec_julia("println(:success)")
+    return