13483409665

Committed 23 Feb 2025 01:16PM UTC coverage: 92.664% (-0.7%) from 93.372%

Build # 13483409665

Build Type

Pull #794

github

Committed by

MilesCranmer

Commit Message

Merge branch 'master' into fix-cluster-manager

Pull Request Pull Request #794: Automatic slurm allocations

Run Details

11 of 26 new or added lines in 3 files covered. (42.31%)

1 existing line in 1 file now uncovered.

1440 of 1554 relevant lines covered (92.66%)

2.61 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.29

/pysr/julia_extensions.py

"""This file installs and loads extensions for SymbolicRegression."""

from typing import Literal

from .julia_helpers import KNOWN_CLUSTERMANAGER_BACKENDS
from .julia_import import Pkg, jl
from .julia_registry_helpers import try_with_registry_fallback
from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec

PACKAGE_UUIDS = {
    "LoopVectorization": "bdcacae8-1622-11e9-2a5c-532679323890",
    "Bumper": "8ce10254-0962-460f-a3d8-1f77fea1446e",
    "Zygote": "e88e6eb3-aa80-5325-afca-941959d7151f",
    "SlurmClusterManager": "c82cd089-7bf7-41d7-976b-6b5d413cbe0a",
    "ClusterManagers": "34f1f09b-3a8b-5176-ab39-66d58a4d544e",
    "TensorBoardLogger": "899adc3e-224a-11e9-021f-63837185c80f",
}


def load_required_packages(
    *,
    turbo: bool = False,
    bumper: bool = False,
    autodiff_backend: Literal["Zygote"] | None = None,
    cluster_manager: str | None = None,
    logger_spec: AbstractLoggerSpec | None = None,
):
    if turbo:
        load_package("LoopVectorization")
    if bumper:
        load_package("Bumper")
    if autodiff_backend is not None:
        load_package("Zygote")
    if cluster_manager is not None:
        if cluster_manager == "slurm_native":
            load_package("SlurmClusterManager")
        elif cluster_manager in KNOWN_CLUSTERMANAGER_BACKENDS:
            load_package("ClusterManagers")
    if isinstance(logger_spec, TensorBoardLoggerSpec):
        load_package("TensorBoardLogger")


def load_all_packages():
    """Install and load all Julia extensions available to PySR."""
    for package_name, uuid_s in PACKAGE_UUIDS.items():
        load_package(package_name, uuid_s)


# TODO: Refactor this file so we can install all packages at once using `juliapkg`,
#       ideally parameterizable via the regular Python extras API


def isinstalled(uuid_s: str):
    return jl.haskey(Pkg.dependencies(), jl.Base.UUID(uuid_s))


def load_package(package_name: str, uuid_s: str | None = None) -> None:
    uuid_s = uuid_s or PACKAGE_UUIDS[package_name]
    if not isinstalled(uuid_s):

        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.
    jl.seval(f"using {package_name}: {package_name}")
    return None

1	"""This file installs and loads extensions for SymbolicRegression."""
2
3	from typing import Literal	3✔
4
5	from .julia_helpers import KNOWN_CLUSTERMANAGER_BACKENDS	3✔
6	from .julia_import import Pkg, jl	3✔
7	from .julia_registry_helpers import try_with_registry_fallback	3✔
8	from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec	3✔
9
10	PACKAGE_UUIDS = {	3✔
11	"LoopVectorization": "bdcacae8-1622-11e9-2a5c-532679323890",
12	"Bumper": "8ce10254-0962-460f-a3d8-1f77fea1446e",
13	"Zygote": "e88e6eb3-aa80-5325-afca-941959d7151f",
14	"SlurmClusterManager": "c82cd089-7bf7-41d7-976b-6b5d413cbe0a",
15	"ClusterManagers": "34f1f09b-3a8b-5176-ab39-66d58a4d544e",
16	"TensorBoardLogger": "899adc3e-224a-11e9-021f-63837185c80f",
17	}
18
19
20	def load_required_packages(	3✔
21	*,
22	turbo: bool = False,
23	bumper: bool = False,
24	autodiff_backend: Literal["Zygote"] \| None = None,
25	cluster_manager: str \| None = None,
26	logger_spec: AbstractLoggerSpec \| None = None,
27	):
28	if turbo:	3✔
29	load_package("LoopVectorization")	3✔
30	if bumper:	3✔
31	load_package("Bumper")	3✔
32	if autodiff_backend is not None:	3✔
NEW 33	load_package("Zygote")	×
34	if cluster_manager is not None:	3✔
NEW 35	if cluster_manager == "slurm_native":	×
NEW UNCOV 36	load_package("SlurmClusterManager")	×
NEW 37	elif cluster_manager in KNOWN_CLUSTERMANAGER_BACKENDS:	×
NEW 38	load_package("ClusterManagers")	×
39	if isinstance(logger_spec, TensorBoardLoggerSpec):	3✔
40	load_package("TensorBoardLogger")	3✔
41
42
43	def load_all_packages():	3✔
44	"""Install and load all Julia extensions available to PySR."""
45	for package_name, uuid_s in PACKAGE_UUIDS.items():	3✔
46	load_package(package_name, uuid_s)	3✔
47
48
49	# TODO: Refactor this file so we can install all packages at once using `juliapkg`,
50	# ideally parameterizable via the regular Python extras API
51
52
53	def isinstalled(uuid_s: str):	3✔
54	return jl.haskey(Pkg.dependencies(), jl.Base.UUID(uuid_s))	3✔
55
56
57	def load_package(package_name: str, uuid_s: str \| None = None) -> None:	3✔
58	uuid_s = uuid_s or PACKAGE_UUIDS[package_name]	3✔
59	if not isinstalled(uuid_s):	3✔
60
61	def _add_package():	3✔
62	Pkg.add(name=package_name, uuid=uuid_s)	3✔
63	Pkg.resolve()	3✔
64
65	try_with_registry_fallback(_add_package)	3✔
66
67	# TODO: Protect against loading the same symbol from two packages,
68	# maybe with a @gensym here.
69	jl.seval(f"using {package_name}: {package_name}")	3✔
70	return None	3✔

MilesCranmer / PySR / 13483409665

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous