pal-code 0.1.0__py3-none-any.whl

pal/orchestrator/components/context_interrupt.py

@@ -0,0 +1,220 @@
+"""
+Context interrupt mechanism for Round Robin scheduling (AIOS section
+3.3 / detailed in their Appendix A.3, which this module is the
+counterpart of here).
+
+Long-running LLM inference is exactly the kind of system call that can
+monopolize a shared resource -- one slow generation blocks everything
+queued behind it unless something can pause it. FIFO scheduling doesn't
+need that (it runs each request to completion by design); Round Robin
+does, since giving every request a bounded time slice only works if a
+request that isn't done yet can be suspended and picked back up later
+without losing the work already done.
+
+Two snapshot/restore approaches, matching what's actually available
+per core:
+
+- **Text-based** (`_run_slice_text` / `ContextSnapshot(mode="text")`):
+  for cores with no logits access -- every cloud API and every local
+  server this project talks to (Anthropic, Gemini, OpenAI-compatible,
+  Ollama). None of them return raw per-step logits over HTTP, so the
+  only state available to snapshot is the decoded text generated so
+  far. Resumption re-injects that text as a partial assistant turn plus
+  a short "continue exactly where you left off" instruction, and lets
+  the model carry on from there.
+
+- **Logits-based** (`_run_slice_logits` / `ContextSnapshot(mode="logits")`):
+  for cores with direct logits access -- currently only
+  HuggingFaceLocalLLM, the in-process local model. Generation is
+  driven one token at a time via manual beam-search stepping
+  (HuggingFaceLocalLLM.beam_search_step), and the snapshot is the
+  actual search-tree state (every beam's token sequence + cumulative
+  score, see BeamSearchState in huggingface_local.py) rather than
+  just its decoded text. Resuming continues the search from that exact
+  state -- no re-tokenization, no risk of the model repeating or
+  drifting from where it left off the way a text-based "please
+  continue" prompt can.
+
+Both are request-level approximations of "pause an in-flight
+computation" -- neither literally suspends a running matrix
+multiplication. What differs between them is how much of the model's
+own state survives the pause: text-based keeps only what came out the
+other end; logits-based keeps the search state that produced it.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from pal.llm_wrapper.components.base import LLMSyscall, LLMWrapper
+
+_CONTINUE_INSTRUCTION = (
+    "[Continue your previous reply exactly where it left off. Do not "
+    "repeat any earlier part of it, do not restart, and do not "
+    "acknowledge this instruction.]"
+)
+
+
+@dataclass
+class ContextSnapshot:
+    """The suspended state of one interrupted generation.
+
+    mode="text": text_so_far is the only state -- decoding must
+    restart from a continuation prompt built around it.
+    mode="logits": beam_state carries the actual search-tree state
+    (an opaque, core-specific object -- BeamSearchState for
+    HuggingFaceLocalLLM) that resumption steps forward directly,
+    with text_so_far kept alongside purely for human-readable
+    inspection/logging, not used to drive resumption.
+    """
+
+    mode: str  # "text" | "logits"
+    text_so_far: str = ""
+    beam_state: Any = None
+
+
+@dataclass
+class InterruptResult:
+    complete: bool
+    text_so_far: str
+    elapsed: float
+    snapshot: ContextSnapshot
+    # True when text_so_far is the FULL cumulative text generated across
+    # every slice so far (logits mode -- beam decoding naturally tracks
+    # whole sequences). False when it's only this slice's delta, to be
+    # appended to whatever came before (text mode -- each streamed
+    # continuation only yields its own new chunk).
+    cumulative: bool = False
+    chunks_yielded: int = 0
+
+
+class ContextInterruptController:
+    """Runs one LLMSyscall against one core for at most `quantum`
+    seconds, choosing the text-based or logits-based approach based on
+    the core's capabilities(), and returns either a complete result or
+    a ContextSnapshot the scheduler can use to resume later."""
+
+    def run_slice(
+        self,
+        core: LLMWrapper,
+        syscall: LLMSyscall,
+        quantum: float,
+        resume_snapshot: Optional[ContextSnapshot] = None,
+    ) -> InterruptResult:
+        use_logits = (resume_snapshot is not None and resume_snapshot.mode == "logits") or (
+            resume_snapshot is None
+            and core.capabilities().supports_logits_snapshot
+            and hasattr(core, "init_beam_search")
+            and syscall.params.get("use_logits_snapshot", True)
+        )
+        if use_logits:
+            return self._run_slice_logits(core, syscall, quantum, resume_snapshot)
+        return self._run_slice_text(core, syscall, quantum)
+
+    # -- text-based: closed-source / no-logits-access cores --------------------
+
+    def _run_slice_text(self, core: LLMWrapper, syscall: LLMSyscall, quantum: float) -> InterruptResult:
+        text_parts: list[str] = []
+        t0 = time.time()
+        chunks = 0
+
+        for chunk in core.stream_chat(syscall):
+            text_parts.append(chunk)
+            chunks += 1
+            elapsed = time.time() - t0
+            if elapsed >= quantum:
+                # Time slice exhausted mid-generation -- stop pulling
+                # from the stream here. Cores without real incremental
+                # streaming (LLMWrapper.stream_chat's default fallback,
+                # which yields the whole response as one chunk) will
+                # already have finished by the time the first chunk
+                # arrives, so this branch only actually interrupts
+                # something for cores with true token-level streaming.
+                text_so_far = "".join(text_parts)
+                return InterruptResult(
+                    complete=False,
+                    text_so_far=text_so_far,
+                    elapsed=elapsed,
+                    snapshot=ContextSnapshot(mode="text", text_so_far=text_so_far),
+                    cumulative=False,
+                    chunks_yielded=chunks,
+                )
+
+        text_so_far = "".join(text_parts)
+        return InterruptResult(
+            complete=True,
+            text_so_far=text_so_far,
+            elapsed=time.time() - t0,
+            snapshot=ContextSnapshot(mode="text", text_so_far=text_so_far),
+            cumulative=False,
+            chunks_yielded=chunks,
+        )
+
+    def build_continuation(self, syscall: LLMSyscall, partial: InterruptResult) -> LLMSyscall:
+        """Text-based resumption: folds the partial output back into the
+        message list as prior context and returns a new LLMSyscall ready
+        to be re-queued. Only meaningful for mode="text" results -- the
+        scheduler doesn't call this for logits-mode interrupts, since
+        those resume from `beam_state` directly instead of rebuilding
+        the prompt."""
+        continued_messages = list(syscall.messages) + [
+            {"role": "assistant", "content": partial.text_so_far},
+            {"role": "user", "content": _CONTINUE_INSTRUCTION},
+        ]
+        return LLMSyscall(
+            messages=continued_messages,
+            model=syscall.model,
+            temperature=syscall.temperature,
+            max_tokens=syscall.max_tokens,
+            stream=syscall.stream,
+            params=syscall.params,
+        )
+
+    # -- logits-based: direct logits access, search-tree state -----------------
+
+    def _run_slice_logits(
+        self,
+        core: LLMWrapper,
+        syscall: LLMSyscall,
+        quantum: float,
+        resume_snapshot: Optional[ContextSnapshot],
+    ) -> InterruptResult:
+        t0 = time.time()
+        beam_width = syscall.params.get("beam_width", 1)
+        max_new_tokens = syscall.max_tokens or 512
+
+        state = resume_snapshot.beam_state if resume_snapshot is not None else core.init_beam_search(
+            syscall, beam_width=beam_width
+        )
+
+        steps_this_slice = 0
+        while not state.finished:
+            state = core.beam_search_step(state, max_new_tokens=max_new_tokens)
+            steps_this_slice += 1
+            elapsed = time.time() - t0
+
+            if elapsed >= quantum and not state.finished:
+                # Suspended mid-generation -- snapshot the intermediate
+                # search-tree state (not just its decoded text) so
+                # resumption continues the beam search exactly where
+                # this step left off.
+                decoded = core.decode_beam_search(state)
+                return InterruptResult(
+                    complete=False,
+                    text_so_far=decoded,
+                    elapsed=elapsed,
+                    snapshot=ContextSnapshot(mode="logits", text_so_far=decoded, beam_state=state),
+                    cumulative=True,
+                    chunks_yielded=steps_this_slice,
+                )
+
+        final_text = core.decode_beam_search(state)
+        return InterruptResult(
+            complete=True,
+            text_so_far=final_text,
+            elapsed=time.time() - t0,
+            snapshot=ContextSnapshot(mode="logits", text_so_far=final_text, beam_state=state),
+            cumulative=True,
+            chunks_yielded=steps_this_slice,
+        )

pal/orchestrator/components/limits.py

@@ -0,0 +1,9 @@
+from pal.orchestrator.components.cgroup_manager import ResourceLimits
+
+# Preset resource envelopes for common agent workload sizes. cpu.max
+# format is "<quota> <period>" in microseconds -- e.g. "50000 100000" =
+# 50% of one CPU core.
+
+LIGHT = ResourceLimits(memory_max="256M", cpu_max="25000 100000", pids_max="20")
+STANDARD = ResourceLimits(memory_max="512M", cpu_max="50000 100000", pids_max="50")
+HEAVY = ResourceLimits(memory_max="2G", cpu_max="100000 100000", pids_max="100")

pal/orchestrator/components/scheduler.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+
+from pal.orchestrator.components.agent import Agent, AgentConfig
+from pal.orchestrator.components.cgroup_manager import CgroupManager
+
+
+@dataclass
+class ScheduledAgent:
+    config: AgentConfig
+    priority: int = 0   # higher runs first when slots are contended
+    weight: int = 100    # relative cpu.weight if backend supports it
+
+
+class Scheduler:
+    """
+    Bounds how many agents run concurrently (max_concurrent), and
+    decides which queued agents get a slot next based on priority.
+    Doesn't touch cgroup internals directly -- that stays
+    CgroupManager's job. This is purely admission control: who gets to
+    start, and in what order.
+    """
+
+    def __init__(self, cgroup_manager: CgroupManager, max_concurrent: int = 3):
+        self.cgroup_manager = cgroup_manager
+        self.max_concurrent = max_concurrent
+        self.running: dict[str, Agent] = {}
+        self.queue: list[ScheduledAgent] = []
+
+    def submit(self, config: AgentConfig, priority: int = 0, weight: int = 100) -> None:
+        if config.name in self.running:
+            raise ValueError(f"agent {config.name} already running")
+        if any(s.config.name == config.name for s in self.queue):
+            raise ValueError(f"agent {config.name} already queued")
+        self.queue.append(ScheduledAgent(config=config, priority=priority, weight=weight))
+        self.queue.sort(key=lambda s: s.priority, reverse=True)
+
+    def tick(self) -> list[str]:
+        """
+        Admits as many queued agents as available slots allow, in
+        priority order. Returns the names of agents newly started this
+        tick. Call this after submit(), and again after any stop() to
+        backfill freed slots.
+        """
+        started = []
+        while self.queue and len(self.running) < self.max_concurrent:
+            scheduled = self.queue.pop(0)
+            agent = Agent(scheduled.config, self.cgroup_manager)
+            agent.start()
+            self.running[scheduled.config.name] = agent
+            started.append(scheduled.config.name)
+        return started
+
+    def stop(self, agent_name: str) -> None:
+        agent = self.running.pop(agent_name, None)
+        if agent is None:
+            raise ValueError(f"agent {agent_name} not running")
+        agent.stop()
+
+    def stop_all(self) -> None:
+        for name in list(self.running.keys()):
+            self.stop(name)
+
+    def status(self) -> dict:
+        return {
+            "running": list(self.running.keys()),
+            "queued": [s.config.name for s in self.queue],
+            "capacity": self.max_concurrent,
+            "slots_free": self.max_concurrent - len(self.running),
+        }

pal/orchestrator/components/syscall_scheduler.py

@@ -0,0 +1,200 @@
+"""
+Centralized system-call scheduler (AIOS section 3.3).
+
+All LLM inference requests -- from Pipeline.run(), from orchestrator
+Agents, from anywhere else in the system -- get submitted here as
+SyscallRequests instead of each caller hitting an LLMWrapper directly.
+That's the actual point of centralizing: request management
+(queueing, ordering, time-slicing) lives in exactly one place, so
+every other module only has to focus on producing/consuming a
+request/result pair, not on how or when its turn comes.
+
+Two scheduling strategies, same as AIOS:
+- FIFO: process queued requests in arrival order, straight through to
+  completion. Simple, but a long request can make everything behind it
+  wait its full duration.
+- Round Robin: give each queued request a bounded time slice
+  (`quantum`), then move to the next one, cycling back to unfinished
+  requests later. This is what needs the context-interrupt mechanism in
+  context_interrupt.py -- pausing an LLM generation mid-stream and
+  resuming it later without restarting from scratch.
+"""
+from __future__ import annotations
+
+import time
+import uuid
+from collections import deque
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+from pal.llm_wrapper.components.base import LLMResponse, LLMSyscall, LLMWrapper
+from pal.orchestrator.components.context_interrupt import ContextInterruptController, ContextSnapshot
+
+
+class SchedulingStrategy(str, Enum):
+    FIFO = "fifo"
+    ROUND_ROBIN = "round_robin"
+
+
+class SyscallStatus(str, Enum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    DONE = "done"
+    ERROR = "error"
+
+
+@dataclass
+class SyscallRequest:
+    """One queued LLM inference request. `core` is whichever LLMWrapper
+    the submitting module wants this run against -- the scheduler
+    doesn't care which deployment option it is, only when it runs."""
+
+    agent_name: str
+    core: LLMWrapper
+    syscall: LLMSyscall
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    submitted_at: float = field(default_factory=time.time)
+    priority: int = 0  # FIFO ignores this; reserved for future strategies
+
+    # internal scheduling state
+    status: SyscallStatus = SyscallStatus.QUEUED
+    accumulated_text: str = ""
+    slices_run: int = 0
+    result: Optional[LLMResponse] = None
+    error: Optional[str] = None
+    # Set when a Round Robin slice is interrupted, so the next tick
+    # knows to RESUME (text-based: via the continuation baked into
+    # `syscall` already; logits-based: via snapshot.beam_state) rather
+    # than starting a fresh generation.
+    snapshot: Optional[ContextSnapshot] = None
+
+
+class SyscallScheduler:
+    """
+    One centralized queue for every LLM syscall in the system,
+    regardless of which module submitted it. Callers submit() a request
+    and get back an id; results are collected via get_result() once
+    processing has advanced far enough (run_fifo() / run_round_robin()
+    drive that, they aren't automatic/background by default -- call
+    tick() from whatever event loop the rest of the system already
+    has, e.g. the orchestrator's own tick()).
+    """
+
+    def __init__(self, strategy: SchedulingStrategy = SchedulingStrategy.FIFO, rr_quantum: float = 2.0):
+        self.strategy = strategy
+        self.rr_quantum = rr_quantum
+        self.queue: deque[SyscallRequest] = deque()
+        self.requests: dict[str, SyscallRequest] = {}
+        self._interrupt_controller = ContextInterruptController()
+
+    def submit(self, agent_name: str, core: LLMWrapper, syscall: LLMSyscall, priority: int = 0) -> str:
+        req = SyscallRequest(agent_name=agent_name, core=core, syscall=syscall, priority=priority)
+        self.requests[req.id] = req
+        self.queue.append(req)
+        return req.id
+
+    def get_result(self, request_id: str) -> SyscallRequest:
+        return self.requests[request_id]
+
+    def pending_count(self) -> int:
+        return len(self.queue)
+
+    # -- scheduling strategies -------------------------------------------------
+
+    def run_fifo_once(self) -> Optional[SyscallRequest]:
+        """Pops the head of the queue and runs it straight through to
+        completion. No interruption -- a long request occupies the
+        core for its full duration before the next one gets to start,
+        by design."""
+        if not self.queue:
+            return None
+
+        req = self.queue.popleft()
+        req.status = SyscallStatus.RUNNING
+        try:
+            req.result = req.core.chat(req.syscall)
+            req.accumulated_text = req.result.text
+            req.status = SyscallStatus.DONE
+        except Exception as e:
+            req.error = str(e)
+            req.status = SyscallStatus.ERROR
+        return req
+
+    def run_round_robin_tick(self) -> Optional[SyscallRequest]:
+        """Runs the head-of-queue request for at most `rr_quantum`
+        seconds, via ContextInterruptController -- which picks the
+        text-based or logits-based snapshot approach depending on what
+        the request's core supports. If it finishes within that slice,
+        it's marked DONE and NOT re-queued. If the slice runs out
+        first:
+        - text-based: the partial output is folded into a continuation
+          syscall (build_continuation) and pushed to the back of the
+          queue.
+        - logits-based: the request's `syscall` is left UNCHANGED
+          (nothing to rebuild -- the search-tree state itself is what
+          gets resumed) and its `snapshot` is stored for the next tick
+          to resume from directly.
+        Either way, every other queued request gets its own slice
+        before this one is revisited -- that's the "round" in Round
+        Robin."""
+        if not self.queue:
+            return None
+
+        req = self.queue.popleft()
+        req.status = SyscallStatus.RUNNING
+        try:
+            partial = self._interrupt_controller.run_slice(
+                req.core, req.syscall, self.rr_quantum, resume_snapshot=req.snapshot
+            )
+        except Exception as e:
+            req.error = str(e)
+            req.status = SyscallStatus.ERROR
+            return req
+
+        req.slices_run += 1
+        if partial.cumulative:
+            # Logits mode: beam decoding naturally tracks the full
+            # sequence generated so far, so this IS the total, not a
+            # delta to append.
+            req.accumulated_text = partial.text_so_far
+        else:
+            # Text mode: each resumed stream only yields its own new
+            # continuation text.
+            req.accumulated_text += partial.text_so_far
+
+        if partial.complete:
+            req.result = LLMResponse(text=req.accumulated_text)
+            req.status = SyscallStatus.DONE
+            req.snapshot = None
+        else:
+            if partial.snapshot.mode == "text":
+                req.syscall = self._interrupt_controller.build_continuation(req.syscall, partial)
+                req.snapshot = None
+            else:
+                req.snapshot = partial.snapshot
+            req.status = SyscallStatus.QUEUED
+            self.queue.append(req)
+
+        return req
+
+    def tick(self) -> Optional[SyscallRequest]:
+        """Advances the scheduler by exactly one step, using whichever
+        strategy this instance was configured with."""
+        if self.strategy == SchedulingStrategy.FIFO:
+            return self.run_fifo_once()
+        return self.run_round_robin_tick()
+
+    def drain(self, max_ticks: int = 10_000) -> list[SyscallRequest]:
+        """Runs tick() until the queue is empty (every request reaches
+        DONE or ERROR) or max_ticks is hit, as a safety bound against
+        an unbounded RR cycle. Returns the requests in the order they
+        finished."""
+        finished = []
+        ticks = 0
+        while self.queue and ticks < max_ticks:
+            req = self.tick()
+            ticks += 1
+            if req is not None and req.status in (SyscallStatus.DONE, SyscallStatus.ERROR):
+                finished.append(req)
+        return finished