mlxsmith 0.1.0__py3-none-any.whl

mlxsmith/util.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from rich.console import Console
+
+console = Console()
+
+def sha1_text(s: str) -> str:
+    return hashlib.sha1(s.encode("utf-8")).hexdigest()
+
+def ensure_dir(p: Path) -> Path:
+    p.mkdir(parents=True, exist_ok=True)
+    return p
+
+def write_jsonl(path: Path, rows):
+    ensure_dir(path.parent)
+    with path.open("a", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+def now_ts() -> str:
+    return time.strftime("%Y%m%d_%H%M%S")
+
+def run_cmd(cmd: list[str], cwd: Optional[Path] = None, timeout_s: Optional[int] = None) -> subprocess.CompletedProcess:
+    return subprocess.run(cmd, cwd=str(cwd) if cwd else None, timeout=timeout_s, capture_output=True, text=True)
+
+@dataclass
+class SystemInfo:
+    python: str
+    python_arch: str
+    platform: str
+    macos_version: Optional[str]
+    machine: str
+    cpu_count: int
+    has_metal: Optional[bool]
+    has_mlx: bool
+    mlx_version: Optional[str]
+    has_zmlx: bool
+
+def detect_system() -> SystemInfo:
+    has_mlx = False
+    mlx_version = None
+    try:
+        import mlx  # type: ignore
+        has_mlx = True
+        mlx_version = getattr(mlx, "__version__", None)
+    except Exception:
+        pass
+
+    import importlib.util
+
+    has_zmlx = importlib.util.find_spec("zmlx") is not None
+
+    # Metal detection (best-effort): on macOS we assume Metal is present; for CI, this is not reliable.
+    has_metal = None
+    if sys.platform == "darwin":
+        has_metal = True
+
+    macos_version = None
+    if sys.platform == "darwin":
+        macos_version = platform.mac_ver()[0] or None
+
+    py_arch = platform.architecture()[0] or "unknown"
+
+    return SystemInfo(
+        python=sys.version.split()[0],
+        python_arch=py_arch,
+        platform=platform.platform(),
+        macos_version=macos_version,
+        machine=platform.machine(),
+        cpu_count=os.cpu_count() or 0,
+        has_metal=has_metal,
+        has_mlx=has_mlx,
+        mlx_version=mlx_version,
+        has_zmlx=has_zmlx,
+    )
+
+def require(cond: bool, msg: str):
+    if not cond:
+        raise RuntimeError(msg)
+
+def copytree(src: Path, dst: Path):
+    if dst.exists():
+        shutil.rmtree(dst)
+    shutil.copytree(src, dst)
+
+
+def tree_map(fn, tree):
+    if tree is None:
+        return None
+    if isinstance(tree, dict):
+        return {k: tree_map(fn, v) for k, v in tree.items()}
+    if isinstance(tree, (list, tuple)):
+        return type(tree)(tree_map(fn, v) for v in tree)
+    return fn(tree)
+
+
+def tree_add(a, b):
+    if a is None:
+        return b
+    if b is None:
+        return a
+    if isinstance(a, dict) and isinstance(b, dict):
+        return {k: tree_add(a.get(k), b.get(k)) for k in a.keys() | b.keys()}
+    if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
+        return type(a)(tree_add(x, y) for x, y in zip(a, b))
+    return a + b
+
+
+def tree_scale(tree, scale: float):
+    return tree_map(lambda x: x * scale, tree)
+
+
+def tree_leaves(tree) -> list:
+    leaves = []
+    if tree is None:
+        return leaves
+    if isinstance(tree, dict):
+        for v in tree.values():
+            leaves.extend(tree_leaves(v))
+    elif isinstance(tree, (list, tuple)):
+        for v in tree:
+            leaves.extend(tree_leaves(v))
+    else:
+        leaves.append(tree)
+    return leaves
+
+
+def clip_grad_norm(grads, max_norm: float):
+    """Clip gradients by global L2 norm. Returns clipped grads."""
+    import mlx.core as mx
+
+    leaves = tree_leaves(grads)
+    if not leaves:
+        return grads
+    total_norm_sq = mx.array(0.0)
+    for g in leaves:
+        total_norm_sq = total_norm_sq + (g * g).sum()
+    total_norm = mx.sqrt(total_norm_sq)
+    clip_coef = mx.minimum(mx.array(max_norm) / mx.maximum(total_norm, mx.array(1e-8)), mx.array(1.0))
+    return tree_map(lambda g: g * clip_coef, grads)
+
+
+def latency_summary_ms(samples: list[float]) -> dict[str, float]:
+    if not samples:
+        return {}
+    items = sorted(samples)
+    n = len(items)
+    mean = sum(items) / n
+
+    def _pct(p: float) -> float:
+        if n == 1:
+            return items[0]
+        idx = int((p / 100.0) * (n - 1))
+        return items[max(0, min(idx, n - 1))]
+
+    return {
+        "mean": mean,
+        "p50": _pct(50),
+        "p90": _pct(90),
+        "p99": _pct(99),
+        "max": items[-1],
+    }

mlxsmith/verifiers/__init__.py

@@ -0,0 +1,3 @@
+from .types import VerifyResult
+
+__all__ = ["VerifyResult"]

mlxsmith/verifiers/compose.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import importlib.util
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .types import VerifyResult
+
+
+def _load_verifier(path: str):
+    verifier_path = Path(path)
+    spec = importlib.util.spec_from_file_location(verifier_path.stem, verifier_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load verifier: {verifier_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)  # type: ignore
+    verify_fn = getattr(module, "verify", None)
+    if not callable(verify_fn):
+        raise RuntimeError(f"Verifier must define verify(...): {verifier_path}")
+    return verify_fn
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    verifiers: List[Any],
+    mode: str = "all",
+    weights: Optional[List[float]] = None,
+    per_verifier_kwargs: Optional[Dict[str, Dict[str, Any]]] = None,
+    **kwargs,
+) -> VerifyResult:
+    """Compose multiple verifiers with AND/OR/weighted reward aggregation.
+
+    verifiers: list of paths or dicts {path, kwargs}.
+    mode: all | any | weighted
+    weights: optional weights for weighted reward aggregation.
+    per_verifier_kwargs: optional mapping of path -> kwargs.
+    """
+    results = []
+    latencies: Dict[str, float] = {}
+    per_kwargs = per_verifier_kwargs or {}
+
+    for idx, entry in enumerate(verifiers):
+        if isinstance(entry, dict):
+            path = entry.get("path")
+            extra = entry.get("kwargs") or {}
+        else:
+            path = entry
+            extra = {}
+
+        if not path:
+            continue
+        verify_fn = _load_verifier(str(path))
+        merged = dict(kwargs)
+        merged.update(per_kwargs.get(str(path), {}))
+        merged.update(extra)
+
+        t0 = time.time()
+        res = verify_fn(prompt, completion, workdir, **merged)
+        latencies[str(path)] = (time.time() - t0) * 1000.0
+        results.append((str(path), res))
+
+    if not results:
+        return VerifyResult(
+            reward=0.0,
+            passed=False,
+            info={"error": "no_verifiers"},
+            artifacts_dir=workdir,
+        )
+
+    mode = (mode or "all").lower()
+    passes = [bool(getattr(r, "passed", False)) for _p, r in results]
+    rewards = [float(getattr(r, "reward", 0.0)) for _p, r in results]
+
+    if mode == "any":
+        passed = any(passes)
+        reward = max(rewards) if rewards else 0.0
+    elif mode == "weighted":
+        if weights and len(weights) == len(rewards):
+            reward = sum(w * r for w, r in zip(weights, rewards))
+        else:
+            reward = sum(rewards) / max(1, len(rewards))
+        passed = reward > 0.0 and any(passes)
+    else:
+        passed = all(passes)
+        reward = sum(rewards) / max(1, len(rewards))
+
+    return VerifyResult(
+        reward=reward,
+        passed=passed,
+        info={
+            "mode": mode,
+            "verifiers": [
+                {
+                    "path": path,
+                    "passed": bool(getattr(res, "passed", False)),
+                    "reward": float(getattr(res, "reward", 0.0)),
+                    "info": getattr(res, "info", {}),
+                }
+                for path, res in results
+            ],
+            "verifier_latencies_ms": latencies,
+        },
+        artifacts_dir=workdir,
+    )
+

mlxsmith/verifiers/docker_verifier.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+
+from .types import VerifyResult
+from .pytest_verifier import verify as local_verify
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    tests_subdir: str = "tests",
+    timeout_s: int = 30,
+    reward_pass: float = 1.0,
+    reward_fail: float = 0.0,
+    image: str = "python:3.11-slim",
+    memory_mb: int = 512,
+    cpus: float = 1.0,
+    pids: int = 128,
+    use_local_fallback: bool = True,
+) -> VerifyResult:
+    """Run pytest inside a locked-down Docker container.
+
+    Falls back to local pytest verifier if Docker is unavailable and use_local_fallback is True.
+    """
+    if shutil.which("docker") is None:
+        if use_local_fallback:
+            return local_verify(
+                prompt,
+                completion,
+                workdir,
+                tests_subdir=tests_subdir,
+                timeout_s=timeout_s,
+                reward_pass=reward_pass,
+                reward_fail=reward_fail,
+            )
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": "docker_not_found"},
+            artifacts_dir=workdir,
+        )
+
+    wd = Path(workdir)
+    wd.mkdir(parents=True, exist_ok=True)
+    if not any(wd.glob("*.py")):
+        (wd / "main.py").write_text(completion, encoding="utf-8")
+
+    tests_path = wd / tests_subdir
+    if not tests_path.exists():
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": f"Missing tests folder: {tests_subdir}"},
+            artifacts_dir=str(wd),
+        )
+
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "--network",
+        "none",
+        "--read-only",
+        "--pids-limit",
+        str(int(pids)),
+        "--memory",
+        f"{int(memory_mb)}m",
+        "--cpus",
+        str(float(cpus)),
+        "--tmpfs",
+        "/tmp:rw,noexec,nosuid,nodev",
+        "-v",
+        f"{wd}:/workspace:rw",
+        "-w",
+        "/workspace",
+        image,
+        "pytest",
+        "-q",
+    ]
+
+    try:
+        proc = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=timeout_s,
+        )
+        passed = proc.returncode == 0
+        return VerifyResult(
+            reward=reward_pass if passed else reward_fail,
+            passed=passed,
+            info={
+                "returncode": proc.returncode,
+                "stdout": proc.stdout[-4000:],
+                "stderr": proc.stderr[-4000:],
+                "docker_image": image,
+            },
+            artifacts_dir=str(wd),
+        )
+    except subprocess.TimeoutExpired:
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": "timeout", "timeout_s": timeout_s},
+            artifacts_dir=str(wd),
+        )

