def init_compute(self, compute_module):
compute_imp = compute_module.ComputeCls
# Check that all of kernel interface is implemented.
systems_utils.check_implementation(ComputeInterface, compute_imp)
if getattr(compute_module, "RNG", None) is None:
raise Exception("No random number generator implemented "
"for compute module %s" % str(compute_module))
self.rng_cls = compute_module.RNG
# Collect implemented module functions.
module_functions = systems_utils.extract_functions(compute_imp)
# Collect function signatures.
function_signatures: dict = {}
required_methods = inspect.getmembers(ComputeInterface(),
predicate=inspect.ismethod)
for name, func in required_methods:
function_signatures[name] = func
for name, func in module_functions.items():
func_sig = function_signatures[name]
try:
remote_params = func_sig.remote_params
except Exception as _:
remote_params = {}
self.register(name, func, remote_params)
# Add functions as methods of this class.
for name, _ in module_functions.items():
self.methods[name] = self.get_callable(name)
def init(self):
# Compute available nodes, based on CPU resource.
local_ip = self.get_private_ip()
for node in ray.nodes():
node_key = list(
filter(lambda key: "node" in key, node["Resources"].keys()))
assert len(node_key) == 1
node_ip = node_key[0].split(":")[1]
has_cpu_resources = "CPU" in node[
"Resources"] and node["Resources"]["CPU"] >= 1.0
if local_ip == node_ip:
print("head node", node_ip)
self.head_node = node
if self.use_head and has_cpu_resources:
self.available_nodes.append(node)
elif has_cpu_resources:
print("worker node", node_ip)
self.available_nodes.append(node)
# Collect compute functions.
module_functions = extract_functions(self.compute_imp)
function_signatures: dict = {}
required_methods = inspect.getmembers(ComputeInterface(),
predicate=inspect.ismethod)
for name, func in required_methods:
function_signatures[name] = func
for name, func in module_functions.items():
func_sig = function_signatures[name]
try:
remote_params = func_sig.remote_params
except Exception as _:
remote_params = {}
self.remote_functions[name] = self.remote(func, remote_params)
def __init__(self, compute_module):
self.compute_module: ModuleType = compute_module
self.compute_imp = compute_module.ComputeCls
self.methods: dict = {}
# Check that all of kernel interface is implemented.
check_implementation(ComputeInterface, self.compute_imp)
# Collect implemented module functions.
self.module_functions = extract_functions(self.compute_imp)
if getattr(compute_module, "RNG", None) is None:
raise Exception("No random number generator implemented "
"for compute module %s" % str(compute_module))
self.rng_cls = compute_module.RNG
def __init__(self, num_gpus=None):
import cupy as cp
from nums.core.systems import cupy_compute
self.compute_imp_funcs = extract_functions(cupy_compute.ComputeCls)
for name in self.compute_imp_funcs:
raw_func = self.compute_imp_funcs[name]
def _func(raw_func):
@ray.remote(num_gpus=1)
def local_func(*args, **kwargs):
args = [cp.array(v) if isinstance(v, np.ndarray) else v for v in args]
kwargs = {k: cp.array(v) if isinstance(v, np.ndarray) else v for k, v in kwargs.items()}
return raw_func(*args, **kwargs).get()
self.compute_imp_funcs[name] = local_func
_func(raw_func)
super().__init__()
def __init__(self, num_gpus=None):
self.compute_imp_funcs = extract_functions(numpy_compute.ComputeCls)
for name in self.compute_imp_funcs:
self.compute_imp_funcs[name] = ray.remote(self.compute_imp_funcs[name])
super().__init__()
