make wait script more robust

This commit is contained in:
Shreyas Goenka 2025-03-02 16:48:01 +01:00
parent f249a1ac3b
commit 8213c4fa22
No known key found for this signature in database
GPG Key ID: 92A07DF49CCB0622
1 changed files with 36 additions and 23 deletions

View File

@ -5,45 +5,58 @@ import time
import platform
import subprocess
def is_finished_windows(pid):
assert int(pid) > 0
# Run task list but filter the list by pid..
output = subprocess.check_output(f"tasklist /fi \"PID eq {pid}\"", text=True)
# If tasklist does not find any process, that means the process we are
# waiting for has terminated.
unaliveMsg = "No tasks are running which match the specified criteria."
if unaliveMsg in output:
return True
return False
def is_finished_unix(pid):
assert int(pid) > 0
# Send signal 0 to the process to check if it exists.
# From the docs:
# If sig is 0, then no signal is sent, but existence and permission
# checks are still performed; this can be used to check for the
# existence of a process ID or process group ID that the caller is
# permitted to signal.
# ref: https://man7.org/linux/man-pages/man2/kill.2.html
try:
os.kill(pid, 0)
except OSError:
return True
return False
def wait_pid(pid):
max_attempts = 600 # 600 * 0.1 seconds = 1 minute
sleep_time = 0.1
for i in range(max_attempts):
# Check if we are on Windows or a Unix system.
if platform.system().lower() == "windows":
# Windows approach: use tasklist to check for the existence of the process.
try:
# Get the output of 'tasklist'
output = subprocess.check_output(["tasklist"], text=True)
except subprocess.CalledProcessError:
print("[wait_pid] Error retrieving tasklist. Assuming process has ended.")
return 0
# if the PID is not found in the list then assume the process ended.
if str(pid) not in output:
if is_finished_windows(pid):
print("[wait_pid] process has ended")
return 0
else:
# Linux/macOS approach: using os.kill with signal 0 to check if the process exists.
try:
os.kill(pid, 0)
except OSError:
if is_finished_unix(pid):
print("[wait_pid] process has ended")
return 0
time.sleep(sleep_time)
print(f"Timeout: Process {pid} did not end within 1 minute")
return 1
try:
pid = int(sys.argv[1])
except ValueError:
print("Error: PID must be an integer.")
sys.exit(1)
exit_code = wait_pid(pid)
exit_code = wait_pid(int(sys.argv[1]))
sys.exit(exit_code)