API Reference for shared¶

def publish(self, artifact: SubtaskArtifact) -> None:
    """Publish a subtask's artifacts."""
    self._artifacts[artifact.name] = artifact

def get(self, name: str) -> SubtaskArtifact | None:
    """Retrieve a subtask's artifacts by name."""
    return self._artifacts.get(name)

def get_dependency_interfaces(self, subtask: dict[str, object]) -> str:
    """Get concatenated interfaces for a subtask's declared dependencies."""
    raw_deps = subtask.get("depends_on", [])
    deps = raw_deps if isinstance(raw_deps, list) else []
    if not deps:
        return ""
    parts: list[str] = []
    for dep_name in deps:
        artifact = self._artifacts.get(str(dep_name))
        if artifact and artifact.interfaces:
            parts.append(f"--- {artifact.name} ---\n{artifact.interfaces}")
    return "\n".join(parts)

def all_interfaces(self) -> str:
    """Get all published interfaces."""
    parts: list[str] = []
    for artifact in self._artifacts.values():
        if artifact.interfaces:
            parts.append(f"--- {artifact.name} ---\n{artifact.interfaces}")
    return "\n".join(parts)

def mark_running(self, task_hash: str, agent_id: str) -> None:
    """Record that an agent has started working on a task.

    Args:
        task_hash: Hash of the task being started.
        agent_id: Identifier of the agent executing the task.
    """
    now = datetime.now(timezone.utc).isoformat()
    record = CheckpointRecord(
        run_id=self._run_id,
        task_hash=task_hash,
        agent_id=agent_id,
        status="running",
        started_at=now,
    )
    with Session(self._engine, expire_on_commit=False) as session:
        session.add(record)
        session.commit()

def mark_completed(self, task_hash: str, outcome: str) -> None:
    """Record that a task has been completed.

    Args:
        task_hash: Hash of the completed task.
        outcome: Description of the result.
    """
    now = datetime.now(timezone.utc).isoformat()
    with Session(self._engine, expire_on_commit=False) as session:
        stmt = (
            select(CheckpointRecord)
            .where(
                CheckpointRecord.task_hash == task_hash,
                CheckpointRecord.run_id == self._run_id,
                CheckpointRecord.status == "running",
            )
            .limit(1)
        )
        record = session.exec(stmt).first()
        if record is not None:
            record.status = "completed"
            record.outcome = outcome
            record.completed_at = now
            session.add(record)
            session.commit()
        else:
            # No running record found — this means mark_running() was never
            # called for this task. Log a warning so the bug is visible,
            # then create a completed record to keep the DB consistent.
            logger.warning(
                "mark_completed called for task_hash=%r but no running record "
                "was found (run_id=%r). mark_running() may not have been called. "
                "Creating orphaned completed record with agent_id='unknown'.",
                task_hash,
                self._run_id,
            )
            new_record = CheckpointRecord(
                run_id=self._run_id,
                task_hash=task_hash,
                agent_id="unknown",
                status="completed",
                outcome=outcome,
                completed_at=now,
            )
            session.add(new_record)
            session.commit()

def mark_failed(self, task_hash: str, agent_id: str) -> None:
    """Record that a task execution has failed.

    Args:
        task_hash: Hash of the failed task.
        agent_id: Identifier of the agent that failed.
    """
    now = datetime.now(timezone.utc).isoformat()
    with Session(self._engine, expire_on_commit=False) as session:
        stmt = (
            select(CheckpointRecord)
            .where(
                CheckpointRecord.task_hash == task_hash,
                CheckpointRecord.run_id == self._run_id,
                CheckpointRecord.agent_id == agent_id,
                CheckpointRecord.status == "running",
            )
            .limit(1)
        )
        record = session.exec(stmt).first()
        if record is not None:
            record.status = "failed"
            record.completed_at = now
            session.add(record)
            session.commit()
        else:
            # No running record found — log a warning and still create the
            # failed record so the failure is visible in the DB.
            logger.warning(
                "mark_failed called for task_hash=%r agent_id=%r but no running "
                "record was found (run_id=%r). mark_running() may not have been "
                "called. Creating orphaned failed record.",
                task_hash,
                agent_id,
                self._run_id,
            )
            new_record = CheckpointRecord(
                run_id=self._run_id,
                task_hash=task_hash,
                agent_id=agent_id,
                status="failed",
                completed_at=now,
            )
            session.add(new_record)
            session.commit()

