17
17
import os
18
18
import psutil
19
19
import subprocess
20
+ import contextlib
20
21
import pathlib
21
22
import time
22
23
import typing
@@ -229,6 +230,24 @@ def callback_function(status_code: int, std_out: str, std_err: str) -> None:
229
230
name = f"{ identifier } _exit_status" , source = "user"
230
231
)
231
232
233
+ @property
234
+ def processes (self ) -> list [psutil .Process ]:
235
+ """Create an array containing a list of processes"""
236
+ if not self ._processes :
237
+ return []
238
+
239
+ _all_processes : list [psutil .Process ] = [
240
+ psutil .Process (process .pid ) for process in self ._processes .values ()
241
+ ]
242
+
243
+ with contextlib .suppress (psutil .NoSuchProcess , psutil .ZombieProcess ):
244
+ for process in _all_processes :
245
+ for child in process .children (recursive = True ):
246
+ if child not in _all_processes :
247
+ _all_processes .append (child )
248
+
249
+ return list (set (_all_processes ))
250
+
232
251
@property
233
252
def success (self ) -> int :
234
253
"""Return whether all attached processes completed successfully"""
@@ -320,15 +339,37 @@ def _save_output(self) -> None:
320
339
f"{ self ._runner .name } _{ proc_id } .out" , category = "output"
321
340
)
322
341
323
- def kill_process (self , process_id : str ) -> None :
324
- """Kill a running process by ID"""
325
- if not (process := self ._processes .get (process_id )):
326
- logger .error (
327
- f"Failed to terminate process '{ process_id } ', no such identifier."
328
- )
329
- return
342
+ def kill_process (
343
+ self , process_id : typing .Union [int , str ], kill_children_only : bool = False
344
+ ) -> None :
345
+ """Kill a running process by ID
330
346
331
- parent = psutil .Process (process .pid )
347
+ If argument is a string this is a process handled by the client,
348
+ else it is a PID of a external monitored process
349
+
350
+ Parameters
351
+ ----------
352
+ process_id : typing.Union[int, str]
353
+ either the identifier for a client created process or the PID
354
+ of an external process
355
+ kill_children_only : bool, optional
356
+ if process_id is an integer, whether to kill only its children
357
+ """
358
+ if isinstance (process_id , str ):
359
+ if not (process := self ._processes .get (process_id )):
360
+ logger .error (
361
+ f"Failed to terminate process '{ process_id } ', no such identifier."
362
+ )
363
+ return
364
+ try :
365
+ parent = psutil .Process (process .pid )
366
+ except psutil .NoSuchProcess :
367
+ return
368
+ elif isinstance (process_id , int ):
369
+ try :
370
+ parent = psutil .Process (process_id )
371
+ except psutil .NoSuchProcess :
372
+ return
332
373
333
374
for child in parent .children (recursive = True ):
334
375
logger .debug (f"Terminating child process { child .pid } : { child .name ()} " )
@@ -337,17 +378,13 @@ def kill_process(self, process_id: str) -> None:
337
378
for child in parent .children (recursive = True ):
338
379
child .wait ()
339
380
340
- logger .debug (f"Terminating child process { process .pid } : { process .args } " )
341
- process .kill ()
342
- process .wait ()
343
-
344
- if trigger := self ._completion_triggers .get (process_id ):
345
- trigger .set ()
346
-
347
- if trigger_process := self ._completion_processes .get (process_id ):
348
- trigger_process .join ()
381
+ if not kill_children_only :
382
+ logger .debug (f"Terminating process { process .pid } : { process .args } " )
383
+ process .kill ()
384
+ process .wait ()
349
385
350
- self ._execute_callback (process_id )
386
+ if isinstance (process_id , str ):
387
+ self ._execute_callback (process_id )
351
388
352
389
def kill_all (self ) -> None :
353
390
"""Kill all running processes"""
0 commit comments