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

81.48

/pysr/julia_helpers.py

"""Functions for initializing the Julia environment and installing deps."""

from typing import Any, Callable, cast, overload

import numpy as np
from juliacall import convert as jl_convert  # type: ignore
from numpy.typing import NDArray

from .deprecated import init_julia, install
from .julia_import import AnyValue, jl

jl_convert = cast(Callable[[Any, Any], Any], jl_convert)

jl.seval("using Serialization: Serialization")
jl.seval("using PythonCall: PythonCall")

Serialization = jl.Serialization
PythonCall = jl.PythonCall

jl.seval("using SymbolicRegression: plus, sub, mult, div, pow")


def _escape_filename(filename):
    """Turn a path into a string with correctly escaped backslashes."""
    if filename is None:
        return None
    str_repr = str(filename)
    str_repr = str_repr.replace("\\", "\\\\")
    return str_repr


KNOWN_CLUSTERMANAGER_BACKENDS = ["slurm", "pbs", "lsf", "sge", "qrsh", "scyld", "htc"]


def load_cluster_manager(cluster_manager: str) -> AnyValue:
    if cluster_manager == "slurm_native":
        jl.seval("using SlurmClusterManager: SlurmManager")
        # TODO: Is this the right way to do this?
        jl.seval(f"using Distributed: addprocs")
        jl.seval("addprocs_slurm_native(args...; kws...) = addprocs(SlurmManager())")
        return jl.addprocs_slurm_native
    elif cluster_manager in KNOWN_CLUSTERMANAGER_BACKENDS:
        jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")
        return jl.seval(f"addprocs_{cluster_manager}")
    else:
        # Assume it's a function
        return jl.seval(cluster_manager)


def jl_array(x, dtype=None):
    if x is None:
        return None
    elif dtype is None:
        return jl_convert(jl.Array, x)
    else:
        return jl_convert(jl.Array[dtype], x)


def jl_dict(x):
    return jl_convert(jl.Dict, x)


def jl_is_function(f) -> bool:
    return cast(bool, jl.seval("op -> op isa Function")(f))


def jl_serialize(obj: Any) -> NDArray[np.uint8]:
    buf = jl.IOBuffer()
    Serialization.serialize(buf, obj)
    return np.array(jl.take_b(buf))


@overload
def jl_deserialize(s: NDArray[np.uint8]) -> AnyValue: ...
@overload
def jl_deserialize(s: None) -> None: ...
def jl_deserialize(s):
    if s is None:
        return s
    buf = jl.IOBuffer()
    jl.write(buf, jl_array(s))
    jl.seekstart(buf)
    return Serialization.deserialize(buf)

1	"""Functions for initializing the Julia environment and installing deps."""
2
3	from typing import Any, Callable, cast, overload	3✔
4
5	import numpy as np	3✔
6	from juliacall import convert as jl_convert # type: ignore	3✔
7	from numpy.typing import NDArray	3✔
8
9	from .deprecated import init_julia, install	3✔
10	from .julia_import import AnyValue, jl	3✔
11
12	jl_convert = cast(Callable[[Any, Any], Any], jl_convert)	3✔
13
14	jl.seval("using Serialization: Serialization")	3✔
15	jl.seval("using PythonCall: PythonCall")	3✔
16
17	Serialization = jl.Serialization	3✔
18	PythonCall = jl.PythonCall	3✔
19
20	jl.seval("using SymbolicRegression: plus, sub, mult, div, pow")	3✔
21
22
23	def _escape_filename(filename):	3✔
24	"""Turn a path into a string with correctly escaped backslashes."""
25	if filename is None:	3✔
26	return None	×
27	str_repr = str(filename)	3✔
28	str_repr = str_repr.replace("\\", "\\\\")	3✔
29	return str_repr	3✔
30
31
32	KNOWN_CLUSTERMANAGER_BACKENDS = ["slurm", "pbs", "lsf", "sge", "qrsh", "scyld", "htc"]	3✔
33
34
35	def load_cluster_manager(cluster_manager: str) -> AnyValue:	3✔
NEW 36	if cluster_manager == "slurm_native":	×
NEW 37	jl.seval("using SlurmClusterManager: SlurmManager")	×
38	# TODO: Is this the right way to do this?
NEW 39	jl.seval(f"using Distributed: addprocs")	×
NEW 40	jl.seval("addprocs_slurm_native(args...; kws...) = addprocs(SlurmManager())")	×
NEW 41	return jl.addprocs_slurm_native	×
NEW 42	elif cluster_manager in KNOWN_CLUSTERMANAGER_BACKENDS:	×
NEW 43	jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")	×
NEW 44	return jl.seval(f"addprocs_{cluster_manager}")	×
45	else:
46	# Assume it's a function
NEW 47	return jl.seval(cluster_manager)	×
48
49
50	def jl_array(x, dtype=None):	3✔
51	if x is None:	3✔
52	return None	3✔
53	elif dtype is None:	3✔
54	return jl_convert(jl.Array, x)	3✔
55	else:
56	return jl_convert(jl.Array[dtype], x)	3✔
57
58
59	def jl_dict(x):	3✔
60	return jl_convert(jl.Dict, x)	3✔
61
62
63	def jl_is_function(f) -> bool:	3✔
64	return cast(bool, jl.seval("op -> op isa Function")(f))	3✔
65
66
67	def jl_serialize(obj: Any) -> NDArray[np.uint8]:	3✔
68	buf = jl.IOBuffer()	3✔
69	Serialization.serialize(buf, obj)	3✔
70	return np.array(jl.take_b(buf))	3✔
71
72
73	@overload	3✔
74	def jl_deserialize(s: NDArray[np.uint8]) -> AnyValue: ...	3✔
75	@overload	3✔
76	def jl_deserialize(s: None) -> None: ...	3✔
77	def jl_deserialize(s):	3✔
78	if s is None:	3✔
79	return s	3✔
80	buf = jl.IOBuffer()	3✔
81	jl.write(buf, jl_array(s))	3✔
82	jl.seekstart(buf)	3✔
83	return Serialization.deserialize(buf)	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