mlxsmith/verifiers/jsonschema.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from jsonschema import validate, ValidationError
+
+from .types import VerifyResult
+
+
+def _load_schema(schema: Any) -> Any:
+    if isinstance(schema, str):
+        p = Path(schema)
+        if p.exists():
+            return json.loads(p.read_text(encoding="utf-8"))
+    return schema
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    schema: Any,
+    reward_pass: float = 1.0,
+    reward_fail: float = 0.0,
+) -> VerifyResult:
+    try:
+        data = json.loads(completion)
+    except json.JSONDecodeError as e:
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": "invalid_json", "detail": str(e)},
+            artifacts_dir=None,
+        )
+
+    schema_obj = _load_schema(schema)
+    try:
+        validate(instance=data, schema=schema_obj)
+        return VerifyResult(
+            reward=reward_pass,
+            passed=True,
+            info={"schema": schema_obj},
+            artifacts_dir=None,
+        )
+    except ValidationError as e:
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": "schema_validation", "detail": str(e)},
+            artifacts_dir=None,
+        )

mlxsmith/verifiers/pytest_verifier.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+from .types import VerifyResult
+
+def _sandbox_env(base_env: dict | None = None, *, workdir: Path) -> dict:
+    env = dict(os.environ)
+    if base_env:
+        env.update(base_env)
+    env.setdefault("PYTHONHASHSEED", "0")
+    env["HOME"] = str(workdir)
+    env["TMPDIR"] = str(workdir)
+    env["PYTHONPATH"] = str(workdir)
+    env.pop("PYTHONSTARTUP", None)
+    env.pop("VIRTUAL_ENV", None)
+    return env
+
+
+def verify(
+    prompt: str,
+    completion: str,
+    workdir: str,
+    *,
+    tests_subdir: str = "tests",
+    timeout_s: int = 30,
+    reward_pass: float = 1.0,
+    reward_fail: float = 0.0,
+) -> VerifyResult:
+    """Run pytest in a sandbox directory.
+
+    Expected usage: environment builder writes files into `workdir` and this verifier runs tests there.
+    For v0, we simply run `pytest -q` and treat exit code 0 as pass.
+
+    IMPORTANT: This runs locally. For stronger isolation, replace with a sandbox runner.
+    """
+    wd = Path(workdir)
+    wd.mkdir(parents=True, exist_ok=True)
+
+    # If user is doing code tasks, they can write completion into a file convention, e.g., main.py
+    # Here we create a default file if none exist:
+    if not any(wd.glob("*.py")):
+        (wd / "main.py").write_text(completion, encoding="utf-8")
+
+    # Ensure tests folder exists; if not, fail deterministically.
+    tests_path = wd / tests_subdir
+    if not tests_path.exists():
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": f"Missing tests folder: {tests_subdir}"},
+            artifacts_dir=str(wd),
+        )
+
+    try:
+        proc = subprocess.run(
+            ["pytest", "-q"],
+            cwd=str(wd),
+            capture_output=True,
+            text=True,
+            timeout=timeout_s,
+            env=_sandbox_env(workdir=wd),
+        )
+        passed = proc.returncode == 0
+        return VerifyResult(
+            reward=reward_pass if passed else reward_fail,
+            passed=passed,
+            info={
+                "returncode": proc.returncode,
+                "stdout": proc.stdout[-4000:],
+                "stderr": proc.stderr[-4000:],
+            },
+            artifacts_dir=str(wd),
+        )
+    except subprocess.TimeoutExpired:
+        return VerifyResult(
+            reward=reward_fail,
+            passed=False,
+            info={"error": "timeout", "timeout_s": timeout_s},
+            artifacts_dir=str(wd),
+        )

mlxsmith/verifiers/regex.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+import re
+
+from .types import VerifyResult
+
+def verify(prompt: str, completion: str, workdir: str, *, pattern: str, flags: int = 0, reward_pass: float = 1.0, reward_fail: float = 0.0) -> VerifyResult:
+    m = re.search(pattern, completion, flags=flags)
+    passed = m is not None
+    return VerifyResult(
+        reward=reward_pass if passed else reward_fail,
+        passed=passed,
+        info={"pattern": pattern, "match": m.group(0) if m else None},
+        artifacts_dir=None,
+    )

mlxsmith/verifiers/types.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+@dataclass
+class VerifyResult:
+    reward: float
+    passed: bool
+    info: Dict[str, Any]
+    artifacts_dir: Optional[str] = None