mlxsmith 0.1.0__py3-none-any.whl

mlxsmith/envs/system.py

@@ -0,0 +1,388 @@
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import tarfile
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import yaml
+
+from ..util import ensure_dir, now_ts, copytree
+
+
+@dataclass
+class EnvManifest:
+    name: str
+    version: str
+    description: Optional[str] = None
+    verifier: Optional[str] = None
+    tasks: Optional[list] = None
+    token_env: Optional[object] = None
+
+
+@dataclass
+class EnvRef:
+    name: str
+    version: Optional[str] = None
+
+
+def _envs_root(project_root: Path) -> Path:
+    return project_root / "envs"
+
+
+def _registry_path(project_root: Path) -> Path:
+    return _envs_root(project_root) / "registry.json"
+
+
+def load_manifest(env_path: Path) -> EnvManifest:
+    data = yaml.safe_load(env_path.read_text(encoding="utf-8")) or {}
+    return EnvManifest(
+        name=str(data.get("name") or env_path.parent.name),
+        version=str(data.get("version") or "0.1.0"),
+        description=data.get("description"),
+        verifier=data.get("verifier"),
+        tasks=data.get("tasks"),
+        token_env=data.get("token_env"),
+    )
+
+
+def _normalize_package_module(name: str) -> str:
+    cleaned = re.sub(r"[^a-zA-Z0-9_]+", "_", name.replace("-", "_"))
+    cleaned = cleaned.strip("_") or "env"
+    if cleaned[0].isdigit():
+        cleaned = f"env_{cleaned}"
+    return cleaned.lower()
+
+
+def _env_scaffold_pyproject(name: str, version: str, description: str) -> str:
+    return f"""[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "{name}"
+version = "{version}"
+description = "{description}"
+readme = "README.md"
+requires-python = ">=3.10"
+
+[tool.setuptools]
+package-dir = {{"" = "src"}}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+"""
+
+
+def _env_scaffold_readme(name: str) -> str:
+    return f"""# {name}
+
+Local MLXSmith environment package.
+
+## Files
+- `env.yaml`: task manifest consumed by mlxsmith.
+- `pyproject.toml`: Python package metadata for Hub publishing.
+
+## Usage
+```bash
+mlxsmith env package {name}
+mlxsmith env publish envs/packages/{name}-0.1.0.tar.gz
+```
+"""
+
+
+def _env_scaffold_module() -> str:
+    return """from pathlib import Path
+
+ENV_MANIFEST = Path(__file__).resolve().parents[2] / "env.yaml"
+
+
+def load_environment() -> Path:
+    return ENV_MANIFEST
+"""
+
+
+def _env_scaffold_init() -> str:
+    return """from .environment import ENV_MANIFEST, load_environment
+
+__all__ = ["ENV_MANIFEST", "load_environment"]
+"""
+
+
+def resolve_env_path(project_root: Path, env_ref: str) -> Path:
+    ref = Path(env_ref)
+    if ref.exists():
+        if ref.is_dir():
+            candidate = ref / "env.yaml"
+            return candidate if candidate.exists() else ref
+        return ref
+    candidate = _envs_root(project_root) / env_ref / "env.yaml"
+    return candidate
+
+
+def init_env(project_root: Path, name: str) -> Path:
+    env_root = _envs_root(project_root) / name
+    ensure_dir(env_root)
+    manifest = {
+        "name": name,
+        "version": "0.1.0",
+        "description": "Sample environment",
+        "verifier": "verifiers/regex.py",
+        "tasks": [
+            {
+                "id": "add",
+                "prompt": "Write a Python function add(a, b) that returns the sum.",
+                "tests": "from main import add\\n\\n\\n"
+                "def test_add():\\n"
+                "    assert add(2, 3) == 5\\n",
+            }
+        ],
+    }
+    (env_root / "env.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False), encoding="utf-8")
+
+    pkg_module = _normalize_package_module(name)
+    (env_root / "pyproject.toml").write_text(
+        _env_scaffold_pyproject(name, manifest["version"], manifest["description"]),
+        encoding="utf-8",
+    )
+    (env_root / "README.md").write_text(_env_scaffold_readme(name), encoding="utf-8")
+    pkg_dir = ensure_dir(env_root / "src" / pkg_module)
+    (pkg_dir / "environment.py").write_text(_env_scaffold_module(), encoding="utf-8")
+    (pkg_dir / "__init__.py").write_text(_env_scaffold_init(), encoding="utf-8")
+    return env_root
+
+
+def _load_registry(project_root: Path) -> dict:
+    path = _registry_path(project_root)
+    if not path.exists():
+        return {"packages": [], "updated_at": None}
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _save_registry(project_root: Path, data: dict) -> None:
+    data["updated_at"] = now_ts()
+    path = _registry_path(project_root)
+    ensure_dir(path.parent)
+    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+
+
+def _parse_env_ref(env_ref: str, version: Optional[str] = None) -> EnvRef:
+    name = env_ref.strip()
+    parsed_version = None
+    if "==" in name:
+        name, parsed_version = name.split("==", 1)
+    elif "@" in name:
+        name, parsed_version = name.rsplit("@", 1)
+    name = name.strip()
+    if parsed_version is not None:
+        parsed_version = parsed_version.strip()
+    if parsed_version in {"", "latest"}:
+        parsed_version = None
+    if version and parsed_version and version != parsed_version:
+        raise RuntimeError(f"Conflicting versions: {parsed_version} vs {version}")
+    return EnvRef(name=name, version=version or parsed_version)
+
+
+def _version_key(version: str) -> tuple:
+    if not version:
+        return tuple()
+    base, _, _build = version.partition("+")
+    main, _, pre = base.partition("-")
+    parts = []
+    for part in main.split("."):
+        if part.startswith("v") and part[1:].isdigit():
+            part = part[1:]
+        if part.isdigit():
+            parts.append((0, int(part)))
+        else:
+            parts.append((1, part))
+    if pre:
+        parts.append((1, pre))
+    else:
+        parts.append((2, ""))
+    return tuple(parts)
+
+
+def _select_registry_package(packages: list[dict], name: str, version: Optional[str]) -> dict:
+    matches = [p for p in packages if p.get("name") == name]
+    if not matches:
+        raise RuntimeError(f"Env not found in registry: {name}")
+    if version:
+        exact = [p for p in matches if p.get("version") == version]
+        if not exact:
+            available = sorted({p.get("version", "") for p in matches})
+            raise RuntimeError(f"Env {name} has no version {version}. Available: {', '.join(available)}")
+        return exact[0]
+    return sorted(matches, key=lambda p: _version_key(p.get("version", "")))[-1]
+
+
+def _load_manifest_from_package(package_path: Path) -> EnvManifest:
+    with tarfile.open(package_path, "r:gz") as tf:
+        for m in tf.getmembers():
+            if m.name.endswith("env.yaml"):
+                f = tf.extractfile(m)
+                if f is None:
+                    break
+                data = yaml.safe_load(f.read().decode("utf-8")) or {}
+                return EnvManifest(
+                    name=str(data.get("name") or Path(m.name).parent.name),
+                    version=str(data.get("version") or "0.1.0"),
+                    description=data.get("description"),
+                    verifier=data.get("verifier"),
+                    tasks=data.get("tasks"),
+                    token_env=data.get("token_env"),
+                )
+    raise RuntimeError("Package missing env.yaml")
+
+
+def list_registry_packages(project_root: Path, name: Optional[str] = None, all_versions: bool = False) -> list[dict]:
+    registry = _load_registry(project_root)
+    packages = registry.get("packages", [])
+    if name:
+        packages = [p for p in packages if p.get("name") == name]
+    if all_versions or not packages:
+        return sorted(packages, key=lambda p: (p.get("name", ""), _version_key(p.get("version", ""))))
+
+    latest = {}
+    for pkg in packages:
+        pkg_name = pkg.get("name")
+        if not pkg_name:
+            continue
+        prev = latest.get(pkg_name)
+        if prev is None or _version_key(pkg.get("version", "")) > _version_key(prev.get("version", "")):
+            latest[pkg_name] = pkg
+    return sorted(latest.values(), key=lambda p: p.get("name", ""))
+
+
+def registry_info(project_root: Path, env_ref: str, version: Optional[str] = None) -> tuple[dict, EnvManifest]:
+    registry = _load_registry(project_root)
+    ref = _parse_env_ref(env_ref, version=version)
+    pkg = _select_registry_package(registry.get("packages", []), ref.name, ref.version)
+    pkg_path = Path(pkg["path"])
+    if not pkg_path.exists():
+        raise RuntimeError(f"Registry package missing: {pkg_path}")
+    manifest = _load_manifest_from_package(pkg_path)
+    return pkg, manifest
+
+
+def install_env(
+    project_root: Path,
+    source: str,
+    version: Optional[str] = None,
+) -> Path:
+    src = Path(source)
+    ensure_dir(_envs_root(project_root))
+    if src.exists():
+        if version:
+            raise RuntimeError("Version pinning only applies to registry installs.")
+        if src.is_dir():
+            manifest_path = src / "env.yaml"
+            if not manifest_path.exists():
+                raise RuntimeError(f"Missing env.yaml in {src}")
+            manifest = load_manifest(manifest_path)
+            dest = _envs_root(project_root) / manifest.name
+            copytree(src, dest)
+            return dest
+        if src.suffixes[-2:] == [".tar", ".gz"]:
+            with tarfile.open(src, "r:gz") as tf:
+                members = tf.getmembers()
+                top = members[0].name.split("/")[0] if members else ""
+                tf.extractall(_envs_root(project_root))
+                env_path = _envs_root(project_root) / top / "env.yaml"
+                if env_path.exists():
+                    return env_path.parent
+            raise RuntimeError(f"Invalid package: {src}")
+        raise RuntimeError(f"Unsupported env source: {src}")
+
+    # treat as registry name
+    registry = _load_registry(project_root)
+    ref = _parse_env_ref(source, version=version)
+    pkg = _select_registry_package(registry.get("packages", []), ref.name, ref.version)
+    pkg_path = Path(pkg["path"])
+    return install_env(project_root, str(pkg_path))
+
+
+def package_env(project_root: Path, env_name: str, out_path: Optional[str] = None) -> Path:
+    env_dir = _envs_root(project_root) / env_name
+    manifest_path = env_dir / "env.yaml"
+    if not manifest_path.exists():
+        raise RuntimeError(f"Missing env.yaml in {env_dir}")
+    manifest = load_manifest(manifest_path)
+
+    out_dir = Path(out_path) if out_path else _envs_root(project_root) / "packages"
+    ensure_dir(out_dir)
+    tar_path = out_dir / f"{manifest.name}-{manifest.version}.tar.gz"
+    with tarfile.open(tar_path, "w:gz") as tf:
+        tf.add(env_dir, arcname=env_dir.name)
+    return tar_path
+
+
+def publish_env(project_root: Path, package_path: str) -> Path:
+    pkg = Path(package_path)
+    if not pkg.exists():
+        raise RuntimeError(f"Missing package: {pkg}")
+
+    manifest = _load_manifest_from_package(pkg)
+
+    registry_dir = _envs_root(project_root) / "registry"
+    ensure_dir(registry_dir)
+    dest = registry_dir / f"{manifest.name}-{manifest.version}.tar.gz"
+    dest.write_bytes(pkg.read_bytes())
+
+    registry = _load_registry(project_root)
+    registry.setdefault("packages", [])
+    registry["packages"] = [
+        p
+        for p in registry["packages"]
+        if not (p.get("name") == manifest.name and p.get("version") == manifest.version)
+    ]
+    registry["packages"].append(
+        {
+            "name": manifest.name,
+            "version": manifest.version,
+            "description": manifest.description,
+            "verifier": manifest.verifier,
+            "path": str(dest),
+        }
+    )
+    _save_registry(project_root, registry)
+    return dest
+
+
+def pull_env(
+    project_root: Path,
+    env_ref: str,
+    out_dir: Optional[str] = None,
+    version: Optional[str] = None,
+    force: bool = False,
+) -> Path:
+    registry = _load_registry(project_root)
+    ref = _parse_env_ref(env_ref, version=version)
+    pkg = _select_registry_package(registry.get("packages", []), ref.name, ref.version)
+    pkg_path = Path(pkg["path"])
+    if not pkg_path.exists():
+        raise RuntimeError(f"Registry package missing: {pkg_path}")
+
+    dest = Path(out_dir) if out_dir else Path.cwd() / ref.name
+    if dest.exists():
+        if not force:
+            raise RuntimeError(f"Destination exists: {dest}")
+        shutil.rmtree(dest)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_root = Path(tmpdir)
+        with tarfile.open(pkg_path, "r:gz") as tf:
+            members = tf.getmembers()
+            if not members:
+                raise RuntimeError(f"Empty package: {pkg_path}")
+            top = members[0].name.split("/")[0]
+            tf.extractall(tmp_root)
+        src = tmp_root / top
+        if not src.exists():
+            raise RuntimeError(f"Invalid package layout: {pkg_path}")
+        shutil.copytree(src, dest)
+    return dest

mlxsmith/envs/token_env.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import inspect
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Optional, Protocol
+
+
+@dataclass
+class TokenEnvStep:
+    observation: list[int]
+    reward: float
+    done: bool
+    info: dict[str, Any] = field(default_factory=dict)
+
+
+class TokenEnv(Protocol):
+    def initial_observation(self) -> list[int] | TokenEnvStep:
+        ...
+
+    def step(self, action: int) -> TokenEnvStep:
+        ...
+
+
+@dataclass
+class TokenEnvSpec:
+    factory: Callable[..., TokenEnv]
+    kwargs: dict[str, Any] = field(default_factory=dict)
+    kind: str = "custom"
+
+
+def _filter_kwargs(fn: Callable[..., Any], kwargs: dict[str, Any]) -> dict[str, Any]:
+    try:
+        sig = inspect.signature(fn)
+    except (TypeError, ValueError):
+        return kwargs
+    for param in sig.parameters.values():
+        if param.kind == param.VAR_KEYWORD:
+            return kwargs
+    return {k: v for k, v in kwargs.items() if k in sig.parameters}
+
+
+def _load_from_path(path: Path):
+    spec = importlib.util.spec_from_file_location(path.stem, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load token env module: {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+def _resolve_factory(module, class_name: Optional[str]) -> Callable[..., TokenEnv]:
+    if class_name:
+        factory = getattr(module, class_name, None)
+        if factory is None:
+            raise RuntimeError(f"Token env factory not found: {class_name}")
+        if not callable(factory):
+            raise RuntimeError(f"Token env factory not callable: {class_name}")
+        return factory
+
+    for fallback in ("Env", "TokenEnv", "make_env", "load_env"):
+        factory = getattr(module, fallback, None)
+        if callable(factory):
+            return factory
+
+    raise RuntimeError("Token env factory not found (expected class or make_env/load_env).")
+
+
+def _parse_token_env_spec(project_root: Path, token_env: Any) -> TokenEnvSpec:
+    if isinstance(token_env, str):
+        if token_env in {"tasks", "task_shim"}:
+            return TokenEnvSpec(factory=StringTaskTokenEnv, kind="tasks")
+        path_part, _, class_part = token_env.partition(":")
+        class_name = class_part or None
+        if Path(path_part).suffix == ".py" or Path(path_part).exists():
+            path = Path(path_part)
+            if not path.is_absolute():
+                path = project_root / path
+            module = _load_from_path(path)
+        else:
+            module = importlib.import_module(path_part)
+        factory = _resolve_factory(module, class_name)
+        return TokenEnvSpec(factory=factory, kind="custom")
+
+    if isinstance(token_env, dict):
+        if token_env.get("type") in {"tasks", "task_shim"}:
+            return TokenEnvSpec(factory=StringTaskTokenEnv, kind="tasks")
+        class_name = token_env.get("class") or token_env.get("cls")
+        kwargs = token_env.get("kwargs") or {}
+        if "path" in token_env:
+            path = Path(token_env["path"])
+            if not path.is_absolute():
+                path = project_root / path
+            module = _load_from_path(path)
+        elif "module" in token_env:
+            module = importlib.import_module(str(token_env["module"]))
+        else:
+            raise RuntimeError("token_env requires 'path' or 'module'")
+        factory = _resolve_factory(module, class_name)
+        return TokenEnvSpec(factory=factory, kwargs=dict(kwargs), kind="custom")
+
+    raise RuntimeError("token_env must be a string or mapping")
+
+
+def load_token_env_spec(project_root: Path, env_data: dict) -> Optional[TokenEnvSpec]:
+    token_env = env_data.get("token_env")
+    if not token_env:
+        return None
+    return _parse_token_env_spec(project_root, token_env)
+
+
+def create_token_env(spec: TokenEnvSpec, **kwargs) -> TokenEnv:
+    params = dict(spec.kwargs)
+    params.update(kwargs)
+    params = _filter_kwargs(spec.factory, params)
+    return spec.factory(**params)
+
+
+class StringTaskTokenEnv:
+    def __init__(
+        self,
+        *,
+        prompt: str,
+        tests: str,
+        verifier_fn: Callable[..., Any],
+        workdir: Path,
+        max_steps: int,
+        encode: Callable[[str], list[int]],
+        decode: Callable[[list[int]], str],
+        verifier_kwargs: Optional[dict[str, Any]] = None,
+        eos_token_id: Optional[int] = None,
+    ):
+        self.prompt = prompt
+        self.tests = tests
+        self.verifier_fn = verifier_fn
+        self.workdir = Path(workdir)
+        self.max_steps = max_steps
+        self.encode = encode
+        self.decode = decode
+        self.verifier_kwargs = verifier_kwargs or {}
+        self.eos_token_id = eos_token_id
+        self._prompt_ids: list[int] = []
+        self._generated: list[int] = []
+        self._steps = 0
+
+    def initial_observation(self) -> list[int]:
+        self._prompt_ids = list(self.encode(self.prompt))
+        self._generated = []
+        self._steps = 0
+        tests_dir = self.workdir / "tests"
+        tests_dir.mkdir(parents=True, exist_ok=True)
+        (tests_dir / "test_task.py").write_text(self.tests or "", encoding="utf-8")
+        return list(self._prompt_ids)
+
+    def step(self, action: int) -> TokenEnvStep:
+        if self.eos_token_id is not None and action == self.eos_token_id:
+            done = True
+        else:
+            done = False
+        self._generated.append(int(action))
+        self._steps += 1
+
+        if self._steps >= self.max_steps:
+            done = True
+
+        reward = 0.0
+        info: dict[str, Any] = {}
+        if done:
+            completion_ids = list(self._generated)
+            if self.eos_token_id is not None and completion_ids and completion_ids[-1] == self.eos_token_id:
+                completion_ids = completion_ids[:-1]
+            completion = self.decode(completion_ids)
+            (self.workdir / "main.py").write_text(completion, encoding="utf-8")
+            t0 = time.time()
+            res = self.verifier_fn(self.prompt, completion, str(self.workdir), **self.verifier_kwargs)
+            latency_ms = (time.time() - t0) * 1000.0
+            reward = float(getattr(res, "reward", 0.0))
+            info = dict(getattr(res, "info", {}) or {})
+            info["passed"] = bool(getattr(res, "passed", False))
+            info["verifier_latency_ms"] = latency_ms
+
+        observation = list(self._prompt_ids) + list(self._generated)
+        return TokenEnvStep(
+            observation=observation,
+            reward=reward,
+            done=done,
+            info=info,
+        )

mlxsmith/eval.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+
+from rich.console import Console
+
+from .util import ensure_dir, now_ts
+from .config import ProjectConfig, load_config
+from .models import resolve_model_spec
+from .llm.registry import get_llm_backend
+
+console = Console()
+
+
+def _load_verifier(verifier_path: Path):
+    import importlib.util
+
+    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 run_eval(project_root: Path, suite_path: Path, model_path: Path) -> Path:
+    import yaml
+
+    suite = yaml.safe_load(suite_path.read_text(encoding="utf-8")) or {}
+    out_dir = ensure_dir(project_root / "eval" / "last")
+    out_path = out_dir / "results.json"
+
+    cfg_path = project_root / "mlxsmith.yaml"
+    if cfg_path.exists():
+        cfg = load_config(cfg_path)
+    else:
+        cfg = ProjectConfig()
+    if suite.get("config"):
+        merged = cfg.model_dump()
+        merged.update(suite.get("config") or {})
+        cfg = ProjectConfig.model_validate(merged)
+
+    llm = get_llm_backend(cfg.model.backend)
+    base_model, adapter_path, _meta = resolve_model_spec(project_root, str(model_path), cfg)
+    llm.load(
+        base_model,
+        max_seq_len=cfg.model.max_seq_len,
+        dtype=cfg.model.dtype,
+        trust_remote_code=cfg.model.trust_remote_code,
+    )
+    if adapter_path:
+        llm.apply_adapter(str(adapter_path))
+
+    tasks = suite.get("tasks") or []
+    if not tasks:
+        results = {
+            "model": str(model_path),
+            "suite": suite.get("name", suite_path.name),
+            "error": "no tasks",
+        }
+        out_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
+        return out_path
+
+    summaries = []
+    for task in tasks:
+        prompt = task.get("prompt", "")
+        k = int(task.get("k", 1))
+        verifier_path = task.get("verifier")
+        verify_fn = None
+        if verifier_path:
+            verify_fn = _load_verifier(project_root / verifier_path)
+        passes = 0
+        responses = []
+        t0 = time.time()
+        for i in range(k):
+            gen = llm.generate(
+                prompt,
+                max_new_tokens=int(task.get("max_new_tokens", 256)),
+                temperature=float(task.get("temperature", 0.7)),
+                top_p=float(task.get("top_p", 1.0)),
+                seed=int(task.get("seed", 0)) if task.get("seed") is not None else None,
+            )
+            completion = gen.text[len(prompt) :] if gen.text.startswith(prompt) else gen.text
+            responses.append(completion)
+            if verify_fn:
+                res = verify_fn(prompt, completion, str(out_dir), **(task.get("verifier_kwargs") or {}))
+                if bool(getattr(res, "passed", False)):
+                    passes += 1
+        elapsed = max(time.time() - t0, 1e-6)
+        summaries.append(
+            {
+                "task_id": task.get("id") or prompt[:32],
+                "k": k,
+                "pass@k": passes / max(1, k),
+                "latency_s": elapsed,
+            }
+        )
+
+    results = {
+        "model": str(model_path),
+        "suite": suite.get("name", suite_path.name),
+        "ts": now_ts(),
+        "summary": summaries,
+    }
+    out_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
+    console.print(f"[green]Wrote[/green] {out_path}")
+    return out_path