"""Multi-step pipeline runner for the deploy/* subcommands. Each Coach UI deploy action (GitHub push, droplet sync, droplet provision) expands to a list of `Step`s that must run serially. This module owns the serial-execution loop; everything else is just a step list. Design notes: - We deliberately do NOT reuse `runs.spawn_subprocess_streaming` directly. That helper assumes a single Popen per run and emits `StatusEvent('succeeded')` + closes subscribers as soon as the Popen exits. Chaining N steps would prematurely close the SSE stream after step 1. Instead, this module runs its own daemon thread, manually publishes `LogEvent`s line-by-line, and emits exactly one terminal `StatusEvent` after all steps complete (or one fails). - Cancellation routes through the existing `RunRegistry.cancel()` API. We call `registry.attach_process()` each time a new step starts, so the registry's Popen pointer always tracks the live subprocess. - For non-subprocess work (httpx calls to the Dev Cloud API), the provision pipeline does not use `Step` objects — it publishes `LogEvent`s inline and then hands off to the step runner for the SSH phases. """ from __future__ import annotations import contextlib import subprocess import threading from collections.abc import Callable from pathlib import Path from mindxtrain.deploy.amd_dev_cloud import ( AmdDevCloudClient, AmdDevCloudConfig, extract_public_ip, ) from mindxtrain.deploy.cloud_init import BOOTSTRAP_SENTINEL, render from mindxtrain.deploy.droplet import ( DropletConfig, build_scp_plan_back, build_ssh_probe, build_tail_cloud_init, build_tail_training_log, sync_steps, ) from mindxtrain.deploy.github_push import ( GithubConfig, Step, bootstrap_steps, remote_url, write_sha_file, ) from mindxtrain.operator.runs import ( LogEvent, RunRegistry, StatusEvent, default_registry, parse_trainer_log_line, ) # ---- run_pipeline -------------------------------------------------------- def _publish_log(registry: RunRegistry, run_id: str, line: str) -> None: registry.publish_threadsafe(run_id, LogEvent(run_id=run_id, line=line)) def _open_log_file(log_path: Path) -> object: log_path.parent.mkdir(parents=True, exist_ok=True) return log_path.open("w", buffering=1) def _stream_step( *, step: Step, registry: RunRegistry, run_id: str, log_file: object, capture: list[str], parse_trainer: bool = False, ) -> int: """Run a single step's subprocess, tee stdout to log_file + events. By default each line becomes a `LogEvent`. With `parse_trainer=True`, lines that match the HF Trainer JSON format are upgraded to `StepEvent` (drives Coach's loss chart and metrics table); non-matching lines still become `LogEvent`s. This is how remote training output bridges into the same SSE channel the in-process trainer would publish to. Captures stdout into `capture` if `step.capture_stdout`. Returns the rc. """ proc = subprocess.Popen( step.cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=step.env or None, text=True, bufsize=1, ) registry.attach_process(run_id, proc) assert proc.stdout is not None step_ctr = 0 for raw in proc.stdout: line = raw.rstrip("\n") log_file.write(raw) # type: ignore[attr-defined] log_file.flush() # type: ignore[attr-defined] if step.capture_stdout: capture.append(line) if parse_trainer: step_ev = parse_trainer_log_line(line, fallback_step=step_ctr + 1) if step_ev is not None: step_ctr = step_ev.step registry.publish_threadsafe( run_id, step_ev.model_copy(update={"run_id": run_id}), ) continue _publish_log(registry, run_id, line) return proc.wait() def run_pipeline( steps: list[Step], *, run_id: str, out_dir: Path, registry: RunRegistry | None = None, on_done: Callable[[int, dict[str, str]], None] | None = None, ) -> threading.Thread: """Run `steps` serially in a daemon thread. Returns the thread immediately. Predicate logic: if a step's `predicate_step` was run earlier and that step's rc is *not* in `predicate_rc_in`, the step is skipped (with a LogEvent) and its rc is recorded as -1 ("skipped"). Captured stdout: any step with `capture_stdout=True` has its full stdout joined into a string and stashed in `captured[step.label]`. The `on_done` callback receives (final_rc, captured). Final status: emits exactly one `StatusEvent("succeeded"|"failed")` and closes subscribers. Cancel via `registry.cancel(run_id)`. """ reg = registry if registry is not None else default_registry() log_file = _open_log_file(out_dir / "pipeline.log") def _runner() -> None: rcs: dict[str, int] = {} captured: dict[str, str] = {} final_rc = 0 try: for i, step in enumerate(steps, 1): # Predicate: skip if the gating step's rc is not in the allowed set. if step.predicate_step is not None: gate_rc = rcs.get(step.predicate_step) if gate_rc is None or gate_rc not in step.predicate_rc_in: _publish_log(reg, run_id, f"=== skip {i}/{len(steps)}: {step.label} (predicate {step.predicate_step}={gate_rc}) ===") rcs[step.label] = -1 continue _publish_log(reg, run_id, f"=== step {i}/{len(steps)}: {step.label} ===") buffer: list[str] = [] rc = _stream_step( step=step, registry=reg, run_id=run_id, log_file=log_file, capture=buffer, ) rcs[step.label] = rc if step.capture_stdout: captured[step.label] = "\n".join(buffer) if rc != 0 and not step.allow_failure: _publish_log(reg, run_id, f"=== step {step.label} exited rc={rc}; aborting pipeline ===") final_rc = rc break if rc != 0: _publish_log(reg, run_id, f" (probe {step.label} rc={rc} — continuing)") status = "succeeded" if final_rc == 0 else "failed" reg.publish_threadsafe( run_id, StatusEvent(run_id=run_id, status=status, message=f"rc={final_rc}"), ) finally: with contextlib.suppress(Exception): log_file.close() # type: ignore[attr-defined] reg.close_subscribers(run_id) if on_done is not None: with contextlib.suppress(Exception): on_done(final_rc, captured) t = threading.Thread(target=_runner, daemon=True, name=f"deploy-{run_id}") t.start() return t # ---- GitHub push pipeline ------------------------------------------------ def github_push_pipeline( cfg: GithubConfig, *, run_id: str, out_dir: Path, commit_message: str = "mindXtrain initial push", force: bool = False, registry: RunRegistry | None = None, on_done: Callable[[int, dict[str, str]], None] | None = None, ) -> threading.Thread: """Drives the github_push step list. Captures HEAD sha to git_sha.txt.""" reg = registry if registry is not None else default_registry() steps = bootstrap_steps(cfg, commit_message=commit_message, force=force) # Wrap on_done so we can persist the captured HEAD sha + emit a clean # guidance LogEvent if a stale remote was detected without --force. def _on_done(rc: int, captured: dict[str, str]) -> None: sha = captured.get("head-sha", "").strip() existing_remote = captured.get("probe-remote", "").strip() if sha: target = write_sha_file(out_dir, sha) _publish_log(reg, run_id, f" HEAD sha {sha} written to {target}") if existing_remote and existing_remote != remote_url(cfg.repo) and not force: _publish_log( reg, run_id, f" origin already points at {existing_remote!r}; " f"re-run with force=true to switch to {remote_url(cfg.repo)!r}", ) if on_done is not None: with contextlib.suppress(Exception): on_done(rc, captured) return run_pipeline(steps, run_id=run_id, out_dir=out_dir, registry=reg, on_done=_on_done) # ---- Existing-droplet sync pipeline ------------------------------------- def droplet_sync_pipeline( cfg: DropletConfig, *, repo_root: Path, run_id: str, out_dir: Path, run_bench: bool = True, fetch_plan: bool = True, registry: RunRegistry | None = None, on_done: Callable[[int, dict[str, str]], None] | None = None, ) -> threading.Thread: plan_dest = out_dir / "plan.remote.json" steps = sync_steps( cfg, repo_root=repo_root, run_bench=run_bench, fetch_plan=fetch_plan, plan_dest=plan_dest, ) return run_pipeline(steps, run_id=run_id, out_dir=out_dir, registry=registry, on_done=on_done) # ---- AMD Dev Cloud provision pipeline ----------------------------------- def droplet_provision_pipeline( cloud_cfg: AmdDevCloudConfig, *, name: str, repo: str, branch: str, container: str, extras: str, run_id: str, out_dir: Path, wait_for_bootstrap: bool = True, recipe: str | None = None, registry: RunRegistry | None = None, client_factory: Callable[[AmdDevCloudConfig], AmdDevCloudClient] | None = None, on_done: Callable[[int, dict[str, str]], None] | None = None, ) -> threading.Thread: """Create a droplet, wait for bootstrap, scp plan.json back, optionally train + bridge training events into the run's SSE stream. When `recipe` is set, cloud-init also runs `mindxtrain train` after bench, and the orchestrator adds a fifth step that SSH-tails the training log so Coach's loss chart populates live from the droplet. Without `recipe`, behaviour is exactly as before (bench only). """ reg = registry if registry is not None else default_registry() out_dir.mkdir(parents=True, exist_ok=True) user_data = render( repo=repo, branch=branch, container=container, extras=extras, recipe=recipe, ) factory = client_factory or AmdDevCloudClient def _log(line: str) -> None: _publish_log(reg, run_id, line) def _runner() -> None: final_rc = 0 captured: dict[str, str] = {} try: with factory(cloud_cfg) as client: _publish_log(reg, run_id, "=== step 1/4: create droplet ===") droplet = client.create(name=name, user_data=user_data, log=_log) droplet_id = int(droplet.get("id", 0)) captured["droplet_id"] = str(droplet_id) (out_dir / "droplet_id.txt").write_text(str(droplet_id) + "\n") _publish_log(reg, run_id, f"=== step 2/4: poll until active ({droplet_id}) ===") droplet = client.poll_until_active(droplet_id, log=_log) ip = extract_public_ip(droplet) or "" captured["public_ip"] = ip (out_dir / "public_ip.txt").write_text(ip + "\n") if not ip: _publish_log(reg, run_id, " (no public IPv4 returned; bailing out)") final_rc = 2 return if not wait_for_bootstrap: _publish_log(reg, run_id, "wait_for_bootstrap=false — exiting after provision") return droplet_cfg = DropletConfig( host=ip, user="root", ssh_key=cloud_cfg_ssh_key(cloud_cfg), container=container, extras=extras, ) _publish_log(reg, run_id, "=== step 3/4: wait for ssh + tail cloud-init ===") # Inline ssh-probe with a few retries; bench-stage cloud-init can take a while. probe_steps = [ Step(label=f"ssh-probe-{i}", cmd=build_ssh_probe(droplet_cfg), allow_failure=True) for i in range(60) # 60 * 5s = 5 min ] log_file = _open_log_file(out_dir / "pipeline.log") ready = False try: for ps in probe_steps: rc = _stream_step(step=ps, registry=reg, run_id=run_id, log_file=log_file, capture=[]) if rc == 0: ready = True break import time as _time _time.sleep(5) finally: with contextlib.suppress(Exception): log_file.close() # type: ignore[attr-defined] if not ready: _publish_log(reg, run_id, " ssh did not come up in 5 minutes; aborting") final_rc = 3 return tail_step = Step(label="cloud-init-tail", cmd=build_tail_cloud_init(droplet_cfg)) scp_step = Step( label="scp-plan", cmd=build_scp_plan_back(droplet_cfg, out_dir / "plan.remote.json"), allow_failure=True, ) total_steps = 5 if recipe else 4 log_file = _open_log_file(out_dir / "pipeline.log") try: _publish_log(reg, run_id, f" (waiting for {BOOTSTRAP_SENTINEL})") rc = _stream_step(step=tail_step, registry=reg, run_id=run_id, log_file=log_file, capture=[]) if rc != 0: final_rc = rc return _publish_log(reg, run_id, f"=== step 4/{total_steps}: scp plan.json back ===") _stream_step(step=scp_step, registry=reg, run_id=run_id, log_file=log_file, capture=[]) if recipe: _publish_log( reg, run_id, f"=== step 5/{total_steps}: bridge training log " f"for recipe={recipe} ===", ) train_tail = Step( label="train-tail", cmd=build_tail_training_log(droplet_cfg), ) train_rc = _stream_step( step=train_tail, registry=reg, run_id=run_id, log_file=log_file, capture=[], parse_trainer=True, ) if train_rc != 0: _publish_log( reg, run_id, f" (remote train exit rc={train_rc})", ) final_rc = train_rc finally: with contextlib.suppress(Exception): log_file.close() # type: ignore[attr-defined] except Exception as exc: _publish_log(reg, run_id, f"!! {type(exc).__name__}: {exc}") final_rc = 1 finally: status = "succeeded" if final_rc == 0 else "failed" reg.publish_threadsafe( run_id, StatusEvent(run_id=run_id, status=status, message=f"rc={final_rc}"), ) reg.close_subscribers(run_id) if on_done is not None: with contextlib.suppress(Exception): on_done(final_rc, captured) t = threading.Thread(target=_runner, daemon=True, name=f"provision-{run_id}") t.start() return t def cloud_cfg_ssh_key(cloud_cfg: AmdDevCloudConfig) -> str: """Default ssh key path used to SSH into a freshly-provisioned droplet. The Dev Cloud control plane only stores the public key; the matching private key must be on the operator host. We assume it lives at `~/.ssh/id_ed25519` unless DROPLET_SSH_KEY is set. """ import os as _os return _os.environ.get("DROPLET_SSH_KEY", "~/.ssh/id_ed25519") __all__ = [ "droplet_provision_pipeline", "droplet_sync_pipeline", "github_push_pipeline", "run_pipeline", ]