def is_completed(self, task_hash: str) -> bool:
    """Check if a task has been completed.

    Args:
        task_hash: Hash of the task to check.

    Returns:
        True if the task has a "completed" checkpoint record.
    """
    with Session(self._engine, expire_on_commit=False) as session:
        stmt = (
            select(CheckpointRecord)
            .where(
                CheckpointRecord.task_hash == task_hash,
                CheckpointRecord.run_id == self._run_id,
                CheckpointRecord.status == "completed",
            )
            .limit(1)
        )
        return session.exec(stmt).first() is not None

@classmethod
def detect(cls) -> HardwareProbe:
    """Detect hardware resources on the current machine.

    Uses psutil for CPU/RAM detection (always available). Attempts to
    import pynvml for GPU detection; falls back to empty GPU list if
    pynvml is not installed or NVML initialization fails.

    Returns:
        A HardwareProbe with detected resource information.
    """
    cpu_count = psutil.cpu_count(logical=True) or 1
    ram_total_mb = psutil.virtual_memory().total // (1024 * 1024)

    gpus: list[GPUInfo] = []
    try:
        import pynvml

        pynvml.nvmlInit()
        device_count = pynvml.nvmlDeviceGetCount()
        for i in range(device_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(name, bytes):
                name = name.decode("utf-8")
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            vram_mb = mem_info.total // (1024 * 1024)
            gpus.append(GPUInfo(name=name, vram_mb=vram_mb))
        pynvml.nvmlShutdown()
    except ImportError:
        pass  # pynvml not installed; CPU-only mode
    except Exception as e:
        logger.debug("GPU detection failed (will run CPU-only): %s", e)

    return cls(cpu_count=cpu_count, ram_total_mb=ram_total_mb, gpus=gpus)

def compute_budget(self, base_model_vram_mb: int = 0) -> HardwareBudget:
    """Compute concurrency limits based on detected hardware.

    Args:
        base_model_vram_mb: VRAM consumed by the base inference model.
            Used to calculate how much VRAM is left for LoRA adapters
            and training.

    Returns:
        A HardwareBudget with computed concurrency limits.
    """
    num_gpus = len(self.gpus)
    single_gpu_mode = num_gpus <= 1

    if num_gpus == 0:
        return HardwareBudget(
            max_agents=max(1, self.cpu_count // 2),
            max_concurrent_loras=0,
            training_slots=0,
            vram_per_gpu_mb=0,
            single_gpu_mode=True,
        )

    vram_per_gpu = self.gpus[0].vram_mb
    free_vram = max(0, vram_per_gpu - base_model_vram_mb)

    max_loras = max(1, free_vram // VRAM_PER_LORA_MB)
    training_slots = 1 if free_vram >= MIN_TRAINING_VRAM_MB else 0

    # Agents limited by CPU cores and available VRAM
    max_agents = max(1, min(self.cpu_count // 2, num_gpus * 4))

    return HardwareBudget(
        max_agents=max_agents,
        max_concurrent_loras=max_loras,
        training_slots=training_slots,
        vram_per_gpu_mb=vram_per_gpu,
        single_gpu_mode=single_gpu_mode,
    )

def to_dict(self) -> dict[str, Any]:
    """Serialize to a plain dict."""
    d = asdict(self)
    # Convert tuple back for JSON compatibility
    d["calibration"]["scaling_range"] = list(d["calibration"]["scaling_range"])
    return d

def save(self, path: Path | None = None) -> Path:
    """Write config to JSON file."""
    path = path or (_DEFAULT_CONFIG_DIR / _CONFIG_FILENAME)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(self.to_dict(), indent=2))
    return path

def override(self, **kwargs: Any) -> PipelineConfig:
    """Return a new config with selected fields replaced.

    Accepts dotted keys like ``adapter.scaling=0.1`` or flat
    section dicts like ``adapter={"scaling": 0.1}``.
    """
    d = self.to_dict()
    for key, value in kwargs.items():
        if "." in key:
            section, field_name = key.split(".", 1)
            d.setdefault(section, {})[field_name] = value
        elif isinstance(value, dict):
            d.setdefault(key, {}).update(value)
        else:
            d[key] = value
    return _from_dict(d)