109 lines
4.7 KiB
Bash
Executable file
109 lines
4.7 KiB
Bash
Executable file
#!/usr/bin/env bash
|
|
# docker-gpu-node.sh — shim that pp-orchestrator can spawn instead of the
|
|
# pp-worker binary directly. Boots one pp-worker container per stage
|
|
# on the host network so iroh can dial without NAT.
|
|
#
|
|
# Required env (forwarded by pp-orchestrator):
|
|
# STAGE, NUM_STAGES, SEED_ADDR, SEED_DIRECT, MAX_TOKENS
|
|
# Optional env (forwarded if present):
|
|
# MODEL, PP_WORKER_STUB, WORKER_CMD,
|
|
# PEER_NODE_ID, PEER_DIRECT,
|
|
# FIRST_PEER_NODE_ID, FIRST_PEER_DIRECT,
|
|
# PP_BOOT_DELAY_STAGE, PP_BOOT_DELAY_SECS,
|
|
# PP_FLEET_SINK (datastream fleet sink the stage ships telemetry to)
|
|
#
|
|
# Configurable via this shim's own environment:
|
|
# PP_IMAGE — image tag (default: swactor-pp-gpu:latest)
|
|
# PP_CONTAINER_PREFIX — name prefix (default: pp-stage)
|
|
# PP_DEV — tinygrad device override; empty/unset leaves the
|
|
# image's CUDA=1 default in place (default: CPU,
|
|
# which is right for both the stub image and
|
|
# workstation real-mode runs)
|
|
# PP_CACHE_DIR — host path for tinygrad's cache (default:
|
|
# $HOME/.cache/tinygrad). Created if missing.
|
|
# PP_STAGE_NETWORK — docker network mode for the stage container
|
|
# (default: host). Set to `container:<id>` to make
|
|
# stages join an existing container's netns; this is
|
|
# what lets the E2E run inside a nested-container
|
|
# sandbox where `--network host` would land in the
|
|
# daemon's namespace instead of the orchestrator's.
|
|
set -euo pipefail
|
|
|
|
IMAGE="${PP_IMAGE:-swactor-pp-gpu:latest}"
|
|
PREFIX="${PP_CONTAINER_PREFIX:-pp-stage}"
|
|
DEV="${PP_DEV:-CPU}"
|
|
CACHE_DIR="${PP_CACHE_DIR:-$HOME/.cache/tinygrad}"
|
|
NAME="${PREFIX}-${STAGE}"
|
|
|
|
# Optional GPU passthrough for real-mode (tinygrad CUDA) runs. The slim
|
|
# runtime image has no host compiler, so the CPU backend can't JIT — real
|
|
# inference needs a GPU. Set PP_GPUS (e.g. "all" or "device=0") to attach
|
|
# one; left unset the container runs CPU-only, which is right for the stub
|
|
# image.
|
|
GPU_ARGS=()
|
|
if [ -n "${PP_GPUS:-}" ]; then
|
|
GPU_ARGS+=(--gpus "$PP_GPUS")
|
|
fi
|
|
# PP_MODEL_CACHE_DIR, when set, is bind-mounted so a stage reuses its
|
|
# downloaded GGUF shard across runs. The mount MUST be per-stage: the
|
|
# worker keys its cache paths (`<file>.gguf`, `.partial`, `.pp_meta`) by
|
|
# the source URL alone, not by stage, so several stages sharing one cache
|
|
# dir would collide on the same .partial — one stage's os.replace/cleanup
|
|
# races another's meta-write (FileNotFoundError on the .partial). On
|
|
# vast.ai each stage owns its machine so this never bites; locally we give
|
|
# each stage its own subdir. Each stage only fetches its own tensor byte
|
|
# ranges, so the per-stage dirs stay small.
|
|
if [ -n "${PP_MODEL_CACHE_DIR:-}" ]; then
|
|
STAGE_MODEL_CACHE="$PP_MODEL_CACHE_DIR/stage-$STAGE"
|
|
mkdir -p "$STAGE_MODEL_CACHE"
|
|
GPU_ARGS+=(-v "$STAGE_MODEL_CACHE":/root/.cache/pp-pipeline -e PP_MODEL_CACHE_DIR=/root/.cache/pp-pipeline)
|
|
fi
|
|
|
|
# Idempotent cleanup of any stale container with the same name so a
|
|
# re-run after a crash never collides with a leftover.
|
|
docker rm -f "$NAME" >/dev/null 2>&1 || true
|
|
|
|
# The tinygrad JIT/kernel cache is a sqlite db; several stages writing one
|
|
# shared copy concurrently can hit "database is locked". Give each stage
|
|
# its own when a GPU is attached (real mode); stub/CPU runs keep the shared
|
|
# default since they never touch it.
|
|
if [ -n "${PP_GPUS:-}" ]; then
|
|
CACHE_DIR="$CACHE_DIR/stage-$STAGE"
|
|
fi
|
|
|
|
# Ensure the cache dir exists so the volume mount doesn't create a
|
|
# root-owned dir under $HOME.
|
|
mkdir -p "$CACHE_DIR"
|
|
|
|
# Forward every env var the orchestrator sets. `-e VAR` (no `=value`)
|
|
# tells `docker run` to copy the value from this shim's environment.
|
|
# Variables that are unset on the host are simply omitted by docker.
|
|
#
|
|
# --entrypoint runs the binary directly, bypassing the image's default
|
|
# pp_entrypoint.sh (sshd + postmortem hold). That supervisor is for remote
|
|
# vast.ai nodes; a local docker stage should exit cleanly when pp-worker
|
|
# does so --rm reaps it and the E2E's no-leftover-container check holds.
|
|
exec docker run --rm --init \
|
|
--entrypoint /usr/local/bin/pp-worker \
|
|
--name "$NAME" \
|
|
--network "${PP_STAGE_NETWORK:-host}" \
|
|
"${GPU_ARGS[@]}" \
|
|
-e STAGE \
|
|
-e NUM_STAGES \
|
|
-e SEED_ADDR \
|
|
-e SEED_DIRECT \
|
|
-e MAX_TOKENS \
|
|
-e MODEL \
|
|
-e PP_WORKER_STUB \
|
|
-e WORKER_CMD \
|
|
-e PEER_NODE_ID \
|
|
-e PEER_DIRECT \
|
|
-e FIRST_PEER_NODE_ID \
|
|
-e FIRST_PEER_DIRECT \
|
|
-e PP_BOOT_DELAY_STAGE \
|
|
-e PP_BOOT_DELAY_SECS \
|
|
-e PP_FLEET_SINK \
|
|
-e DEV="$DEV" \
|
|
-e WORKER_SCRIPT=/usr/local/share/pp_tinygrad_worker.py \
|
|
-v "$CACHE_DIR":/root/.cache/tinygrad \
|
|
"$IMAGE"
|