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 platform
|
||||||
import subprocess
|
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):
|
def wait_pid(pid):
|
||||||
max_attempts = 600 # 600 * 0.1 seconds = 1 minute
|
max_attempts = 600 # 600 * 0.1 seconds = 1 minute
|
||||||
sleep_time = 0.1
|
sleep_time = 0.1
|
||||||
|
|
||||||
for i in range(max_attempts):
|
for i in range(max_attempts):
|
||||||
# Check if we are on Windows or a Unix system.
|
|
||||||
if platform.system().lower() == "windows":
|
if platform.system().lower() == "windows":
|
||||||
# Windows approach: use tasklist to check for the existence of the process.
|
if is_finished_windows(pid):
|
||||||
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:
|
|
||||||
print("[wait_pid] process has ended")
|
print("[wait_pid] process has ended")
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
else:
|
else:
|
||||||
# Linux/macOS approach: using os.kill with signal 0 to check if the process exists.
|
if is_finished_unix(pid):
|
||||||
try:
|
|
||||||
os.kill(pid, 0)
|
|
||||||
except OSError:
|
|
||||||
print("[wait_pid] process has ended")
|
print("[wait_pid] process has ended")
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
time.sleep(sleep_time)
|
time.sleep(sleep_time)
|
||||||
|
|
||||||
|
|
||||||
print(f"Timeout: Process {pid} did not end within 1 minute")
|
print(f"Timeout: Process {pid} did not end within 1 minute")
|
||||||
return 1
|
return 1
|
||||||
|
|
||||||
|
exit_code = wait_pid(int(sys.argv[1]))
|
||||||
try:
|
|
||||||
pid = int(sys.argv[1])
|
|
||||||
except ValueError:
|
|
||||||
print("Error: PID must be an integer.")
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
exit_code = wait_pid(pid)
|
|
||||||
sys.exit(exit_code)
|
sys.exit(exit_code)
|
||||||
|
|
Loading…
Reference in New Issue