cross-py-builder/ctrl/scheduler.py

import datetime
import os
import threading

import requests

import ctrl.db as db
import ctrl.settings as settings
import ctrl.workers as wrk


class Scheduler:
    def __init__(self):
        self._stop = threading.Event()
        self._thread = None
        self._active = 0
        self._active_lock = threading.Lock()
        self._log_lock = threading.Lock()
        self._stream_subscribers = {}  # job_id -> set of queue.Queue

    # ---- SSE helpers -----------------------------------------------------
    def subscribe(self, job_id):
        import queue
        q = queue.Queue()
        self._stream_subscribers.setdefault(job_id, set()).add(q)
        return q

    def unsubscribe(self, job_id, q):
        subs = self._stream_subscribers.get(job_id)
        if subs:
            subs.discard(q)
            if not subs:
                self._stream_subscribers.pop(job_id, None)

    def _publish(self, job_id, event, data):
        subs = self._stream_subscribers.get(job_id)
        if subs:
            frame = f"event: {event}\ndata: {data}\n\n"
            for q in list(subs):
                q.put_nowait(frame)

    # ---- lifecycle -------------------------------------------------------
    def start(self):
        if self._thread and self._thread.is_alive():
            return
        settings.ensure_dirs()
        db.init_db()
        self._stop.clear()
        self._thread = threading.Thread(target=self._loop, name="scheduler", daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=5)

    # ---- dispatch loop ---------------------------------------------------
    def _loop(self):
        while not self._stop.is_set():
            try:
                self._tick()
            except Exception as e:
                print(f"[scheduler] tick error: {e}")
            self._stop.wait(settings.SCHEDULER_INTERVAL)

    def _tick(self):
        with self._active_lock:
            if self._active >= settings.MAX_CONCURRENT:
                return
            free = settings.MAX_CONCURRENT - self._active
        jobs = [j for j in db.list_jobs(limit=200) if j["status"] == "queued"]
        for job in jobs[:free]:
            self._dispatch(job)

    # ---- job dispatch ----------------------------------------------------
    def _dispatch(self, job):
        job_id = job["id"]
        candidates = [w for w in wrk.probe_all() if wrk.matches(job, w)]
        if not candidates:
            self._append(job_id, f"[scheduler] No worker matched requirements (os='{job['os_req']}', cpu='{job['cpu_req']}').\n")
            return

        worker = candidates[0]
        with self._active_lock:
            self._active += 1
        try:
            self._run_on_worker(job, worker)
        finally:
            with self._active_lock:
                self._active -= 1

    def _run_on_worker(self, job, worker):
        job_id = job["id"]
        db.update_job(job_id, status="dispatching", worker_host=worker["url"],
                      started_at=datetime.datetime.now().isoformat())
        self._append(job_id, f"[scheduler] Dispatching to {worker['url']} ({worker.get('os')} {worker.get('cpu')})\n")

        try:
            if job["source_type"] == "git":
                worker_job = self._request_git(job, worker)
            else:
                worker_job = self._request_upload(job, worker)
        except Exception as e:
            db.update_job(job_id, status="failed", finished_at=datetime.datetime.now().isoformat(),
                          error=str(e))
            self._append(job_id, f"[scheduler] FAILED: {e}\n")
            return

        worker_job_id = worker_job.get("id")
        worker_url = worker["url"]
        db.update_job(job_id, status="building", worker_job_id=worker_job_id, worker_url=worker_url)

        # download artifacts
        try:
            self._fetch_artifacts(job, worker, worker_job_id)
            db.update_job(job_id, status="done", finished_at=datetime.datetime.now().isoformat())
            self._append(job_id, "[scheduler] Build complete, artifacts saved.\n")
        except Exception as e:
            db.update_job(job_id, status="failed", finished_at=datetime.datetime.now().isoformat(),
                          error=str(e))
            self._append(job_id, f"[scheduler] FAILED: {e}\n")
        finally:
            self._publish(job["id"], "job", '"done"')

    def _request_git(self, job, worker):
        self._append(job["id"], f"[scheduler] Cloning {job['source']} on {worker['url']}\n")
        resp = requests.post(f"{worker['url']}/checkout_git",
                             json={"repo_url": job["source"]},
                             timeout=settings.WORKER_BUILD_TIMEOUT,
                             stream=False)
        return self._worker_response(resp, "checkout")

    def _request_upload(self, job, worker):
        path = self._source_path(job)
        if not path or not os.path.exists(path):
            raise FileNotFoundError(f"Source file not found: {path}")
        self._append(job["id"], f"[scheduler] Uploading {os.path.basename(path)} to {worker['url']}\n")
        with open(path, "rb") as f:
            resp = requests.post(f"{worker['url']}/upload",
                                 files={"file": (os.path.basename(path), f, "application/zip")},
                                 timeout=settings.WORKER_BUILD_TIMEOUT,
                                 stream=False)
        worker_job = self._worker_response(resp, "upload")
        if not worker_job.get("id"):
            # Agent may have started building before returning; try progress fallback.
            pass
        return worker_job

    @staticmethod
    def _worker_response(resp, action):
        try:
            data = resp.json()
        except ValueError:
            raise RuntimeError(f"{action} returned non-JSON ({resp.status_code})")
        if resp.status_code >= 400:
            raise RuntimeError(f"{action} failed ({resp.status_code}): {data.get('error', data)}")
        return data

    def _source_path(self, job):
        jobs_dir = os.path.join(settings.DATA_DIR, "sources")
        name = job.get("source")
        if not name:
            return None
        return os.path.join(jobs_dir, name)

    def _fetch_artifacts(self, job, worker, worker_job_id):
        job_id = job["id"]
        self._append(job_id, f"[scheduler] Fetching artifacts from {worker['url']}/download/{worker_job_id}\n")
        resp = requests.get(f"{worker['url']}/download/{worker_job_id}", timeout=settings.WORKER_BUILD_TIMEOUT)
        resp.raise_for_status()

        out_dir = os.path.join(settings.BUILDS_DIR, job_id)
        os.makedirs(out_dir, exist_ok=True)
        base = os.path.splitext(os.path.basename(job.get("source") or job_id))[0]
        fname = f"{base}-{str(worker.get('os')).lower()}-{str(worker.get('cpu')).lower()}.zip"
        dest = os.path.join(out_dir, fname)
        with open(dest, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)

        db.update_job(job_id, artifacts=[fname])

    # ---- logging / status -------------------------------------------------
    def _append(self, job_id, text):
        with self._log_lock:
            db.append_log(job_id, text)
        self._publish(job_id, "log", '"%s"' % text.rstrip().replace('"', '\\"'))


def make_scheduler():
    return Scheduler()