| 43 | |
| 44 | |
| 45 | class ClusterHelper: |
| 46 | def __init__(self, cluster_info: ClusterInfo, workdir: Path): |
| 47 | self.cluster_info = cluster_info |
| 48 | self.workdir = workdir |
| 49 | self.workdir.mkdir(exist_ok=True, parents=True) |
| 50 | |
| 51 | self.cluster = Cluster(str(self.workdir)) |
| 52 | |
| 53 | @property |
| 54 | def active_nodes(self) -> List[str]: |
| 55 | return list(self.cluster.nodes.keys()) |
| 56 | |
| 57 | @property |
| 58 | def processes(self) -> List[ProcessInfo]: |
| 59 | processes = [] |
| 60 | for node in self.cluster.nodes.values(): |
| 61 | processes += node.processes |
| 62 | return processes |
| 63 | |
| 64 | def wait_for_process_end( |
| 65 | self, filter_fn: Callable[[ProcessInfo], bool], duration: datetime.timedelta = datetime.timedelta(seconds=5) |
| 66 | ): |
| 67 | """ |
| 68 | Wait until processed that pass through the given `filter_fn` are stopped. |
| 69 | :param filter_fn: Filter function to select a process. |
| 70 | :param duration: How long to wait (for all processes together) at most. |
| 71 | """ |
| 72 | start = time.time() |
| 73 | for process_info in self.processes: |
| 74 | if is_local(process_info.hostname) and filter_fn(process_info): |
| 75 | try: |
| 76 | process = psutil.Process(process_info.pid) |
| 77 | except psutil.NoSuchProcess: |
| 78 | logging.warning(f"Process {process_info.pid} has already stopped") |
| 79 | continue |
| 80 | |
| 81 | while process.is_running(): |
| 82 | if time.time() - start < duration.total_seconds(): |
| 83 | time.sleep(0.1) |
| 84 | else: |
| 85 | return |
| 86 | |
| 87 | def commit(self): |
| 88 | with open(self.workdir / CLUSTER_FILENAME, "w") as f: |
| 89 | self.cluster.serialize(f) |
| 90 | |
| 91 | def stop(self, use_sigint=False): |
| 92 | start = time.time() |
| 93 | |
| 94 | fn = functools.partial(kill_fn, use_sigint) |
| 95 | self.cluster.kill(fn) |
| 96 | logging.debug(f"Cluster killed in {time.time() - start} seconds") |
| 97 | |
| 98 | def start_processes(self, processes: List[StartProcessArgs]): |
| 99 | def prepare_workdir(workdir: Path) -> Path: |
| 100 | workdir = workdir if workdir else self.workdir |
| 101 | workdir.mkdir(parents=True, exist_ok=True) |
| 102 | return workdir.absolute() |