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__()