diff --git a/aorlta/jobs/aorlta-run.sbatch b/aorlta/jobs/aorlta-run.sbatch index 8cd9c1659..5935fd696 100644 --- a/aorlta/jobs/aorlta-run.sbatch +++ b/aorlta/jobs/aorlta-run.sbatch @@ -2,8 +2,21 @@ # AORLTA: one run of the GRPO loop (`aorlta.orchestrator.run`) on one node, such as configs/standin.yaml's 3-iteration # stand-in run. The TokenSpeed image is built in-job if the node lacks it, then the trainer image from # docker/Dockerfile.trainer; an EXIT/TERM trap removes this job's containers, its /dev/shm slots and every image it -# built or pulled. docker runs in the background with `wait`, so a TERM (scancel) runs the trap at once: bash defers -# traps while a foreground command runs, and a killed docker client leaves its container running. +# built or pulled that no other job on the node may be using. docker runs in the background with `wait`, so a TERM +# (scancel) runs the trap at once: bash defers traps while a foreground command runs, and a killed docker client +# leaves its container running. +# +# Two runs on one node. The job asks for the whole node (the --exclusive directive below); a command-line option +# overrides it, so `sbatch --exclusive=user -w ...` lets your own jobs share a node, on a partition that allows +# it (one with OverSubscribe=EXCLUSIVE never puts two jobs on a node). Their configs must name +# disjoint GPUs (rollout.engines, trainer.devices, reward.remote's GPUs, and any GPUs a task's own checks use), and +# each its own run id and reward.remote.socket_dir. The job claims the GPUs and its two /dev/shm directories +# (/ and the socket directory) in files under /dev/shm/aorlta-claims, and ends with exit 89 when a +# job Slurm still runs holds one, or a directory inside or around one (a GPU a task's own checks use is not claimed: +# the job cannot see it). Engine ports +# need nothing with rollout.port_base null (the default: each engine start picks a free window); two runs that both +# fix a port_base must give windows that do not overlap. A container or claim of a job Slurm no longer runs is what +# a cancelled or killed job left: the job removes or takes it over, and never touches another live job's. # # With `reward.remote` in the config (the standin-gpu profile's ConSan and ASan legs), the checkers run in # rocm/dev-ubuntu-24.04:7.0.2 (glibc 2.39 for the ConSan hook, the ROCm 7.0.2 ASan runtime): first aorta's ASan @@ -57,6 +70,8 @@ # AORLTA_EXCLUDE_NODES default none. Nodes this job refuses to run on (exit 90), as a Slurm host list. #SBATCH # lines cannot read the environment, so give Slurm the same list when submitting: # sbatch --exclude="$AORLTA_EXCLUDE_NODES" jobs/aorlta-run.sbatch +# AORLTA_CLAIMS_DIR default /dev/shm/aorlta-claims. Where jobs on this node claim GPUs and directories: node- +# local, and the same for every job of the node, or two jobs cannot see each other's claims # A container sees each directory at its host path. Mounted besides the directories above: the package root, /dev/shm # (reward.remote.socket_dir must be under it, and sync.dir is by default), and from the config run.root, # model.base_dir and model.adapter_dir (the adapter read-only; for a Hugging Face snapshot, the model's whole cache @@ -114,7 +129,118 @@ BUILT_TRAINER=0 PULLED_CHK=0 SOCK_DIR= VRAM_PID= +CLAIMS=${AORLTA_CLAIMS_DIR:-/dev/shm/aorlta-claims} +CLAIMED=() +OWN_DIRS=0 # this job holds the claims on / and the socket directory, so its trap may remove them stamp() { echo "=== [$(date -u +%FT%TZ)] $*"; } +job_ended() { # job_ended ID: Slurm no longer runs job ID: it does not know the job, or names a final state. Any other + # answer, and no answer at all (squeue missing, the controller unreachable), counts as running, so a live job's + # containers, images and claims are never taken for left-overs. + local out + if ! out=$(squeue -h -j "$1" -o %T 2>&1); then + grep -qi 'invalid job id' <<< "$out" + return + fi + case "$out" in + ''|BOOT_FAIL|CANCELLED|COMPLETED|DEADLINE|FAILED|NODE_FAIL|OUT_OF_MEMORY|PREEMPTED|REVOKED|SPECIAL_EXIT|TIMEOUT) + return 0 ;; + esac + return 1 +} +stale_containers() { # the aorlta-run-[-part] containers on this node whose job Slurm no longer runs + local name id + local -A ended=() + while read -r name; do + id=${name#aorlta-run-} + id=${id%%-*} + [ "$id" = "$SLURM_JOB_ID" ] && continue + if [ -z "${ended[$id]:-}" ]; then + if job_ended "$id"; then ended[$id]=1; else ended[$id]=0; fi + fi + [ "${ended[$id]}" = 1 ] && echo "$name" + done < <(docker ps -a --format '{{.Names}}' | grep -xE 'aorlta-run-[0-9]+(-[a-z]+)?') +} +other_jobs() { # the other Slurm jobs running on this node, one per line; "unknown" if Slurm cannot say + local out + out=$(squeue -h -w "${SLURMD_NODENAME:-$(hostname -s)}" -o %i 2>/dev/null) || { echo unknown; return; } + grep -vxF -- "$SLURM_JOB_ID" <<< "$out" | grep . +} +remove_image() { # remove_image IMAGE: remove an image this job built or pulled, unless another job may use it: a + # container of another job was made from it (or from an image built on it), or another job runs on this node (it + # may have found the image when it started and not have run it yet). A kept image is recorded under + # $CLAIMS/images, and every job's cleanup tries it again (retry_kept_images), so the last job to leave removes it. + local users jobs rec + rec=$CLAIMS/images/$(printf %s "$1" | sha256sum | cut -c1-16) + users=$(docker ps -a --filter "ancestor=$1" --format '{{.Names}}' 2>/dev/null | grep -vxE "$CNAME(-[a-z]+)?") + jobs=$(other_jobs) + if [ -n "$users$jobs" ]; then + echo "kept $1: another job may use it (containers: ${users:-none}; jobs on this node: ${jobs:-none})" | tr '\n' ' ' + echo + mkdir -p "$CLAIMS/images" 2>/dev/null && chmod 1777 "$CLAIMS/images" 2>/dev/null + printf '%s\n' "$1" > "$rec" 2>/dev/null + return + fi + if docker rmi -f "$1" >/dev/null 2>&1; then + echo "removed $1" + rm -f "$rec" + elif ! docker image inspect "$1" >/dev/null 2>&1; then + rm -f "$rec" # gone already + fi +} +retry_kept_images() { # the images an earlier job built or pulled and kept for another job (remove_image) + local f img + for f in "$CLAIMS"/images/*; do + [ -f "$f" ] || continue + img=$(head -n 1 "$f") + if [ -z "$img" ] || ! docker image inspect "$img" >/dev/null 2>&1; then rm -f "$f"; continue; fi + remove_image "$img" + done +} +claim() { # claim WHAT KEY...: hold each KEY for this job (a file under $CLAIMS naming its job id). A key held by a job + # Slurm still runs ends this job, exit 89; one whose job ended is taken over. A directory (WHAT "directory PATH", + # PATH canonical: canon_dir) also ends the job when a live job holds one inside it or one it is inside, since + # either job's cleanup removes its directory recursively. Under flock, so two jobs taking over the same left-over + # claim cannot both get it. + local what=$1 key f owner fd held mine other + shift + mkdir -p "$CLAIMS" && chmod 1777 "$CLAIMS" 2>/dev/null # other users' jobs on a shared node claim here too + touch "$CLAIMS/.lock" 2>/dev/null && chmod 666 "$CLAIMS/.lock" 2>/dev/null + exec {fd}>>"$CLAIMS/.lock" && flock -w 60 "$fd" || { echo "[guard] cannot lock $CLAIMS/.lock" >&2; exit 89; } + for key in "$@"; do + f=$CLAIMS/$key + if [ "${what#directory /}" != "$what" ]; then + mine=${what#directory } + for held in "$CLAIMS"/dir-*; do + [ -f "$held" ] && [ "$held" != "$f" ] || continue + owner=$(cut -d' ' -f1 "$held" 2>/dev/null) + [ -n "$owner" ] && [ "$owner" != "$SLURM_JOB_ID" ] || continue + other=$(cut -d' ' -f3- "$held" 2>/dev/null) + case "$mine/" in "$other/"*) ;; *) case "$other/" in "$mine/"*) ;; *) continue ;; esac ;; esac + if ! job_ended "$owner"; then + echo "[guard] job $owner, which Slurm still runs (or Slurm cannot say), holds directory $other on this node," \ + "which overlaps $mine: give this run another run id and socket directory" >&2 + exec {fd}>&- + exit 89 + fi + done + fi + owner=$(cut -d' ' -f1 "$f" 2>/dev/null) + if [ -n "$owner" ] && [ "$owner" != "$SLURM_JOB_ID" ]; then + if ! job_ended "$owner"; then + echo "[guard] job $owner, which Slurm still runs (or Slurm cannot say), holds $what on this node ($f):" \ + "give this run other GPUs, or another run id and socket directory" >&2 + exec {fd}>&- + exit 89 + fi + stamp "taking over $key from job $owner, which has ended" + fi + echo "$SLURM_JOB_ID $what" > "$f" && CLAIMED+=("$f") || { echo "[guard] cannot write $f" >&2; exec {fd}>&-; exit 89; } + done + exec {fd}>&- +} +canon_dir() { realpath -m -- "$1"; } # one spelling per directory: symlinks resolved, no trailing slash, no . or .. +dir_key() { echo "dir-$(canon_dir "$1" | tr -d '\n' | sha256sum | cut -c1-16)"; } +claim_dir() { local d; d=$(canon_dir "$1"); claim "directory $d" "$(dir_key "$d")"; } vram_sampler() { # every 10 s, from sysfs (no GPU context): UTC seconds, PCI address, VRAM used and total in bytes local c while :; do @@ -217,7 +343,11 @@ clean_img() { # an image this node has, for removing root-owned files local i for i in "$BASE" "$CHK_IMG"; do docker image inspect "$i" >/dev/null 2>&1 && { echo "$i"; return; }; done } -cleanup() { +CLEANED=0 +cleanup() { # once: the TERM/INT trap's exit runs the EXIT trap too, and by then this job's claims are released, so a + # second pass could remove the directories of a job that claimed them in between + [ "$CLEANED" = 1 ] && return + CLEANED=1 [ -n "$VRAM_PID" ] && kill "$VRAM_PID" 2>/dev/null docker ps -a --format '{{.Names}}' | grep -xE "$CNAME(-[a-z]+)?" | xargs -r docker rm -f >/dev/null 2>&1 find "$OUT" -maxdepth 3 -name 'gpucore.*' -type f -delete 2>/dev/null @@ -226,25 +356,52 @@ cleanup() { add_mounts mounts rw /dev/shm "$OUT" "$RUN_ROOT" 2>/dev/null # The GPU legs' logs, harnesses and the candidates' parsed sources (stand-in data) are kept, small, before the # node directory goes: they are how an unsettled ConSan or ASan leg is read afterwards. + # (a job that did not get its claims leaves the directories alone: another live job holds them) [ -n "$img" ] && docker run --rm --network none --ulimit core=0 "${mounts[@]}" --entrypoint bash "$img" -c \ - "if cd '$SYNC_DIR/$RUN_ID/scratch/reward' 2>/dev/null; then + "if [ $OWN_DIRS = 1 ] && cd '$SYNC_DIR/$RUN_ID/scratch/reward' 2>/dev/null; then find pkg \( -path '*/san.*' \( -name '*.log' -o -name aorlta_harness.hip -o -name '*.json' \) \ -o -path '*/gate.parse-*' -name source.hip \) -size -2M -print0 | tar --null -czf '$OUT/gpu-leg-evidence.tgz' -T -; cd /; fi - rm -rf '$SYNC_DIR/$RUN_ID' ${SOCK_DIR:+'$SOCK_DIR'}; chown -R $(id -u):$(id -g) '$OUT'; + [ $OWN_DIRS = 1 ] && rm -rf '$SYNC_DIR/$RUN_ID' ${SOCK_DIR:+'$SOCK_DIR'}; chown -R $(id -u):$(id -g) '$OUT'; [ -d '$RUN_DIR' ] && chown -R $(id -u):$(id -g) '$RUN_DIR'" >/dev/null 2>&1 - [ "$BUILT_TRAINER" = 1 ] && docker rmi -f "$IMG" >/dev/null 2>&1 && echo "removed $IMG" - [ "$BUILT_BASE" = 1 ] && docker rmi -f "$BASE" >/dev/null 2>&1 && echo "removed $BASE" - [ "$PULLED_CHK" = 1 ] && docker rmi -f "$CHK_IMG" >/dev/null 2>&1 && echo "removed $CHK_IMG (pulled by this job)" + [ "$BUILT_TRAINER" = 1 ] && remove_image "$IMG" + [ "$BUILT_BASE" = 1 ] && remove_image "$BASE" + [ "$PULLED_CHK" = 1 ] && remove_image "$CHK_IMG" + retry_kept_images + local f + for f in "${CLAIMED[@]}"; do + [ "$(cut -d' ' -f1 "$f" 2>/dev/null)" = "$SLURM_JOB_ID" ] && rm -f "$f" + done echo "containers left: $(docker ps -a --format '{{.Names}}' | grep -cxE "$CNAME(-[a-z]+)?")" } trap cleanup EXIT trap 'cleanup; exit 143' TERM INT stamp "node $(hostname -s) job $SLURM_JOB_ID run $RUN_ID config $CONFIG components $COMPONENTS resume $RESUME commit $COMMIT" -# The node is exclusive: a container of this run left by an earlier, cancelled job is removed before starting. -stale=$(docker ps -a --format '{{.Names}}' | grep -xE 'aorlta-run-[0-9]+(-[a-z]+)?') -[ -n "$stale" ] && { stamp "removing stale containers"; echo "$stale" | xargs -r docker rm -f >/dev/null 2>&1; } +# Before anything uses them (the checker canary empties the socket directory): this run's /dev/shm directories, and +# the GPUs it uses on this node, which are the engines', the trainer ranks' own (with trainer.devices null they share +# the engines') and the checker servers'. The fakes use no GPU. +claim_dir "$SYNC_DIR/$RUN_ID" +[ -n "$SOCK_DIR" ] && claim_dir "$SOCK_DIR" +OWN_DIRS=1 +RUN_GPUS=() +if [ "$COMPONENTS" = real ]; then + read -ra RUN_GPUS < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c " +from aorlta.config import load_config +c = load_config('$CODE/$CONFIG') +gpus = set(c.rollout.engines) +if c.orchestrator.mode != 'score': + gpus |= set(c.trainer.devices or ()) +if c.reward.remote is not None: + for name in c.reward.remote.adapters: + gpus |= set(c.reward.remote.gpus(name)) +print(*sorted(gpus))") || { echo "[guard] $CONFIG does not load" >&2; exit 94; } +fi +for g in "${RUN_GPUS[@]}"; do claim "GPU $g" "gpu$g"; done +stamp "claimed GPUs ${RUN_GPUS[*]:-none} and the run's /dev/shm directories" +# A container left by a job Slurm no longer runs (cancelled, killed) is removed; another live job's are left alone. +stale=$(stale_containers) +[ -n "$stale" ] && { stamp "removing stale containers: $(echo $stale)"; echo "$stale" | xargs -r docker rm -f >/dev/null 2>&1; } df -h /dev/shm | tail -1 # --- the remote checkers, first: a node whose canary or pre-flight fails ends the job before anything is built --- diff --git a/aorlta/src/aorlta/config.py b/aorlta/src/aorlta/config.py index 8c85980f4..c21fc9e5e 100644 --- a/aorlta/src/aorlta/config.py +++ b/aorlta/src/aorlta/config.py @@ -33,6 +33,7 @@ from .atomic import atomic_write_json, is_partial from .ids import IdError, sha256_file, sha256_json, validate_run_id +from .rollout.ports import PORT_STRIDE Source = Literal["standin"] # the stand-in data source; a new source adds a value here RUN_PLAN_SCHEMA = "aorlta.run_plan/1" @@ -439,9 +440,19 @@ class RolloutSection(_Section): # with near-unique answers and few prompts may raise it: one such group in 8 is 12.5%. The engines' start-up # self-test checks the sampler on its own. identical_groups_stop: float = 0.05 + # the engines' host ports. null: each engine start picks a free window of 2000 ports at random, outside the + # kernel's ephemeral range, and tries another if one is taken before the engine binds it (rollout/ports.py), so + # two runs can share a node. An int in [1024, 65536 - 2000 x engines]: engine i on port_base + 2000 i, started + # once; 20000 is the layout every run used before this key. The ports used are in /engines.jsonl. + port_base: int | None = None def _validate(self) -> None: _check(len(self.engines) >= 1, "engines: at least one engine GPU is needed") + if self.port_base is not None: + top = 65536 - PORT_STRIDE * len(self.engines) + _check(1024 <= self.port_base <= top, + f"port_base: must be null or in [1024, {top}] ({len(self.engines)} engines, {PORT_STRIDE} ports " + f"each), got {self.port_base}") _check(all(g >= 0 for g in self.engines), f"engines: GPU indices must be >= 0, got {list(self.engines)}") _check(len(set(self.engines)) == len(self.engines), f"engines: a GPU is listed twice: {list(self.engines)}") _check(self.K >= 2, f"K: must be >= 2 so a group can have an advantage, got {self.K}") diff --git a/aorlta/src/aorlta/orchestrator/wiring.py b/aorlta/src/aorlta/orchestrator/wiring.py index e22c46307..2204257eb 100644 --- a/aorlta/src/aorlta/orchestrator/wiring.py +++ b/aorlta/src/aorlta/orchestrator/wiring.py @@ -48,6 +48,7 @@ PROBE_KERNELS = 40 # WP3's probes are the last `sync.probes` of these stand-in kernels WORKER_THREADS = 8 # CPU threads per spawned trainer rank: its CPU work is small, and the compile pool shares the node ENGINE_CALL_TIMEOUT_S = 600.0 # a probe, a greedy draw or a reload, as EngineProc's own reload default +ENGINES_LOG = "engines.jsonl" # in the run directory: one line per engine start, with its GPU and ports # A plugin task's engine self-test chats (plugins.Task.selftest_messages): open-ended, so two seeds' first 128 tokens # differ on a sound sampler whatever the task's own prompts are. SELF_TEST_CHATS = tuple( @@ -319,7 +320,8 @@ def startup_max_len(cfg: RunConfig, tokenizer, records: Sequence, selftest_messa def engine_handles(cfg: RunConfig, max_len: int) -> list: - """One TokenSpeed engine process per configured GPU, as WP2's harness starts them, on `served_model_dir`.""" + """One TokenSpeed engine process per configured GPU, as WP2's harness starts them, on `served_model_dir`: on + `rollout.port_base + 2000 i`, or with no port, so that each start chooses one (rollout/ports.py).""" from ..rollout.engine_pool import engine_kwargs, engine_port from ..rollout.engine_proc import EngineProc @@ -327,7 +329,8 @@ def engine_handles(cfg: RunConfig, max_len: int) -> list: kw = engine_kwargs(model_dir=str(served_model_dir(cfg)), dtype="bfloat16", P=r.P, K=r.K, n_engines=len(r.engines), max_len=max_len, gpu_memory_utilization=r.gpu_memory_utilization, run_seed=cfg.orchestrator.run_seed) - return [EngineProc(i, str(g), dict(kw, port=engine_port(i))) for i, g in enumerate(r.engines)] + return [EngineProc(i, str(g), kw if r.port_base is None else dict(kw, port=engine_port(i, r.port_base))) + for i, g in enumerate(r.engines)] def rollout_compiler(cfg: RunConfig): @@ -352,9 +355,29 @@ def __init__(self, cfg: RunConfig, tokenizer, handles: Sequence[Any], *, max_len self.clock = clock # the orchestrator's clock: its deadlines are read on it self.handles = [CountedHandle(h) for h in handles] self.pool = EnginePool(self.handles, run_seed=cfg.orchestrator.run_seed, - request_timeout_s=cfg.rollout.request_timeout_s, clock=spans.now) + request_timeout_s=cfg.rollout.request_timeout_s, on_ready=self._record_start, + clock=spans.now) self.manager = None self._on_shutdown = list(on_shutdown) + self._record_lock = threading.Lock() + + def _record_start(self, i: int, ready: dict, restart: bool) -> None: + """One line of /engines.jsonl per engine start, and one in the log: where it runs and the ports it + uses, so that two runs on one node can be told apart afterwards. A run has written its plan, so its directory + exists, by then; a measurement without a run directory (measure/t10.py) gets only the log line.""" + import datetime as dt + + row = {"t": dt.datetime.now(dt.UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z"), + "engine": i, "gpu": getattr(self.handles[i], "gpu", None), "restart": restart, + "port_base": self.cfg.rollout.port_base, "port": ready.get("port"), "ports": ready.get("ports"), + "attempt": ready.get("attempt"), "earlier_attempts": ready.get("earlier_attempts") or [], + "pid": ready.get("pid"), "startup_sec": ready.get("startup_sec")} + log.info("engine %d on GPU %s started%s on port %s (attempt %s): %s", i, row["gpu"], + " again" if restart else "", row["port"], row["attempt"], row["ports"]) + if not self.cfg.run_dir.is_dir(): + return + with self._record_lock, open(self.cfg.run_dir / ENGINES_LOG, "a", encoding="utf-8") as f: + f.write(json.dumps(row, sort_keys=True) + "\n") def start(self, selftest_prompts: Sequence) -> dict: from ..rollout.chain import RolloutManager, RolloutSettings, make_leak_check @@ -378,7 +401,8 @@ def start(self, selftest_prompts: Sequence) -> dict: self.manager = RolloutManager(settings, self.pool, self.tok, self.compiler, leak_check=make_leak_check(r.leak_check_adapter), spans=self.spans, compile_workers=min(96, os.cpu_count() or 8)) - return {"engines": len(self.handles), "max_model_len": self.max_len, "self_tests": self.pool.self_tests} + return {"engines": len(self.handles), "max_model_len": self.max_len, "self_tests": self.pool.self_tests, + "ports": [{"port": (rd or {}).get("port"), "ports": (rd or {}).get("ports")} for rd in self.pool.ready]} def run(self, prompts, iteration: int, policy_version: str, *, on_candidate, on_chain=None, deadline=None): """Each chain goes to `on_chain` as it ends (its candidate with it), so groups train during generation. diff --git a/aorlta/src/aorlta/rollout/engine_pool.py b/aorlta/src/aorlta/rollout/engine_pool.py index 2cdb530aa..3e8a999ba 100644 --- a/aorlta/src/aorlta/rollout/engine_pool.py +++ b/aorlta/src/aorlta/rollout/engine_pool.py @@ -11,9 +11,9 @@ - `ideal_cached_tokens` is the longest prefix of a request this engine has already seen since its last flush: earlier inputs (from their submit) and earlier input + output (from their reply). - A dead or timed-out engine fails its in-flight requests (`engine_error`), is restarted with the same - kwargs, passes start-up and self-test again, and serves again only after the `on_restart` hook - (the reload and landing check of Section 6.5, owned by WP3/WP7) returns True. A second death in one - iteration stops the run. + kwargs (with `rollout.port_base` null, on a newly chosen window of ports), passes start-up and + self-test again, and serves again only after the `on_restart` hook (the reload and landing check of + Section 6.5, owned by WP3/WP7) returns True. A second death in one iteration stops the run. """ from __future__ import annotations @@ -26,13 +26,13 @@ from typing import Any from ..ids import seed +from .ports import PORT_BASE, PORT_STRIDE, engine_port # noqa: F401 - the fixed port layout, imported from here too SELF_TEST_PROMPTS = 4 SELF_TEST_TOKENS = 128 SELF_TEST_TEMPERATURE = 0.7 SELF_TEST_TOP_P = 0.95 SELF_TEST_MIN_DIFFER = 3 -PORT_BASE, PORT_STRIDE = 20000, 2000 class RolloutStop(RuntimeError): @@ -186,7 +186,7 @@ def engine_kwargs( *, model_dir: str, dtype: str, P: int, K: int, n_engines: int, max_len: int, gpu_memory_utilization: float, run_seed: int, log_level: str = "warning", ) -> dict: - """The Section 6.1 engine kwargs, without `port` (set per engine).""" + """The Section 6.1 engine kwargs, without `port` (set per engine, or chosen at each start: rollout/ports.py).""" return { "model": model_dir, "tokenizer": model_dir, "dtype": dtype, "sampling_backend": "triton", "enable_output_logprobs": True, "enable_prefix_caching": True, "disable_kvstore": True, @@ -195,10 +195,6 @@ def engine_kwargs( } -def engine_port(index: int) -> int: - return PORT_BASE + PORT_STRIDE * index - - def lpt_assign(estimates: Sequence[float], n_engines: int) -> list[int]: """Engine index per group: largest estimate first, each to the least-loaded engine that still has room (at most ceil(G/E) groups each). Ties go to the lower engine index and then to the earlier @@ -233,7 +229,9 @@ def ok(self) -> bool: class EnginePool: - """The engines one rollout manager drives. `handles` are EngineProc or InProcessEngine objects.""" + """The engines one rollout manager drives. `handles` are EngineProc or InProcessEngine objects. + `on_ready(i, ready, restart)` gets each engine's ready message as soon as it has started, before its checks, at + start-up and after each restart, so the run can record every engine's ports.""" def __init__( self, @@ -245,6 +243,7 @@ def __init__( self_test_timeout_s: float = 300.0, on_restart: Callable[[int], bool] | None = None, make_handle: Callable[[int], Any] | None = None, + on_ready: Callable[[int, dict, bool], None] | None = None, clock: Callable[[], float] = time.time, ): if not handles: @@ -256,6 +255,7 @@ def __init__( self.self_test_timeout_s = self_test_timeout_s self.on_restart = on_restart self.make_handle = make_handle + self.on_ready = on_ready self.clock = clock self.prefix = [PrefixIndex() for _ in self.handles] self.available = [False] * len(self.handles) @@ -277,6 +277,10 @@ def _bring_up(self, i: int, selftest_prompts: Sequence[Sequence[int]]) -> None: handle = self.handles[i] ready = handle.start() self.ready[i] = ready + if self.on_ready is not None: + with self._lock: + restart = i in self._restarting + self.on_ready(i, ready, restart) problems = server_arg_problems(ready.get("info", {}), self.expected) if problems: handle.stop() diff --git a/aorlta/src/aorlta/rollout/engine_proc.py b/aorlta/src/aorlta/rollout/engine_proc.py index 2cb2cc99e..b053e40b9 100644 --- a/aorlta/src/aorlta/rollout/engine_proc.py +++ b/aorlta/src/aorlta/rollout/engine_proc.py @@ -9,7 +9,10 @@ - `get_server_info()` and `flush_cache()` are never called (A9): the resolved ServerArgs are read from `engine.tokenizer_manager.server_args`; - the same process body can host the fake engine (`factory="fake"`), so the pipe protocol, the death - path and the timeout path are tested on CPU. + path and the timeout path are tested on CPU; +- an engine whose kwargs name no `port` picks a free window of ports at each start, and tries another + if one is taken before the engine binds it (rollout/ports.py); the ready message names the ports the + engine resolved. Messages (driver -> worker), each a tuple whose first item is the kind: ("gen", rid, input_ids, sampling_params, extra) -> ("gen", rid, slim result | {"error": str}) @@ -30,6 +33,8 @@ from concurrent.futures import TimeoutError as FutureTimeout from typing import Any +from . import ports as port_choice + READY_KEYS = ( "sampling_backend", "enable_output_logprobs", "enable_prefix_caching", "dtype", "disable_kvstore", "disable_autotune", "max_num_seqs", "max_model_len", "gpu_memory_utilization", "port", "seed", "numerics", @@ -111,7 +116,8 @@ def send(msg: tuple) -> None: t0 = time.monotonic() engine = _make_engine(factory, engine_index, engine_kwargs, fake) send(("ready", engine_index, {"startup_sec": round(time.monotonic() - t0, 1), "pid": os.getpid(), - "info": _server_info(engine)})) + "info": _server_info(engine), + "ports": port_choice.bound_ports(getattr(engine, "port_args", None))})) except BaseException: send(("ready", engine_index, {"error": traceback.format_exc()[-4000:]})) return @@ -165,7 +171,8 @@ def __init__(self) -> None: class EngineProc: - """Driver-side handle on one engine process. `generate` is thread-safe and blocking.""" + """Driver-side handle on one engine process. `generate` is thread-safe and blocking. With a `port` in + `engine_kwargs` the engine starts on it; without one, each start chooses one (`start`).""" def __init__( self, @@ -181,6 +188,8 @@ def __init__( self.engine_kwargs = dict(engine_kwargs) self.factory, self.fake = factory, fake self.ready_timeout_s = ready_timeout_s + self.fixed_port = "port" in self.engine_kwargs + self.port: int | None = self.engine_kwargs.get("port") # the port of the last successful start self.ready: dict | None = None self._proc: Any = None self._conn: Any = None @@ -196,8 +205,43 @@ def alive(self) -> bool: return not self._dead def start(self) -> dict: + """Start the engine process and wait until it is ready. With a fixed port it is started once. Otherwise each + start holds a free window of ports (`ports.hold_window`), and a start that fails on a port taken after that + check (`ports.lost_to_another`) is tried again on another window, at most `ports.START_ATTEMPTS` times in + all; any other failure is raised at once. The ready message gains `port` (the base used), `attempt` and + `earlier_attempts` (why the starts before it failed).""" + if self.fixed_port: + ready = self._start_on(self.engine_kwargs) + ready.update(port=self.port, attempt=1, earlier_attempts=[]) + return ready + holder, failed, used = id(self), [], [] + for attempt in range(1, port_choice.START_ATTEMPTS + 1): + try: + base = port_choice.hold_window(holder, avoid=used) + except port_choice.PortsUnavailable as exc: + raise EngineStartError(f"engine {self.index}: {exc}" + "".join(f"\n{f}" for f in failed)) from None + try: + ready = self._start_on(dict(self.engine_kwargs, port=base)) + except EngineStartError as exc: + port_choice.release_window(holder) + if not port_choice.lost_to_another(str(exc), base): + if failed: + raise EngineStartError("\n".join([str(exc), *failed])) from None + raise + used.append(base) + failed.append(f"engine {self.index} attempt {attempt} on port {base}, a port was taken: " + f"{str(exc)[-400:]}") + continue + self.port = base + ready.update(port=base, attempt=attempt, earlier_attempts=failed) + return ready + raise EngineStartError(f"engine {self.index}: a port was taken on each of {port_choice.START_ATTEMPTS} " + "starts\n" + "\n".join(failed)) + + def _start_on(self, engine_kwargs: dict) -> dict: import multiprocessing as mp + port = engine_kwargs.get("port") ctx = mp.get_context("spawn") parent, child = ctx.Pipe() with _ENV_LOCK: @@ -205,7 +249,7 @@ def start(self) -> dict: os.environ["HIP_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu try: proc = ctx.Process( - target=worker_main, args=(child, self.index, self.gpu, self.engine_kwargs, self.factory, self.fake), + target=worker_main, args=(child, self.index, self.gpu, engine_kwargs, self.factory, self.fake), daemon=False, ) proc.start() @@ -219,15 +263,15 @@ def start(self) -> dict: self._proc, self._conn = proc, parent if not parent.poll(self.ready_timeout_s): self._terminate() - raise EngineStartError(f"engine {self.index} not ready after {self.ready_timeout_s:g} s") + raise EngineStartError(f"engine {self.index} (port {port}) not ready after {self.ready_timeout_s:g} s") try: _, _, ready = parent.recv() except (EOFError, OSError) as exc: self._terminate() - raise EngineStartError(f"engine {self.index} died during start-up ({exc})") from None + raise EngineStartError(f"engine {self.index} (port {port}) died during start-up ({exc})") from None if "error" in ready: self._terminate() - raise EngineStartError(f"engine {self.index} failed to start:\n{ready['error']}") + raise EngineStartError(f"engine {self.index} (port {port}) failed to start:\n{ready['error']}") self.ready = ready self._dead, self.death = False, None threading.Thread(target=self._read, name=f"engine-{self.index}-reader", daemon=True).start() @@ -308,6 +352,8 @@ def stop(self, timeout_s: float = 60.0) -> None: proc.join(timeout=timeout_s) self._terminate() self._fail_all(f"engine {self.index} stopped") + if not self.fixed_port: + port_choice.release_window(id(self)) class InProcessEngine: diff --git a/aorlta/src/aorlta/rollout/fake_engine.py b/aorlta/src/aorlta/rollout/fake_engine.py index d89a0b4e6..b5ffe9fa1 100644 --- a/aorlta/src/aorlta/rollout/fake_engine.py +++ b/aorlta/src/aorlta/rollout/fake_engine.py @@ -20,6 +20,7 @@ from __future__ import annotations import math +import socket import threading import time import types @@ -28,6 +29,7 @@ from typing import Any from ..ids import INT32_MAX, HashRng, sha256_ids +from .ports import CLUSTER_OFFSETS STOP_TOKEN_ID = 248046 # <|im_end|>, the token the engine stops on (F2, Appendix C) BYTE_TOKEN_BASE = 10_000 # ByteTokenizer: token id = BYTE_TOKEN_BASE + byte value @@ -131,6 +133,10 @@ class FakeEngineConfig: # since the last flush(), at most len(input) - 1, as measured on TokenSpeed (F15). It depends on # request order, so it is off by default and the responses themselves never depend on it. prefix_cache: bool = False + # Bind the ports a one-node TokenSpeed engine binds for its `port` (rollout/ports.py) on 127.0.0.1 until + # shutdown(), so a taken port fails the start with "Address already in use", as on the engine. Off: the ports + # are only reported (`port_args`). + bind_ports: bool = False def __post_init__(self) -> None: if not (isinstance(self.seed, int) and not isinstance(self.seed, bool)): @@ -207,6 +213,22 @@ def __init__( self.engine_id = engine_id self.server_args = FakeServerArgs(**engine_kwargs) self.tokenizer_manager = types.SimpleNamespace(server_args=self.server_args) + d = self.server_args.port + CLUSTER_OFFSETS[0] + self.port_args = types.SimpleNamespace( # as TokenSpeed's PortArgs when none of the ports is taken + dist_init_addr=f"127.0.0.1:{d}", tokenizer_ipc_name=f"tcp://127.0.0.1:{d + 1}", + rpc_ipc_name=f"tcp://127.0.0.1:{d + 3}", metrics_ipc_name=f"tcp://127.0.0.1:{d + 4}", + scheduler_input_ipc_name=f"tcp://127.0.0.1:{d + 5}", nccl_port=self.server_args.port + 100) + self._sockets: list[socket.socket] = [] + if self.fake.bind_ports: + try: + for p in (d, d + 1, d + 3, d + 4, d + 5): + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self._sockets.append(s) + s.bind(("127.0.0.1", p)) + s.listen(1) + except OSError: + self._close_sockets() + raise self.tokenizer = ByteTokenizer() self.requests: list[FakeRequest] = [] self._sleep = sleep @@ -236,6 +258,12 @@ def set_weights(self, tag: str) -> None: def shutdown(self) -> None: self._alive = False + self._close_sockets() + + def _close_sockets(self) -> None: + for s in self._sockets: + s.close() + self._sockets = [] def flush_prefix_cache(self) -> None: with self._lock: diff --git a/aorlta/src/aorlta/rollout/ports.py b/aorlta/src/aorlta/rollout/ports.py new file mode 100644 index 000000000..08a51250e --- /dev/null +++ b/aorlta/src/aorlta/rollout/ports.py @@ -0,0 +1,169 @@ +"""The host ports a rollout engine binds, and how they are chosen (`rollout.port_base`). + +A TokenSpeed `Engine` at 3b684074 binds nothing at its `port` kwarg itself: on one node, +`PortArgs.init_new` (runtime/utils/server_args.py) derives what it binds from it, all on the host's +network (the run's containers use the host network): + +- the rendezvous port d = port + 233, which the scheduler's torch.distributed TCP store listens on, and + ZMQ sockets at d+1 (tokenizer), d+3 (rpc), d+4 (metrics) and d+5 (scheduler input). It also wants + d+6 free. While any of the six is taken it moves d up by 10 and checks again, with no bound; +- an "nccl" port at a random offset in [100, 1000], moved up by 42 while taken, which it checks but a + one-node engine does not bind (its rendezvous is d). + +So an engine's ports stay inside [port, port + PORT_STRIDE) unless about 170 clusters there are taken. +Each check is a bind test that lets the port go again before the engine binds it, so another process +can take a port in between; the engine then fails to start with "Address already in use" (from ZMQ in +the engine's process) or "Initialization failed" (the TCP store, in the scheduler's). + +With `rollout.port_base` null, each start of an engine picks a window here: one of the fixed slots +`AUTO_LOW + PORT_STRIDE * n`, in a random order, whose cluster passes the same bind test, whose window +lies outside the kernel's ephemeral port range (`ip_local_port_range`, where outgoing connections get +their ports at any time) and which another engine of this process does not hold. With the default +ephemeral range there are 12 slots, so eight engines of a node always find room on an idle host. A start that fails on a taken port tries another window, at +most START_ATTEMPTS times in all (EngineProc.start). With an int, engine i uses port_base + 2000 i, +as every run did before the key existed (20000), and is started once. +""" + +from __future__ import annotations + +import random +import re +import socket +import threading +from collections.abc import Callable, Iterable +from pathlib import Path + +PORT_BASE, PORT_STRIDE = 20000, 2000 # the fixed layout: engine i at PORT_BASE + PORT_STRIDE * i +CLUSTER_OFFSETS = (233, 234, 236, 237, 238, 239) # d = port + 233 and the ports TokenSpeed derives from it +AUTO_LOW, AUTO_HIGH = 10000, 65536 # automatic windows are chosen in [AUTO_LOW, AUTO_HIGH), less the ephemeral range +EPHEMERAL_DEFAULT = (32768, 60999) # Linux's default ip_local_port_range, if the file cannot be read +EPHEMERAL_FILE = Path("/proc/sys/net/ipv4/ip_local_port_range") +CHOOSE_TRIES = 200 # slots bind-tested for one window before giving up +START_ATTEMPTS = 3 # engine starts with automatic ports, the first included, before a taken port stops the run +_IN_USE = re.compile(r"address already in use|EADDRINUSE", re.IGNORECASE) + +_LOCK = threading.Lock() +_HELD: dict[int, int] = {} # holder id -> base of the window it holds, for the engines of this process + + +def engine_port(index: int, base: int = PORT_BASE) -> int: + """Engine `index`'s port in the fixed layout.""" + return base + PORT_STRIDE * index + + +def cluster(port: int) -> list[int]: + """The ports an engine started with `port` binds (and checks) when none of them is taken.""" + return [port + o for o in CLUSTER_OFFSETS] + + +def port_free(port: int) -> bool: + """TokenSpeed's own test (`is_port_available`): the port can be bound on every address and listened on.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("", port)) + s.listen(1) + return True + except (OSError, OverflowError): + return False + + +def taken(ports: Iterable[int], free: Callable[[int], bool] = port_free) -> list[int]: + return [p for p in ports if not free(p)] + + +def ephemeral_range(path: Path = EPHEMERAL_FILE) -> tuple[int, int]: + """The kernel's ephemeral port range, inclusive; Linux's default if it cannot be read.""" + try: + lo, hi = (int(x) for x in path.read_text().split()) + return lo, hi + except (OSError, ValueError): + return EPHEMERAL_DEFAULT + + +def _overlaps(a: int, b: int) -> bool: + return a < b + PORT_STRIDE and b < a + PORT_STRIDE + + +def candidate_bases(ephemeral: tuple[int, int], low: int = AUTO_LOW, high: int = AUTO_HIGH) -> list[tuple[int, int]]: + """The intervals [lo, hi] of bases whose whole window lies in [low, high) and outside `ephemeral`.""" + e_lo, e_hi = ephemeral + spans = [(low, min(high, e_lo) - PORT_STRIDE), (max(low, e_hi + 1), high - PORT_STRIDE)] + return [(lo, hi) for lo, hi in spans if lo <= hi] + + +class PortsUnavailable(RuntimeError): + """No free window of engine ports was found.""" + + +def slot_bases(ephemeral: tuple[int, int], low: int = AUTO_LOW, high: int = AUTO_HIGH) -> list[int]: + """The bases automatic windows are chosen from: `low + PORT_STRIDE * n`, each window wholly inside one span of + `candidate_bases`. Fixed slots cannot overlap one another, so the windows held never fragment what is left (any + base in a span could: seven windows at random bases can leave no room for an eighth on an idle host).""" + out = [] + for lo, hi in candidate_bases(ephemeral, low, high): + first = low + -(-(lo - low) // PORT_STRIDE) * PORT_STRIDE + out.extend(range(first, hi + 1, PORT_STRIDE)) + return out + + +def hold_window(holder: int, *, rng: random.Random | None = None, free: Callable[[int], bool] = port_free, + ephemeral: tuple[int, int] | None = None, tries: int = CHOOSE_TRIES, + avoid: Iterable[int] = ()) -> int: + """A base for `holder` (any id, one window each), from `slot_bases` in a random order: its cluster passes the bind + test, its window is outside the ephemeral range and overlaps no window another holder in this process holds, nor + one in `avoid` (bases a failed start of this holder used). At most `tries` slots are bind-tested. Held until + `release_window(holder)`; a holder's earlier window is released first.""" + rng = rng or random.SystemRandom() + slots = slot_bases(ephemeral if ephemeral is not None else ephemeral_range()) + if not slots: + raise PortsUnavailable("no port window lies outside the kernel's ephemeral port range") + rng.shuffle(slots) + avoid = list(avoid) + tested = 0 + with _LOCK: + _HELD.pop(holder, None) + others = list(_HELD.values()) + for base in slots: + if any(_overlaps(base, b) for b in others + avoid): + continue + if tested == tries: + break + tested += 1 + if taken(cluster(base), free): + continue + _HELD[holder] = base + return base + raise PortsUnavailable(f"no free window of {PORT_STRIDE} engine ports among {len(slots)} slots ({tested} " + f"bind-tested; {len(others)} held by other engines of this process, {len(avoid)} avoided)") + + +def release_window(holder: int) -> None: + with _LOCK: + _HELD.pop(holder, None) + + +def held_windows() -> dict[int, int]: + with _LOCK: + return dict(_HELD) + + +def lost_to_another(error: str, base: int, free: Callable[[int], bool] = port_free) -> bool: + """Whether a failed start with `base` looks like a port taken between the check and the engine's bind: the error + says so, or a port of its cluster is held now that the engine's process is gone. Anything else is not retried.""" + return bool(_IN_USE.search(error)) or bool(taken(cluster(base), free)) + + +def bound_ports(port_args) -> dict | None: + """The ports a started engine resolved (TokenSpeed's `Engine.port_args`), by role; None if it has none.""" + if port_args is None: + return None + out: dict[str, int | None] = {} + for role, attr in (("dist_init", "dist_init_addr"), ("tokenizer", "tokenizer_ipc_name"), ("rpc", "rpc_ipc_name"), + ("metrics", "metrics_ipc_name"), ("scheduler_input", "scheduler_input_ipc_name")): + value = getattr(port_args, attr, None) + m = re.search(r":(\d+)$", value) if isinstance(value, str) else None + out[role] = int(m.group(1)) if m else None + nccl = getattr(port_args, "nccl_port", None) + out["nccl_checked"] = nccl if isinstance(nccl, int) else None + return out diff --git a/aorlta/tests/orchestrator/test_run_job.py b/aorlta/tests/orchestrator/test_run_job.py index 5fae50080..344cc7869 100644 --- a/aorlta/tests/orchestrator/test_run_job.py +++ b/aorlta/tests/orchestrator/test_run_job.py @@ -2,9 +2,12 @@ containers that load a profile or the config get the site's tool roots and site file.""" import re +import shutil import subprocess from pathlib import Path +import pytest + PKG = Path(__file__).resolve().parents[2] JOB = (PKG / "jobs" / "aorlta-run.sbatch").read_text() @@ -90,6 +93,213 @@ def test_an_adapter_inside_the_hugging_face_cache_is_mounted_read_only(tmp_path) assert f"{hf}:{hf}" in out and f"{adapter}:{adapter}:ro" in out +# --- two jobs on one node ------------------------------------------------------------------------------------- + +SQUEUE = r"""#!/bin/bash +# squeue -h -j ID -o %T: the state $JOBS names for ID ("ID:STATE ..."), Slurm's error for a job it does not know, or +# the controller's error when $SQUEUE_DOWN is set; squeue -h -w NODE -o %i: the jobs in $NODE_JOBS +[ -n "${SQUEUE_DOWN:-}" ] && { echo "slurm_load_jobs error: Unable to contact slurm controller" >&2; exit 1; } +case "$*" in + *" -w "*) printf '%s\n' $NODE_JOBS; exit 0 ;; +esac +id=$3 +for j in $JOBS; do + [ "${j%%:*}" = "$id" ] && { echo "${j#*:}"; exit 0; } +done +echo "slurm_load_jobs error: Invalid job id specified" >&2 +exit 1 +""" +DOCKER2 = r"""#!/bin/bash +# docker ps -a [--filter ancestor=IMG] --format ...: $CONTAINERS, or the containers $USERS_OF_IMAGE names for IMG; +# docker rmi / rm: logged; docker run ... -c CMD: CMD run here (the cleanup's file work); docker image inspect: found, +# or with $IMAGES_DIR set, found while a file of the image's name is there (rmi removes it) +echo "$*" >> "$DOCKER_LOG" +case "$1" in + ps) if [ "$3" = --filter ]; then printf '%s\n' $USERS_OF_IMAGE; else printf '%s\n' $CONTAINERS; fi ;; + rm|rmi) shift; [ "$1" = -f ] && shift; printf '%s %s\n' "removed" "$*" >> "$DOCKER_LOG" + [ -n "${IMAGES_DIR:-}" ] && rm -f "$IMAGES_DIR/$(printf %s "$*" | tr '/:' '__')" ;; + image) [ -z "${IMAGES_DIR:-}" ] || [ -e "$IMAGES_DIR/$(printf %s "$3" | tr '/:' '__')" ]; exit ;; + run) while [ "$#" -gt 0 ] && [ "$1" != -c ]; do shift; done; [ "$1" = -c ] && bash -c "$2" ;; +esac +exit 0 +""" + + +def _functions(*names: str) -> str: + return "".join(re.search(rf"^{name}\(\) \{{.*?^\}}\n", JOB, re.M | re.S).group(0) for name in names) + + +def _bash(tmp_path, script: str, **env) -> subprocess.CompletedProcess: + bin_dir = tmp_path / "bin" + bin_dir.mkdir(exist_ok=True) + for name, text in (("squeue", SQUEUE), ("docker", DOCKER2)): + (bin_dir / name).write_text(text) + (bin_dir / name).chmod(0o755) + base = {"PATH": f"{bin_dir}:/usr/bin:/bin", "DOCKER_LOG": str(tmp_path / "docker.log"), "SLURM_JOB_ID": "123", + "SLURMD_NODENAME": "node1", "JOBS": "", "NODE_JOBS": "", "CONTAINERS": "", "USERS_OF_IMAGE": ""} + return subprocess.run([shutil.which("bash") or "/bin/bash", "-c", script], env={**base, **env}, capture_output=True, + text=True, timeout=60) + + +def _removed(tmp_path) -> list[str]: + log = tmp_path / "docker.log" + lines = log.read_text().splitlines() if log.exists() else [] + return sorted(" ".join(x.split(" ", 1)[1] for x in lines if x.startswith("removed ")).split()) + + +@pytest.mark.parametrize("state, ended", [("", True), ("RUNNING", False), ("COMPLETING", False), ("PENDING", False), + ("SUSPENDED", False), ("COMPLETED", True), ("CANCELLED", True), + ("FAILED", True), ("TIMEOUT", True), ("NODE_FAIL", True), + ("SOMETHING_NEW", False)]) +def test_a_job_has_ended_only_when_slurm_says_so(tmp_path, state, ended): + out = _bash(tmp_path, f"{_functions('job_ended')}job_ended 77 && echo ended || echo running", + JOBS=f"77:{state}" if state else "77:") + assert out.stdout.strip() == ("ended" if ended else "running"), out.stderr + + +def test_a_job_slurm_does_not_know_has_ended_and_no_answer_from_slurm_is_running(tmp_path): + script = f"{_functions('job_ended')}job_ended 77 && echo ended || echo running" + assert _bash(tmp_path, script).stdout.strip() == "ended" + assert _bash(tmp_path, script, SQUEUE_DOWN="1").stdout.strip() == "running" + tools = tmp_path / "tools" # no squeue at all, the host's included: a PATH with grep alone + tools.mkdir() + (tools / "grep").symlink_to(shutil.which("grep")) + assert _bash(tmp_path, script, PATH=str(tools)).stdout.strip() == "running" + + +def test_only_the_containers_of_jobs_slurm_no_longer_runs_are_stale(tmp_path): + names = ("aorlta-run-100 aorlta-run-100-checkers aorlta-run-200 aorlta-run-200-canary aorlta-run-300 " + "aorlta-run-400-mounts aorlta-run-123 aorlta-run-123-imports aorlta-run-12x other-1 aorlta-run-1000") + out = _bash(tmp_path, f"{_functions('job_ended', 'stale_containers')}stale_containers", CONTAINERS=names, + JOBS="100:RUNNING 300:COMPLETED 400:COMPLETING 1000:RUNNING") + # 200 is unknown to Slurm (ended long ago), 300 completed; 100, 400 and 1000 run, 123 is this job + assert sorted(out.stdout.split()) == ["aorlta-run-200", "aorlta-run-200-canary", "aorlta-run-300"], out.stderr + down = _bash(tmp_path, f"{_functions('job_ended', 'stale_containers')}stale_containers", CONTAINERS=names, + SQUEUE_DOWN="1") + assert down.stdout.split() == [] # no answer from Slurm: nothing is taken for a left-over + calls = (tmp_path / "bin").exists() + assert calls + + +def test_the_job_removes_stale_containers_after_its_claims_and_before_its_first_container(): + assert JOB.index('claim_dir "$SYNC_DIR/$RUN_ID"') < JOB.index("stale=$(stale_containers)") < JOB.index( + '--name "$CNAME-canary"') + assert JOB.index("trap cleanup EXIT") < JOB.index('claim_dir "$SYNC_DIR/$RUN_ID"') + assert "aorlta-run-[0-9]+(-[a-z]+)?')\n" not in JOB[JOB.index("stale=$(stale_containers)"):] # no blanket removal + + +@pytest.mark.parametrize("users, node_jobs, down, removed", [ + ("", "", False, True), + ("aorlta-run-123 aorlta-run-123-checkers", "", False, True), # this job's own containers do not count + ("aorlta-run-456", "", False, False), # another job's container was made from it + ("", "456", False, False), # another job runs on the node and may be about to use it + ("", "", True, False), # Slurm cannot say who else runs here +]) +def test_an_image_this_job_built_is_removed_only_when_no_other_job_may_use_it(tmp_path, users, node_jobs, down, + removed): + env = {"USERS_OF_IMAGE": users, "NODE_JOBS": f"123 {node_jobs}".strip()} + if down: + env["SQUEUE_DOWN"] = "1" + out = _bash(tmp_path, f"{_functions('other_jobs', 'remove_image')}CNAME=aorlta-run-123\nremove_image img:1", **env) + assert _removed(tmp_path) == (["img:1"] if removed else []), out.stdout + out.stderr + assert ("removed img:1" if removed else "kept img:1: another job may use it") in out.stdout + + +def _claim_script(tmp_path, body: str) -> str: + return (f"{_functions('stamp', 'job_ended', 'claim', 'canon_dir', 'dir_key', 'claim_dir')}CLAIMS={tmp_path}/claims\n" + f"CLAIMED=()\n{body}\n" + "echo \"claimed ${#CLAIMED[@]}\"") + + +def test_a_claim_is_taken_kept_by_a_running_job_and_taken_over_from_an_ended_one(tmp_path): + claims = tmp_path / "claims" + out = _bash(tmp_path, _claim_script(tmp_path, 'claim "GPU 0" gpu0\nclaim "GPU 0" gpu0')) + assert out.returncode == 0 and "claimed 2" in out.stdout and (claims / "gpu0").read_text() == "123 GPU 0\n" + (claims / "gpu1").write_text("456 GPU 1\n") + held = _bash(tmp_path, _claim_script(tmp_path, 'claim "GPU 1" gpu1'), SLURM_JOB_ID="789", JOBS="456:RUNNING") + assert held.returncode == 89 and "job 456, which Slurm still runs" in held.stderr and "GPU 1" in held.stderr + assert (claims / "gpu1").read_text() == "456 GPU 1\n" + down = _bash(tmp_path, _claim_script(tmp_path, 'claim "GPU 1" gpu1'), SLURM_JOB_ID="789", SQUEUE_DOWN="1") + assert down.returncode == 89 and (claims / "gpu1").read_text() == "456 GPU 1\n" + over = _bash(tmp_path, _claim_script(tmp_path, 'claim "GPU 1" gpu1'), SLURM_JOB_ID="789", JOBS="456:CANCELLED") + assert over.returncode == 0 and "taking over gpu1 from job 456" in over.stdout + assert (claims / "gpu1").read_text() == "789 GPU 1\n" + + +def test_a_directory_claim_key_is_one_per_path(tmp_path): + real = tmp_path / "shm" / "a" + real.mkdir(parents=True) + (tmp_path / "alias").symlink_to(tmp_path / "shm") + spellings = [real, f"{real}/", f"{tmp_path}/shm/./a", f"{tmp_path}/shm/b/../a", tmp_path / "alias" / "a"] + out = _bash(tmp_path, f"{_functions('canon_dir', 'dir_key')}" + "".join(f"dir_key '{p}'; " for p in spellings) + + f"dir_key '{tmp_path}/shm/b'; dir_key '{tmp_path}/shm/not-yet/c/'; dir_key '{tmp_path}/shm/not-yet/c'") + keys = out.stdout.split() + assert len(set(keys[:5])) == 1 and keys[5] != keys[0] and keys[6] == keys[7], out.stderr + assert keys[0].startswith("dir-") and len(keys[0]) == 20 + + +@pytest.mark.parametrize("mine, clash", [("/x/shm/checkers", True), ("/x/shm/checkers/", True), + ("/x/shm/checkers/run-b", True), ("/x/shm", True), + ("/x/shm/checkers2", False), ("/x/shm/check", False)]) +def test_a_directory_inside_or_around_a_live_jobs_directory_is_refused(tmp_path, mine, clash): + claims = tmp_path / "claims" + claims.mkdir() + (claims / "dir-0123456789abcdef").write_text("456 directory /x/shm/checkers\n") + out = _bash(tmp_path, _claim_script(tmp_path, f'claim_dir "{mine}"'), SLURM_JOB_ID="789", JOBS="456:RUNNING") + assert out.returncode == (89 if clash else 0), out.stdout + out.stderr + if clash: + assert "job 456" in out.stderr and "overlaps" in out.stderr + ended = _bash(tmp_path, _claim_script(tmp_path, f'claim_dir "{mine}"'), SLURM_JOB_ID="789", JOBS="456:COMPLETED") + assert ended.returncode == 0, ended.stderr + own = _bash(tmp_path, _claim_script(tmp_path, f'claim_dir "{mine}"'), SLURM_JOB_ID="456", JOBS="456:RUNNING") + assert own.returncode == 0, own.stderr # a job's own claims never stop it + + +def _cleanup(tmp_path, own: int, built: int = 1, calls: str = "cleanup", **env) -> subprocess.CompletedProcess: + funcs = _functions("add_mounts", "clean_img", "other_jobs", "remove_image", "retry_kept_images") + funcs += re.search(r"^CLEANED=0\ncleanup\(\) \{.*?^\}\n", JOB, re.M | re.S).group(0) + state = (f"CNAME=aorlta-run-{env.get('SLURM_JOB_ID', '123')} VRAM_PID= OUT={tmp_path}/out RUN_ROOT={tmp_path}/runs " + f"SYNC_DIR={tmp_path}/shm RUN_ID=r SOCK_DIR={tmp_path}/shm/sockets RUN_DIR={tmp_path}/runs/r " + f"BUILT_TRAINER={built} BUILT_BASE=0 PULLED_CHK=0 BASE=b CHK_IMG=c IMG=trainer:1 OWN_DIRS={own} " + f"CLAIMS={tmp_path}/claims\nCLAIMED=({tmp_path}/claims/gpu0 {tmp_path}/claims/gpu1)\n") + return _bash(tmp_path, f"{funcs}{state}{calls}", **env) + + +def test_cleanup_removes_the_runs_directories_only_when_this_job_claimed_them(tmp_path): + for d in ("out", "runs", "shm/r/scratch", "shm/sockets", "claims"): + (tmp_path / d).mkdir(parents=True, exist_ok=True) + (tmp_path / "claims" / "gpu0").write_text("123 GPU 0\n") + (tmp_path / "claims" / "gpu1").write_text("456 GPU 1\n") # taken over by another job since: not ours to drop + out = _cleanup(tmp_path, 0, NODE_JOBS="123 456") + assert (tmp_path / "shm" / "r").is_dir() and (tmp_path / "shm" / "sockets").is_dir(), out.stderr + assert not (tmp_path / "claims" / "gpu0").exists() and (tmp_path / "claims" / "gpu1").exists() + assert "kept trainer:1" in out.stdout and _removed(tmp_path) == [] # job 456 still runs on the node + out = _cleanup(tmp_path, 1, NODE_JOBS="123") + assert not (tmp_path / "shm" / "r").exists() and not (tmp_path / "shm" / "sockets").exists(), out.stderr + assert _removed(tmp_path) == ["trainer:1"] and "removed trainer:1" in out.stdout + + +def test_the_job_claims_every_gpu_its_config_uses(tmp_path): + start = 'read -ra RUN_GPUS < <(PYTHONPATH="$CODE/src$XPATH" "$PY" -c "' + snippet = JOB[JOB.index(start) + len(start):JOB.index('") || { echo "[guard] $CONFIG does not load"', + JOB.index("RUN_GPUS"))] + import sys + + def gpus(config: Path, **edits) -> list[int]: + text = snippet.replace("$CODE/$CONFIG", str(config)) + prelude = "import sys; sys.path.insert(0, %r)\n" % str(PKG / "src") + out = subprocess.run([sys.executable, "-c", prelude + text], capture_output=True, text=True, timeout=120, + cwd=PKG) + assert out.returncode == 0, out.stderr + return [int(g) for g in out.stdout.split()] + + assert gpus(PKG / "configs" / "standin.yaml") == list(range(8)) # 4 engines and their ranks, checkers on 4-7 + assert gpus(PKG / "configs" / "arith-score.yaml") == [0] + devices = tmp_path / "devices.yaml" + devices.write_text((PKG / "configs" / "arith-train.yaml").read_text().replace(" ranks: 1", " ranks: 1\n devices: [3]")) + assert gpus(devices) == [0, 3] + + def _aorta_modules_snippet() -> str: """The Python the job runs on the host to pick the aorta modules of the import pre-flight.""" block = JOB[JOB.index("AORTA_MODULES=$("):] @@ -143,3 +353,36 @@ def test_the_import_pre_flight_loads_aorta_only_for_a_profile_that_uses_it(tmp_p report = json.loads(subprocess.run(check + extra, capture_output=True, text=True, env=env).stdout) assert report["aorta"] == extra assert sorted(k for k in report["failed"] if k.startswith("aorta.")) == aorta + + +def test_cleanup_runs_once_so_a_second_pass_never_removes_directories_claimed_since(tmp_path): + for d in ("out", "runs", "shm/r", "shm/sockets", "claims"): + (tmp_path / d).mkdir(parents=True, exist_ok=True) + # the TERM trap's cleanup, a job taking the released directories, then the EXIT trap's cleanup + calls = (f"cleanup\nmkdir -p {tmp_path}/shm/r {tmp_path}/shm/sockets\n" + f"echo '456 directory {tmp_path}/shm/r' > {tmp_path}/claims/dir-r\ncleanup") + out = _cleanup(tmp_path, 1, calls=calls, NODE_JOBS="123") + assert (tmp_path / "shm" / "r").is_dir() and (tmp_path / "shm" / "sockets").is_dir(), out.stdout + out.stderr + assert out.stdout.count("containers left:") == 1 + assert JOB.index("CLEANED=0\ncleanup() {") < JOB.index("trap cleanup EXIT") + assert "trap 'cleanup; exit 143' TERM INT" in JOB + + +def test_an_image_kept_for_another_job_is_removed_by_the_last_job_to_leave(tmp_path): + images = tmp_path / "images" + images.mkdir() + for name in ("trainer_1", "older_1"): + (images / name).write_text("") # older:1 was on the node before either job: no job recorded it + for d in ("out", "runs", "claims"): + (tmp_path / d).mkdir(exist_ok=True) + first = _cleanup(tmp_path, 0, NODE_JOBS="123 456", USERS_OF_IMAGE="aorlta-run-456", IMAGES_DIR=str(images)) + assert "kept trainer:1" in first.stdout and _removed(tmp_path) == [], first.stdout + first.stderr + assert [p.read_text() for p in (tmp_path / "claims" / "images").iterdir()] == ["trainer:1\n"] + # job 456 found the image and did not build it; it leaves last + last = _cleanup(tmp_path, 0, built=0, SLURM_JOB_ID="456", NODE_JOBS="456", IMAGES_DIR=str(images)) + assert _removed(tmp_path) == ["trainer:1"] and "removed trainer:1" in last.stdout, last.stdout + last.stderr + assert list((tmp_path / "claims" / "images").iterdir()) == [] and (images / "older_1").exists() + # a record whose image is gone already is dropped, and nothing is removed for it + (tmp_path / "claims" / "images" / "x").write_text("gone:1\n") + _cleanup(tmp_path, 0, built=0, SLURM_JOB_ID="789", NODE_JOBS="789", IMAGES_DIR=str(images)) + assert list((tmp_path / "claims" / "images").iterdir()) == [] and _removed(tmp_path) == ["trainer:1"] diff --git a/aorlta/tests/orchestrator/test_score.py b/aorlta/tests/orchestrator/test_score.py index 08a68c752..41b0c1c99 100644 --- a/aorlta/tests/orchestrator/test_score.py +++ b/aorlta/tests/orchestrator/test_score.py @@ -3,6 +3,7 @@ import json import os +import re import subprocess import sys from pathlib import Path @@ -433,7 +434,7 @@ def test_the_run_job_prints_a_score_runs_summary_and_says_it_is_never_resumed(): post = text[text.index("Every committed iteration's summary"):] branch = post[post.index('if cfg.orchestrator.mode == "score":'):post.index("for n in complete_iterations")] assert "REPORT_DIR, SUMMARY_NAME" in branch and "sys.exit(0)" in branch - header = text[:text.index("#SBATCH")] + header = text[:re.search(r"^#SBATCH", text, re.M).start()] # up to the first directive, not a mention of one assert "orchestrator.mode: score" in header and "never resumed" in header @@ -494,7 +495,8 @@ def orchestrator_mounts(tmp_path, data: dict) -> list[str]: env = {k: v for k, v in os.environ.items() if not k.startswith(("AORLTA_", "SLURM"))} env.update(PATH=f"{bin_dir}:{env['PATH']}", SLURM_JOB_ID="4242", SLURM_SUBMIT_DIR=str(pkg), CONFIG="configs/score.yaml", COMPONENTS="fake", AORLTA_VENV=str(venv), - AORLTA_TMPDIR=str(tmp_path / "tmpdir"), AORLTA_HF_HOME=str(tmp_path / "hf")) + AORLTA_TMPDIR=str(tmp_path / "tmpdir"), AORLTA_HF_HOME=str(tmp_path / "hf"), + AORLTA_CLAIMS_DIR=str(tmp_path / "claims")) proc = subprocess.run(["bash", str(REPO / "jobs" / "aorlta-run.sbatch")], cwd=pkg, env=env, capture_output=True, text=True, timeout=300) assert proc.returncode == 0, proc.stdout[-2000:] + proc.stderr[-2000:] diff --git a/aorlta/tests/orchestrator/test_wiring.py b/aorlta/tests/orchestrator/test_wiring.py index 94cc0d96a..fe504dca9 100644 --- a/aorlta/tests/orchestrator/test_wiring.py +++ b/aorlta/tests/orchestrator/test_wiring.py @@ -436,6 +436,37 @@ def now(self): assert r.manager.kw["deadline"] is None +def test_the_engines_choose_their_ports_unless_the_config_fixes_a_base(tmp_path): + auto = wiring.engine_handles(_cfg(tmp_path), 4096) + assert [h.gpu for h in auto] == ["0", "1"] and all("port" not in h.engine_kwargs and not h.fixed_port for h in auto) + fixed = wiring.engine_handles(_cfg(tmp_path, rollout={"port_base": 20000}), 4096) + assert [h.engine_kwargs["port"] for h in fixed] == [20000, 22000] and all(h.fixed_port for h in fixed) + assert {k: v for k, v in fixed[1].engine_kwargs.items() if k != "port"} == auto[1].engine_kwargs + + +def test_each_engine_start_is_recorded_with_its_gpu_and_ports(tmp_path, caplog): + from aorlta.rollout.fake_engine import FakeEngine, FakeEngineConfig + + cfg = _cfg(tmp_path) + handles = [InProcessEngine(i, FakeEngine(FakeEngineConfig(), port=20000 + 2000 * i)) for i in range(2)] + r = wiring.RolloutAdapter(cfg, None, handles, max_len=4096, spans=types.SimpleNamespace(now=lambda: 0.0), + compiler=None) + ready = {"port": 24000, "attempt": 2, "earlier_attempts": ["attempt 1 on port 12000, a port was taken"], "pid": 5, + "startup_sec": 1.5, "ports": {"dist_init": 24233}} + r._record_start(1, ready, False) # no run directory yet: the log line only + assert not (cfg.run_dir / wiring.ENGINES_LOG).exists() + cfg.run_dir.mkdir(parents=True) + with caplog.at_level("INFO", logger=wiring.log.name): + r._record_start(1, ready, False) + r._record_start(1, dict(ready, port=26000, attempt=1, earlier_attempts=[]), True) + rows = [json.loads(x) for x in (cfg.run_dir / wiring.ENGINES_LOG).read_text().splitlines()] + assert [(x["engine"], x["gpu"], x["restart"], x["port"], x["attempt"]) for x in rows] == [ + (1, "inproc", False, 24000, 2), (1, "inproc", True, 26000, 1)] + assert rows[0]["ports"] == {"dist_init": 24233} and rows[0]["earlier_attempts"] == ready["earlier_attempts"] + assert rows[0]["port_base"] is None and rows[0]["t"].endswith("Z") and rows[0]["pid"] == 5 + assert "engine 1 on GPU inproc started again on port 26000" in caplog.text + + def test_an_adapter_on_several_gpus_gets_one_factory_over_its_servers(tmp_path): from aorlta.reward.clock import Clock from aorlta.reward.remote import RemoteFactory diff --git a/aorlta/tests/rollout/test_engine_pool.py b/aorlta/tests/rollout/test_engine_pool.py index 19e2e23d4..c1a2871a6 100644 --- a/aorlta/tests/rollout/test_engine_pool.py +++ b/aorlta/tests/rollout/test_engine_pool.py @@ -184,6 +184,21 @@ def test_a_dead_engine_restarts_and_serves_only_after_the_landing_check(rt, land assert ("error" not in after.result) == landing_ok +def test_every_engine_start_is_reported_at_start_up_and_after_a_restart(rt): + h = Mortal(0, FakeEngine(CFG, **rt.ENGINE_KW)) + seen = [] + pool = rt.make_pool([h, FakeEngine(CFG, **rt.ENGINE_KW)], on_restart=lambda i: True, + on_ready=lambda i, ready, restart: seen.append((i, restart, ready))) + assert sorted((i, r) for i, r, _ in seen) == [(0, False), (1, False)] + assert all(ready is pool.ready[i] for i, _, ready in seen) + pool.begin_iteration(1) + h.armed = True + pool.generate(0, [1], _sp()) + pool.wait_restarts(30) + assert pool.available == [True, True] and [(i, r) for i, r, _ in seen[2:]] == [(0, True)] + assert seen[2][2] is pool.ready[0] + + def test_a_second_death_in_one_iteration_stops_the_run(rt): h = Mortal(0, FakeEngine(CFG, **rt.ENGINE_KW)) pool = rt.make_pool([h], on_restart=lambda i: True) diff --git a/aorlta/tests/rollout/test_engine_proc.py b/aorlta/tests/rollout/test_engine_proc.py index a6f0d0659..b5747317c 100644 --- a/aorlta/tests/rollout/test_engine_proc.py +++ b/aorlta/tests/rollout/test_engine_proc.py @@ -3,11 +3,13 @@ from __future__ import annotations +import socket import threading import time import pytest +from aorlta.rollout import ports from aorlta.rollout.engine_pool import engine_kwargs, engine_port from aorlta.rollout.engine_proc import EngineProc, EngineStartError, slim from aorlta.rollout.fake_engine import FakeEngineConfig @@ -81,6 +83,121 @@ def test_a_worker_that_cannot_start_is_an_error(): p.start() +def _auto(index=0) -> EngineProc: + """An engine with no port in its kwargs, whose fake binds TokenSpeed's ports as the engine would.""" + cfg = FakeEngineConfig(seed=1, min_tokens=20, max_tokens=40, bind_ports=True) + return EngineProc(index, str(index), dict(KW), factory="fake", fake=cfg, ready_timeout_s=120) + + +def _listen(port: int) -> socket.socket: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", port)) + s.listen(1) + return s + + +def _free_bases(n: int) -> list[int]: + out, base = [], 20000 + while len(out) < n: + if not ports.taken(ports.cluster(base)): + out.append(base) + base += ports.PORT_STRIDE + return out + + +def test_an_engine_without_a_port_starts_on_a_free_window_and_reports_the_ports_it_bound(): + a, b = _auto(0), _auto(1) + try: + ra, rb = a.start(), b.start() + for p, ready in ((a, ra), (b, rb)): + base = ready["port"] + assert ready["info"]["port"] == base == p.port and ready["attempt"] == 1 and ready["earlier_attempts"] == [] + assert ready["ports"] == {"dist_init": base + 233, "tokenizer": base + 234, "rpc": base + 236, + "metrics": base + 237, "scheduler_input": base + 238, "nccl_checked": base + 100} + assert not ports.port_free(base + 233) # bound by the engine's process while it runs + assert ports.held_windows()[id(p)] == base + assert abs(ra["port"] - rb["port"]) >= ports.PORT_STRIDE + finally: + a.stop() + b.stop() + assert id(a) not in ports.held_windows() and ports.port_free(ra["port"] + 233) + + +def test_a_start_that_loses_a_port_after_the_check_tries_another_window(monkeypatch): + lost, _ = _free_bases(2) + real = ports.hold_window + calls = [] + + def first_loses(holder, **kw): # the first window's check passed, then another process took a port + calls.append(list(kw.get("avoid"))) + return lost if len(calls) == 1 else real(holder, **kw) + + monkeypatch.setattr(ports, "hold_window", first_loses) + p = _auto() + with _listen(lost + 234): + try: + ready = p.start() + assert ready["attempt"] == 2 and ready["port"] != lost and p.port == ready["port"] + assert len(ready["earlier_attempts"]) == 1 and f"on port {lost}, a port was taken" in ready[ + "earlier_attempts"][0] and "Address already in use" in ready["earlier_attempts"][0] + assert calls == [[], [lost]] # the second choice avoids the lost window + assert "error" not in p.generate([1, 2], SP, timeout_s=60) + finally: + p.stop() + + +def test_a_port_lost_on_every_attempt_stops_with_every_attempt_named(monkeypatch): + bases = _free_bases(ports.START_ATTEMPTS + 1) + held = [_listen(b + 233) for b in bases] + monkeypatch.setattr(ports, "hold_window", lambda holder, **kw: bases[len(kw["avoid"])]) + try: + with pytest.raises(EngineStartError, match=f"a port was taken on each of {ports.START_ATTEMPTS} starts") as e: + _auto().start() + assert all(f"on port {b}" in str(e.value) for b in bases[:ports.START_ATTEMPTS]) + assert f"on port {bases[ports.START_ATTEMPTS]}" not in str(e.value) + finally: + for s in held: + s.close() + + +def test_a_start_that_fails_for_another_reason_is_not_retried(monkeypatch): + real, calls = ports.hold_window, [] + monkeypatch.setattr(ports, "hold_window", lambda holder, **kw: calls.append(1) or real(holder, **kw)) + p = EngineProc(0, "0", dict(KW, sampling_backend="flashinfer"), factory="fake", + fake=FakeEngineConfig(bind_ports=True), ready_timeout_s=60) + with pytest.raises(EngineStartError, match="failed to start") as e: + p.start() + assert calls == [1] and "a port was taken" not in str(e.value) and id(p) not in ports.held_windows() + + +def test_a_fixed_port_is_used_as_given_and_started_once(monkeypatch): + base = _free_bases(1)[0] + monkeypatch.setattr(ports, "hold_window", lambda *a, **k: pytest.fail("a fixed port is never chosen")) + cfg = FakeEngineConfig(seed=1, min_tokens=20, max_tokens=40, bind_ports=True) + p = EngineProc(0, "0", dict(KW, port=base), factory="fake", fake=cfg, ready_timeout_s=120) + try: + ready = p.start() + assert ready["port"] == base and ready["ports"]["dist_init"] == base + 233 and ready["attempt"] == 1 + finally: + p.stop() + with _listen(base + 233), pytest.raises(EngineStartError, match=f"(?s)port {base}.*Address already in use"): + EngineProc(0, "0", dict(KW, port=base), factory="fake", fake=cfg, ready_timeout_s=120).start() + + +def test_an_engine_restarted_after_a_death_starts_again_with_its_kwargs(): + p = _auto() + try: + first = p.start() + p.kill("test death") + assert not p.alive + again = p.start() # what the pool's restart calls, on the same handle + assert again["attempt"] == 1 and again["ports"]["dist_init"] == again["port"] + 233 + assert p.engine_kwargs == dict(KW) and ports.held_windows()[id(p)] == again["port"] + assert "error" not in p.generate([1, 2], SP, timeout_s=60) and first["pid"] != again["pid"] + finally: + p.stop() + + def test_slim_reads_triples_and_reports_malformed_logprobs(): out = {"output_ids": [5, 6], "meta_info": {"output_token_logprobs": [(-0.5, 5, None), (-1.0, 6, None)], "cached_tokens": 3, "finish_reason": {"type": "length", "length": 2}}} diff --git a/aorlta/tests/rollout/test_ports.py b/aorlta/tests/rollout/test_ports.py new file mode 100644 index 000000000..4349c6545 --- /dev/null +++ b/aorlta/tests/rollout/test_ports.py @@ -0,0 +1,162 @@ +"""The engines' host ports (rollout/ports.py): TokenSpeed's layout, the bind test, and the automatic choice of a free +window, on real sockets of this host where the test needs a port to be taken (CPU only).""" + +from __future__ import annotations + +import socket +import types + +import pytest + +from aorlta.rollout import ports as P + + +class Seq: + """A random source whose shuffle puts the given bases first, in order (as `hold_window` tries them).""" + + def __init__(self, bases): + self.bases = list(bases) + + def shuffle(self, slots): + rest = [b for b in slots if b not in self.bases] + assert all(b in slots for b in self.bases), (self.bases, slots) + slots[:] = self.bases + rest + + +def listen(port: int) -> socket.socket: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.bind(("127.0.0.1", port)) + s.listen(1) + return s + + +def free_bases(n: int, start: int = 20000) -> list[int]: + """`n` bases in the lower automatic span, windows apart, whose clusters are free on this host now.""" + out, base = [], start + while len(out) < n: + assert base < 30000, "no free port windows on this host" + if not P.taken(P.cluster(base)): + out.append(base) + base += P.PORT_STRIDE + return out + + +@pytest.fixture(autouse=True) +def no_held_windows(): + yield + for holder in list(P.held_windows()): + P.release_window(holder) + + +def test_the_fixed_layout_is_the_one_every_run_used_before_the_key(): + assert [P.engine_port(i) for i in range(4)] == [20000, 22000, 24000, 26000] + assert [P.engine_port(i, 31000) for i in range(3)] == [31000, 33000, 35000] + + +def test_the_cluster_is_what_tokenspeed_derives_from_the_port(): + # PortArgs.init_new at 3b684074: d = port + 233; tokenizer d+1, rpc d+3, metrics d+4, scheduler input d+5; d+6 + # checked too + assert P.cluster(20000) == [20233, 20234, 20236, 20237, 20238, 20239] + assert all(o < P.PORT_STRIDE for o in P.CLUSTER_OFFSETS) + + +def test_the_bind_test_sees_a_listening_port(): + base = free_bases(1)[0] + assert P.port_free(base + 233) + with listen(base + 233): + assert not P.port_free(base + 233) + assert P.taken(P.cluster(base)) == [base + 233] + + +def test_automatic_windows_lie_outside_the_ephemeral_range(): + assert P.candidate_bases((32768, 60999)) == [(10000, 30768), (61000, 63536)] + for lo, hi in P.candidate_bases((32768, 60999)): + for base in (lo, hi): + window = range(base, base + P.PORT_STRIDE) + assert window[-1] < 65536 and (window[-1] < 32768 or window[0] > 60999) + assert P.candidate_bases((1024, 65535)) == [] + with pytest.raises(P.PortsUnavailable, match="ephemeral"): + P.hold_window(1, ephemeral=(1024, 65535)) + + +def test_the_ephemeral_range_is_read_from_the_kernel(tmp_path): + f = tmp_path / "range" + f.write_text("40000\t50000\n") + assert P.ephemeral_range(f) == (40000, 50000) + assert P.ephemeral_range(tmp_path / "missing") == P.EPHEMERAL_DEFAULT + + +def test_a_window_with_a_taken_port_is_passed_over(): + busy, ok = free_bases(2) + with listen(busy + 236): # one port of the cluster, not its first + base = P.hold_window(1, rng=Seq([busy, ok]), ephemeral=(32768, 60999)) + assert base == ok and P.held_windows() == {1: ok} + + +def test_engines_of_one_process_never_share_a_window(): + a = P.hold_window(1, rng=Seq([20000]), free=lambda p: True, ephemeral=(32768, 60999)) + b = P.hold_window(2, rng=Seq([20000, 22000]), free=lambda p: True, ephemeral=(32768, 60999)) + assert (a, b) == (20000, 22000) # 20000 is engine 1's + P.release_window(1) + assert P.hold_window(3, rng=Seq([20000]), free=lambda p: True, ephemeral=(32768, 60999)) == 20000 + # a holder choosing again gives up its own window first, and skips the ones it is told to avoid + assert P.hold_window(2, rng=Seq([22000]), free=lambda p: True, ephemeral=(32768, 60999)) == 22000 + again = P.hold_window(2, rng=Seq([22000, 12000]), free=lambda p: True, ephemeral=(32768, 60999), avoid=[22000]) + assert again == 12000 and P.held_windows() == {3: 20000, 2: 12000} + + +def test_automatic_windows_are_fixed_slots_inside_the_spans(): + slots = P.slot_bases((32768, 60999)) + assert slots == list(range(10000, 30001, P.PORT_STRIDE)) + [62000] + assert all(any(lo <= b <= hi for lo, hi in P.candidate_bases((32768, 60999))) for b in slots) + assert P.slot_bases((1024, 65535)) == [] + + +@pytest.mark.parametrize("seed", [105, *range(40)]) +def test_eight_engines_always_find_windows_on_an_idle_host(seed): + # random bases fragmented the spans: with Random(105) the eighth of eight found no room in 200 tries + import random + + rng = random.Random(seed) + bases = [P.hold_window(h, rng=rng, free=lambda p: True, ephemeral=(32768, 60999)) for h in range(8)] + assert len(set(bases)) == 8 and all(b in P.slot_bases((32768, 60999)) for b in bases) + assert all(not P._overlaps(a, b) for i, a in enumerate(bases) for b in bases[i + 1:]) + # all twelve slots can be held at once, and a thirteenth is an error + more = [P.hold_window(h, rng=rng, free=lambda p: True, ephemeral=(32768, 60999)) for h in range(8, 12)] + assert sorted(bases + more) == P.slot_bases((32768, 60999)) + with pytest.raises(P.PortsUnavailable, match="12 slots"): + P.hold_window(12, rng=rng, free=lambda p: True, ephemeral=(32768, 60999)) + + +def test_no_free_window_is_an_error_not_a_guess(): + with pytest.raises(P.PortsUnavailable, match="no free window"): + P.hold_window(1, free=lambda p: False, tries=20) + assert P.held_windows() == {} + + +def test_a_random_choice_is_free_and_in_range(): + for holder in range(8): + base = P.hold_window(holder) + assert not P.taken(P.cluster(base)) + assert any(lo <= base <= hi for lo, hi in P.candidate_bases(P.ephemeral_range())) + bases = sorted(P.held_windows().values()) + assert all(b - a >= P.PORT_STRIDE for a, b in zip(bases, bases[1:])) + + +def test_a_lost_port_is_told_from_another_failure(): + base = free_bases(1)[0] + assert P.lost_to_another("zmq.error.ZMQError: Address already in use (addr='tcp://127.0.0.1:1')", base) + assert P.lost_to_another("RuntimeError: ... code: -98, name: EADDRINUSE, message: address already in use", base) + assert not P.lost_to_another("RuntimeError: Initialization failed. Please see the error messages above.", base) + with listen(base + 233): # the TCP store's port, which only the scheduler's own log names + assert P.lost_to_another("RuntimeError: Initialization failed. Please see the error messages above.", base) + + +def test_the_bound_ports_are_read_from_tokenspeeds_port_args(): + args = types.SimpleNamespace(dist_init_addr="127.0.0.1:20243", tokenizer_ipc_name="tcp://127.0.0.1:20244", + rpc_ipc_name="tcp://127.0.0.1:20246", metrics_ipc_name="tcp://127.0.0.1:20247", + scheduler_input_ipc_name="tcp://127.0.0.1:20248", nccl_port=20517) + assert P.bound_ports(args) == {"dist_init": 20243, "tokenizer": 20244, "rpc": 20246, "metrics": 20247, + "scheduler_input": 20248, "nccl_checked": 20517} + assert P.bound_ports(None) is None + assert P.bound_ports(types.SimpleNamespace(dist_init_addr="ipc:///x"))["dist_init"] is None diff --git a/aorlta/tests/test_config.py b/aorlta/tests/test_config.py index 1ae490e54..91f6eeab2 100644 --- a/aorlta/tests/test_config.py +++ b/aorlta/tests/test_config.py @@ -306,6 +306,20 @@ def test_value_rules(dotted, value, message): assert message in str(err.value) +def test_the_engine_ports_are_chosen_at_start_unless_a_base_is_given(): + assert config_from_dict(standin()).rollout.port_base is None + four = standin() # 4 engines: the last window must end at 65535 + for ok in (1024, 20000, 65536 - 4 * 2000): + assert config_from_dict(edit(four, "rollout.port_base", ok)).rollout.port_base == ok + for bad in (1023, 65536 - 4 * 2000 + 1, -1): + with pytest.raises(ConfigError, match=r"rollout.port_base: must be null or in \[1024, 57536\] \(4 engines"): + config_from_dict(edit(four, "rollout.port_base", bad)) + one = edit(edit(four, "rollout.engines", [0]), "trainer.ranks", 1) + assert config_from_dict(edit(one, "rollout.port_base", 63536)).rollout.port_base == 63536 + with pytest.raises(ConfigError, match="rollout.port_base: expected an int"): + config_from_dict(edit(four, "rollout.port_base", "20000")) + + def test_accepted_alternatives(): cfg = config_from_dict(edit(edit(standin(), "trainer.aggregation", "constant(4096)"), "trainer.tis", 2.0)) assert cfg.trainer.aggregation == "constant(4096)" and cfg.trainer.tis == 2.0