mirror of https://github.com/databricks/cli.git
make wait script more robust
This commit is contained in:
parent
f249a1ac3b
commit
8213c4fa22
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue