Merge pull request #197 from MilesCranmer/custom-julia-project

Allow custom shared projects for `julia_project`

Author: MilesCranmer

.github/workflows/CI.yml

@@ -69,6 +69,8 @@ jobs:
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
         run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_torch
+      - name: "Run custom env tests"
+        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_env
       - name: Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -121,6 +123,8 @@ jobs:
         run: python3 -m pip install coverage coveralls
       - name: "Run tests"
         run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test
+      - name: "Run custom env tests"
+        run: coverage run --append --source=pysr --omit='*/feynman_problems.py' -m unittest test.test_env
       - name: Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/CI_Windows.yml

@@ -63,3 +63,5 @@ jobs:
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
         run: python -m unittest test.test_torch
+      - name: "Run custom env tests"
+        run: python -m unittest test.test_env

.github/workflows/CI_conda_forge.yml

@@ -34,4 +34,4 @@ jobs:
       - name: "Install pysr-forge"
         run: conda activate test && mamba install pysr
       - name: "Run tests"
-        run: conda activate test && python3 -m unittest test.test
+        run: conda activate test && python3 -m unittest test.test && python3 -m unittest test.test_env

.github/workflows/CI_docker.yml

@@ -38,4 +38,4 @@ jobs:
       - name: Build docker
         run: docker build -t pysr --build-arg ARCH=${{ matrix.arch }} --build-arg VERSION=${{ matrix.julia-version }} --build-arg PYVERSION=${{ matrix.python-version }} .
       - name: Test docker
-        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test'
+        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test && python3 -m unittest test.test_env'

.github/workflows/CI_docker_large_nightly.yml

@@ -31,4 +31,4 @@ jobs:
       - name: Build docker
         run: docker build -t pysr --build-arg ARCH=${{ matrix.arch }} --build-arg VERSION=${{ matrix.julia-version }} --build-arg PYVERSION=${{ matrix.python-version }} .
       - name: Test docker
-        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test'
+        run: docker run --platform=${{ matrix.arch }} --rm pysr /bin/bash -c 'python3 -m unittest test.test && python3 -m unittest test.test_env'

.github/workflows/CI_large_nightly.yml

@@ -41,4 +41,4 @@ jobs:
             python setup.py install
             python -c 'import pysr; pysr.install()'
       - name: "Run tests"
-        run: python -m unittest test.test
+        run: python -m unittest test.test && python -m unittest test.test_env

.github/workflows/CI_mac.yml

@@ -67,3 +67,5 @@ jobs:
         run: pip install torch # (optional import)
       - name: "Run Torch tests"
         run: python -m unittest test.test_torch
+      - name: "Run custom env tests"
+        run: python -m unittest test.test_env

environment.yml

@@ -9,4 +9,5 @@ dependencies:
   - scikit-learn
   - setuptools
   - pyjulia
-  - openlibm<0.8.0
+  - openlibm
+  - openspecfun

pysr/julia_helpers.py

@@ -33,7 +33,9 @@ def _get_julia_env_dir():
     # Have to manually get env dir:
     try:
         julia_env_dir_str = subprocess.run(
-            ["julia", "-e using Pkg; print(Pkg.envdir())"], capture_output=True
+            ["julia", "-e using Pkg; print(Pkg.envdir())"],
+            capture_output=True,
+            env=os.environ,
         ).stdout.decode()
     except FileNotFoundError:
         env_path = os.environ["PATH"]
@@ -54,6 +56,12 @@ def _set_julia_project_env(julia_project, is_shared):
         os.environ["JULIA_PROJECT"] = str(julia_project)
 
 
