Fall back to eager registry when needed

pysr/julia_extensions.py

@@ -3,6 +3,7 @@
 from typing import Literal
 
 from .julia_import import Pkg, jl
+from .julia_registry_helpers import try_with_registry_fallback
 from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
 
 
@@ -47,8 +48,12 @@ def isinstalled(uuid_s: str):
 
 def load_package(package_name: str, uuid_s: str) -> None:
     if not isinstalled(uuid_s):
-        Pkg.add(name=package_name, uuid=uuid_s)
-        Pkg.resolve()
+
+        def _add_package():
+            Pkg.add(name=package_name, uuid=uuid_s)
+            Pkg.resolve()
+
+        try_with_registry_fallback(_add_package)
 
     # TODO: Protect against loading the same symbol from two packages,
     #       maybe with a @gensym here.

pysr/julia_import.py

@@ -4,6 +4,8 @@
 from types import ModuleType
 from typing import cast
 
+from .julia_registry_helpers import try_with_registry_fallback
+
 # Check if JuliaCall is already loaded, and if so, warn the user
 # about the relevant environment variables. If not loaded,
 # set up sensible defaults.
@@ -42,6 +44,14 @@
     # Deprecated; so just pass to juliacall
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
+
+def import_juliacall():
+    import juliacall  # type: ignore
+
+
+try_with_registry_fallback(import_juliacall)
+
+
 from juliacall import AnyValue  # type: ignore
 from juliacall import VectorValue  # type: ignore
 from juliacall import Main as jl  # type: ignore

pysr/julia_registry_helpers.py

@@ -0,0 +1,44 @@
+"""Utilities for managing Julia registry preferences during package operations."""
+
+import os
+import warnings
+from collections.abc import Callable
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def try_with_registry_fallback(f: Callable[..., T], *args, **kwargs) -> T:
+    """Execute function with modified Julia registry preference.
+
+    First tries with existing registry preference. If that fails with a Julia registry error,
+    temporarily modifies the registry preference to 'eager'. Restores original preference after
+    execution.
+    """
+    try:
+        return f(*args, **kwargs)
+    except Exception as initial_error:
+        # Check if this is a Julia registry error by looking at the error message
+        error_str = str(initial_error)
+        if (
+            "JuliaError" not in error_str
+            or "Unsatisfiable requirements detected" not in error_str
+        ):
+            raise initial_error
+
+        old_value = os.environ.get("JULIA_PKG_SERVER_REGISTRY_PREFERENCE", None)
+        if old_value == "eager":
+            raise initial_error
+
+        warnings.warn(
+            "Initial Julia registry operation failed. Attempting to use the `eager` registry flavor of the Julia "
+            "General registry from the Julia Pkg server (via the `JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment variable)."
+        )
+        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
+        try:
+            return f(*args, **kwargs)
+        finally:
+            if old_value is not None:
+                os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value
+            else:
+                del os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]