-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add notifications and refactor pid tracking
- Loading branch information
Showing
12 changed files
with
121 additions
and
48 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| #!/bin/bash | ||
|
|
||
| set -eu | ||
|
|
||
| MESSAGE="${1-}" | ||
|
|
||
| if [ -z "${MESSAGE}" ]; then | ||
| echo "No message provided for notification, quitting early..." >&2 | ||
| exit 1 | ||
| fi | ||
|
|
||
| if [ -z "${PUSHBULLET_API_TOKEN}" ]; then | ||
| echo "No pushbullet token provided via \$PUSHBULLET_API_TOKEN variable, quitting early..." >&2 | ||
| exit 1 | ||
| fi | ||
|
|
||
| curl -u ${PUSHBULLET_API_TOKEN}: \ | ||
| -X POST https://api.pushbullet.com/v2/pushes \ | ||
| --header 'Content-Type: application/json' \ | ||
| --data-binary "{\"type\":\"note\",\"title\":\"sd-ultimate\",\"body\":\"${MESSAGE}\"}" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,6 @@ | ||
| #!/usr/bin/env bash | ||
| set -eu | ||
|
|
||
| echo "Starting Stable Diffusion Web UI" | ||
| cd ${A1111_ROOT} | ||
| source ./venv/bin/activate | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,6 @@ | ||
| #!/usr/bin/env bash | ||
| set -eu | ||
|
|
||
| echo "Starting InvokeAI..." | ||
| source ${INVOKEAI_ROOT}/venv/bin/activate | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,6 @@ | ||
| #!/usr/bin/env bash | ||
| set -eu | ||
|
|
||
| echo "Starting VS Server..." | ||
| PASSWORD=${VS_SERVER_PASSWORD} \ | ||
| nohup code-server \ | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| #!/bin/bash | ||
|
|
||
| set -eu | ||
|
|
||
| if [ -f /app/pid/track-training-pids.pid ];then | ||
| echo "Killing running instance of this script..." | ||
| rkill $(cat /app/pid/track-training-pids.pid) | ||
| fi | ||
|
|
||
| echo $$ > /app/pid/track-training-pids.pid | ||
|
|
||
| # This script runs in the background and tracks common long-running tasks | ||
| # which we might want to quit and/or notify the completion of. | ||
|
|
||
| mkdir -p /app/pid/training-pids | ||
|
|
||
| while true; do | ||
| # 1. locate new processes that match supported tasks (eg, kohya training scripts) | ||
| ps aux | grep "${TRAINING_PID_PATTERN}" | grep -v grep | while read -r line; do | ||
| # 2. extract PID and key info (eg model name) | ||
| pid=$(echo ${line} | awk '{print $2}') | ||
| model_name=$(echo ${line} | grep -Po "output_name=\K\S+" | sed "s/\s+/_/") # dedupe by kohya output names | ||
| [ -z "${model_name}" ] && continue | ||
|
|
||
| # 3. Add a file to /app/pid with PID and info string | ||
| pid_file="/app/pid/training-pids/${model_name}.pid" | ||
| [ -f ${pid_file} ] || (echo "${pid}" > ${pid_file}) | ||
| done | ||
|
|
||
| # 4. Check all currently tracked pids to see if any exited, if so quit and/or notify | ||
| for pid_file in /app/pid/training-pids/*.pid; do | ||
| model_name=$(basename ${pid_file} .pid) | ||
| pid=$(cat ${pid_file}) | ||
|
|
||
| if ! ps -p ${pid} > /dev/null; then | ||
| echo "Model ${model_name} finished training (PID no longer running)." | ||
| rm ${pid_file} | ||
| [ "${NOTIFY_ON_TRAINING_END}" = true ] && send-notification.sh "Model ${model_name} finished training." | ||
| [ "${SHUTDOWN_AFTER_TRAINING}" = true ] && runpodctl remove pod ${RUNPOD_POD_ID} | ||
| fi | ||
| done | ||
| sleep 30 | ||
| done |