+def _get_io_arg(quiet):
+    io = "devnull" if quiet else "stderr"
+    io_arg = f"io={io}" if is_julia_version_greater_eq(version=(1, 6, 0)) else ""
+    return io_arg
+
+
 def install(julia_project=None, quiet=False):  # pragma: no cover
     """
     Install PyCall.jl and all required dependencies for SymbolicRegression.jl.
@@ -64,28 +72,13 @@ def install(julia_project=None, quiet=False):  # pragma: no cover
 
     _version_assertion()
     # Set JULIA_PROJECT so that we install in the pysr environment
-    julia_project, is_shared = _process_julia_project(julia_project)
-    _set_julia_project_env(julia_project, is_shared)
+    processed_julia_project, is_shared = _process_julia_project(julia_project)
+    _set_julia_project_env(processed_julia_project, is_shared)
 
     julia.install(quiet=quiet)
+    Main = init_julia(julia_project, quiet=quiet)
+    io_arg = _get_io_arg(quiet)
 
-    if is_shared:
-        # is_shared is only true if the julia_project arg was None
-        # See _process_julia_project
-        Main = init_julia(None)
-    else:
-        Main = init_julia(julia_project)
-
-    Main.eval("using Pkg")
-
-    io = "devnull" if quiet else "stderr"
-    io_arg = f"io={io}" if is_julia_version_greater_eq(version=(1, 6, 0)) else ""
-
-    # Can't pass IO to Julia call as it evaluates to PyObject, so just directly
-    # use Main.eval:
-    Main.eval(
-        f'Pkg.activate("{_escape_filename(julia_project)}", shared = Bool({int(is_shared)}), {io_arg})'
-    )
     if is_shared:
         # Install SymbolicRegression.jl:
         _add_sr_to_julia_project(Main, io_arg)
@@ -117,11 +110,14 @@ def _import_error_string(julia_project=None):
 def _process_julia_project(julia_project):
     if julia_project is None:
         is_shared = True
-        julia_project = f"pysr-{__version__}"
+        processed_julia_project = f"pysr-{__version__}"
+    elif julia_project[0] == "@":
+        is_shared = True
+        processed_julia_project = julia_project[1:]
     else:
         is_shared = False
-        julia_project = Path(julia_project)
-    return julia_project, is_shared
+        processed_julia_project = Path(julia_project)
+    return processed_julia_project, is_shared
 
 
 def is_julia_version_greater_eq(juliainfo=None, version=(1, 6, 0)):
@@ -151,7 +147,7 @@ def _check_for_conflicting_libraries():  # pragma: no cover
         )
 
 
-def init_julia(julia_project=None):
+def init_julia(julia_project=None, quiet=False):
     """Initialize julia binary, turning off compiled modules if needed."""
     global julia_initialized
 
@@ -161,8 +157,8 @@ def init_julia(julia_project=None):
     from julia.core import JuliaInfo, UnsupportedPythonError
 
     _version_assertion()
-    julia_project, is_shared = _process_julia_project(julia_project)
-    _set_julia_project_env(julia_project, is_shared)
+    processed_julia_project, is_shared = _process_julia_project(julia_project)
+    _set_julia_project_env(processed_julia_project, is_shared)
 
     try:
         info = JuliaInfo.load(julia="julia")
@@ -189,11 +185,24 @@ def init_julia(julia_project=None):
 
         Main = _Main
 
+    if julia_initialized:
+        Main.eval("using Pkg")
+
+        io_arg = _get_io_arg(quiet)
+        # Can't pass IO to Julia call as it evaluates to PyObject, so just directly
+        # use Main.eval:
+        Main.eval(
+            f'Pkg.activate("{_escape_filename(processed_julia_project)}",'
+            f"shared = Bool({int(is_shared)}), "
+            f"{io_arg})"
+        )
+
     julia_initialized = True
     return Main
 
 
 def _add_sr_to_julia_project(Main, io_arg):
+    Main.eval("using Pkg")
     Main.sr_spec = Main.PackageSpec(
         name="SymbolicRegression",
         url="https://github.com/MilesCranmer/SymbolicRegression.jl",

test/test_env.py

@@ -0,0 +1,48 @@
+"""Contains tests for creating and initializing custom Julia projects."""
+
+import unittest
+import os
+from pysr import julia_helpers
+from tempfile import TemporaryDirectory
+
+
+class TestJuliaProject(unittest.TestCase):
+    """Various tests for working with Julia projects."""
+
+    def test_custom_shared_env(self):
+        """Test that we can use PySR in a custom shared env."""
+        with TemporaryDirectory() as tmpdir:
+            # Create a temp depot to store julia packages (and our custom env)
+            Main = julia_helpers.init_julia()
+
+            # Set up env:
+            if "JULIA_DEPOT_PATH" not in os.environ:
+                old_env = None
+                os.environ["JULIA_DEPOT_PATH"] = tmpdir
+            else:
+                old_env = os.environ["JULIA_DEPOT_PATH"]
+                os.environ[
+                    "JULIA_DEPOT_PATH"
+                ] = f"{tmpdir}:{os.environ['JULIA_DEPOT_PATH']}"
+            Main.eval(
+                f'pushfirst!(DEPOT_PATH, "{julia_helpers._escape_filename(tmpdir)}")'
+            )
+            test_env_name = "@pysr_test_env"
+            julia_helpers.install(julia_project=test_env_name)
+            Main = julia_helpers.init_julia(julia_project=test_env_name)
+
+            # Try to use env:
+            Main.eval("using SymbolicRegression")
+            Main.eval("using Pkg")
+
+            # Assert we actually loaded it:
+            cur_project_dir = Main.eval("splitdir(dirname(Base.active_project()))[1]")
+            potential_shared_project_dirs = Main.eval("Pkg.envdir(DEPOT_PATH[1])")
+            self.assertEqual(cur_project_dir, potential_shared_project_dirs)
+
+            # Clean up:
+            Main.eval("pop!(DEPOT_PATH)")
+            if old_env is None:
+                del os.environ["JULIA_DEPOT_PATH"]
+            else:
+                os.environ["JULIA_DEPOT_PATH"] = old_env