From 8fb9d5ece37544baebf05d0eb092ae2fd11ad450 Mon Sep 17 00:00:00 2001 From: Avanika Narayan Date: Wed, 6 May 2026 16:20:27 -0700 Subject: [PATCH] feat(mining): add Pearl model conversion workflow (#323) * feat(mining): add Pearl model conversion workflow * test(mining): tolerate missing Docker device request type * docs(mining): record Gemma and Qwen conversion evidence * docs(mining): record Qwen local validation evidence --------- Co-authored-by: krypticmouse --- docs/development/pearl-model-enablement.md | 65 +++- scripts/mining/pearl_model_converter.py | 359 ++++++++++++++++++++ src/openjarvis/cli/mine_cmd.py | 126 +++++-- src/openjarvis/mining/_docker.py | 47 ++- tests/cli/test_mine_cmd.py | 76 +++++ tests/mining/test_docker.py | 63 ++++ tests/scripts/test_pearl_model_converter.py | 116 +++++++ 7 files changed, 811 insertions(+), 41 deletions(-) create mode 100644 scripts/mining/pearl_model_converter.py create mode 100644 tests/scripts/test_pearl_model_converter.py diff --git a/docs/development/pearl-model-enablement.md b/docs/development/pearl-model-enablement.md index b21ff95d..7f4b8c7e 100644 --- a/docs/development/pearl-model-enablement.md +++ b/docs/development/pearl-model-enablement.md @@ -29,7 +29,7 @@ The H100 smoke run validated the default Llama Pearl model end to end through `jarvis mine start`, vLLM `/v1/models`, OpenJarvis inference routing, Pearl gateway template refresh, and `jarvis mine validate-model`. -The Qwen and Gemma targets remain planned: +The Qwen targets and smaller Gemma target remain planned: - `pearl-ai/Qwen3.5-9B-pearl`, `pearl-ai/Qwen3.6-27B-pearl`, and `pearl-ai/Gemma-4-E4B-it-pearl` are not publicly available artifacts yet. @@ -40,6 +40,26 @@ The Qwen and Gemma targets remain planned: injecting metadata from `google/gemma-4-31B-it`; that is not sufficient for OpenJarvis promotion because a clean user install would still fail. +PR #323 added a local staging path for original checkpoint conversion and H100 +runtime validation. Current local evidence: + +- `google/gemma-4-31B-it` converted to + `/tmp/openjarvis-h100/converted/Gemma-4-31B-it-pearl-experimental`. + The local artifact passes `jarvis mine inspect-model`, starts through + `jarvis mine start --local-model-path`, exposes + `pearl-ai/Gemma-4-31B-it-pearl` at `/v1/models`, completes a chat prompt, + and passes `jarvis mine validate-model --allow-planned`. Validation artifact: + `/tmp/openjarvis-h100/converted/Gemma-4-31B-it-pearl-experimental-validate.json`. +- `Qwen/Qwen3.5-9B` converted to + `/tmp/openjarvis-h100/converted/Qwen3.5-9B-pearl-experimental`. + The local artifact passes `jarvis mine inspect-model`, starts Pearl gateway, + resolves `Qwen3_5ForConditionalGeneration`, selects Pearl kernels, loads all + four safetensors shards, exposes `pearl-ai/Qwen3.5-9B-pearl` at `/v1/models`, + completes a chat prompt, and passes `jarvis mine validate-model + --allow-planned`. Required validation flags: `--gdn-prefill-backend triton` + and a 4096-token context. Validation artifact: + `/tmp/openjarvis-h100/converted/Qwen3.5-9B-pearl-experimental-validate.json`. + ## Enablement Checklist 1. Reproduce the current Llama Pearl model recipe. @@ -55,6 +75,49 @@ The Qwen and Gemma targets remain planned: by vLLM's Gemma4 multimodal profiler. - Publish under the planned `pearl-ai/*-pearl` id or a staging namespace. + OpenJarvis includes an experimental local converter for this work: + + ```bash + python scripts/mining/pearl_model_converter.py \ + Qwen/Qwen3.5-9B \ + /tmp/pearl-ai-Qwen3.5-9B-pearl \ + --device cuda + ``` + + The converter copies Hugging Face metadata, emits + `quantization_config.quant_method = "pearl"`, writes a safetensors index, + converts attention q/k/v and MLP down projections to int8 non-mining layers, + and converts the remaining text linear weights to int7 mining layers. Treat + its output as a staging artifact until `jarvis mine inspect-model` and + `jarvis mine validate-model` pass on H100/H200 hardware. + + Local staging artifacts can be inspected before upload: + + ```bash + jarvis mine inspect-model \ + --model /tmp/pearl-ai-Qwen3.5-9B-pearl + ``` + + To run a local staging artifact through the Docker miner, keep `--model` as + the intended served model name and point `--local-model-path` at the + converted checkpoint directory: + + ```bash + jarvis mine init \ + --provider vllm-pearl \ + --wallet-address \ + --model pearl-ai/Qwen3.5-9B-pearl \ + --local-model-path /tmp/pearl-ai-Qwen3.5-9B-pearl \ + --cuda-visible-devices 1 \ + --vllm-arg=--language-model-only \ + --vllm-arg=--skip-mm-profiling \ + --vllm-arg=--gdn-prefill-backend \ + --vllm-arg=triton + jarvis mine start + ``` + + For Qwen3.5 validation, set `max_model_len = 4096` in `[mining.extra]`. + 3. Validate the Pearl vLLM plugin path. - Run `jarvis mine inspect-model --model --allow-planned` before starting the miner. diff --git a/scripts/mining/pearl_model_converter.py b/scripts/mining/pearl_model_converter.py new file mode 100644 index 00000000..a5b5bacc --- /dev/null +++ b/scripts/mining/pearl_model_converter.py @@ -0,0 +1,359 @@ +#!/usr/bin/env python3 +"""Create experimental Pearl-compatible compressed-tensors checkpoints. + +This utility converts raw Hugging Face safetensors checkpoints into the minimal +Pearl quantization shape consumed by Pearl's vLLM plugin: + +- int7 channel-wise weights for mining layers +- int8 channel-wise weights for non-mining layers +- dynamic token-wise symmetric activation metadata + +It is intentionally conservative and aimed at Gemma4/Qwen3.5 enablement work. +It does not upload to Hugging Face and does not mark a model validated. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import load_file, save_file + +NON_MINING_RE = re.compile( + r"(self_attn\.(q_proj|k_proj|v_proj|qkv_proj)|mlp\.down_proj)\.weight$" +) +IGNORED_TEXT_RE = re.compile( + r"(^|\.)(embed_tokens|lm_head|norm|layernorm|layer_norm)\.weight$" +) +IGNORED_MULTIMODAL_RE = re.compile( + r"(^model\.(vision|audio|embed_vision)|vision_tower|vision_model|visual|audio)" +) + + +@dataclass(frozen=True) +class ConversionStats: + copied: int = 0 + mining: int = 0 + non_mining: int = 0 + + def add(self, kind: str) -> "ConversionStats": + return ConversionStats( + copied=self.copied + (kind == "copied"), + mining=self.mining + (kind == "mining"), + non_mining=self.non_mining + (kind == "non_mining"), + ) + + +def classify_weight(name: str, tensor: torch.Tensor) -> str: + """Classify one safetensors entry as copied, mining, or non_mining.""" + + if not name.endswith(".weight") or tensor.ndim != 2: + return "copied" + if IGNORED_TEXT_RE.search(name) or IGNORED_MULTIMODAL_RE.search(name): + return "copied" + if NON_MINING_RE.search(name): + return "non_mining" + return "mining" + + +def quantize_channelwise( + weight: torch.Tensor, + *, + max_val: int, + device: str, + chunk_rows: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Symmetric per-output-channel int quantization for one 2D weight.""" + + if weight.ndim != 2: + raise ValueError(f"expected 2D weight, got shape {tuple(weight.shape)}") + rows = weight.shape[0] + q_chunks: list[torch.Tensor] = [] + scale_chunks: list[torch.Tensor] = [] + for start in range(0, rows, chunk_rows): + chunk = weight[start : start + chunk_rows].to( + device=device, dtype=torch.float32 + ) + scale = chunk.abs().amax(dim=1, keepdim=True) / float(max_val) + scale = torch.where(scale == 0, torch.ones_like(scale), scale) + quantized = torch.round(chunk / scale).clamp(-max_val, max_val).to(torch.int8) + q_chunks.append(quantized.cpu()) + scale_chunks.append(scale.to(dtype=torch.bfloat16).cpu()) + return torch.cat(q_chunks, dim=0), torch.cat(scale_chunks, dim=0) + + +def quantization_config() -> dict[str, Any]: + """Return the Pearl quantization config for generated checkpoints.""" + + dynamic_token = { + "actorder": None, + "block_structure": None, + "dynamic": True, + "group_size": None, + "num_bits": 7, + "observer": None, + "observer_kwargs": {}, + "strategy": "token", + "symmetric": True, + "type": "int", + } + weight_channel = { + "actorder": None, + "block_structure": None, + "dynamic": False, + "group_size": None, + "num_bits": 7, + "observer": "minmax", + "observer_kwargs": {}, + "strategy": "channel", + "symmetric": True, + "type": "int", + } + + int8_input = {**dynamic_token, "num_bits": 8} + int8_weight = {**weight_channel, "num_bits": 8} + + return { + "config_groups": { + "group_0": { + "format": "int-quantized", + "input_activations": int8_input, + "output_activations": None, + "targets": [ + "re:.*self_attn\\.[qkv]_proj$", + "re:.*self_attn\\.qkv_proj$", + "re:.*\\.down_proj$", + ], + "weights": int8_weight, + }, + "group_1": { + "format": "int-quantized", + "input_activations": dynamic_token, + "output_activations": None, + "targets": ["Linear"], + "weights": weight_channel, + }, + }, + "format": "mixed-precision", + "global_compression_ratio": None, + "ignore": [ + "lm_head", + "re:.*embed_tokens$", + "re:.*vision.*", + "re:.*visual.*", + "re:.*vision_tower.*", + "re:.*vision_model.*", + "re:.*image.*", + "re:.*audio.*", + "re:.*embed_vision.*", + "re:.*multi_modal_projector.*", + ], + "kv_cache_scheme": None, + "quant_method": "pearl", + "quantization_status": "compressed", + "sparsity_config": {}, + "transform_config": {}, + "version": "0.13.0", + } + + +def resolve_source(source: str, token: str | None) -> Path: + """Resolve a local path or download a Hugging Face snapshot.""" + + path = Path(source) + if path.exists(): + return path + from huggingface_hub import snapshot_download + + return Path( + snapshot_download( + source, + token=token, + allow_patterns=[ + "*.json", + "*.jinja", + "*.txt", + "*.model", + "*.safetensors", + ".gitattributes", + ], + ) + ) + + +def copy_metadata_files(source_dir: Path, output_dir: Path) -> None: + """Copy non-safetensors model metadata into the output directory.""" + + for src in source_dir.rglob("*"): + if not src.is_file() or src.suffix == ".safetensors": + continue + rel = src.relative_to(source_dir) + dst = output_dir / rel + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) + + +def patch_config(output_dir: Path) -> None: + config_path = output_dir / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"missing {config_path}") + config = json.loads(config_path.read_text()) + config["quantization_config"] = quantization_config() + output_dir.mkdir(parents=True, exist_ok=True) + config_path.write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") + + +def patch_processor_metadata(output_dir: Path) -> None: + """Provide compatibility filenames required by vLLM Gemma4 profiling.""" + + config_path = output_dir / "config.json" + if not config_path.exists(): + return + config = json.loads(config_path.read_text()) + architectures = config.get("architectures") or [] + is_gemma4 = any( + isinstance(item, str) and "Gemma4" in item for item in architectures + ) + if not is_gemma4: + return + + processor = output_dir / "processor_config.json" + preprocessor = output_dir / "preprocessor_config.json" + if processor.exists() and not preprocessor.exists(): + shutil.copy2(processor, preprocessor) + + +def convert_safetensors_file( + src: Path, + dst: Path, + *, + device: str, + chunk_rows: int, + dry_run: bool, +) -> tuple[int, dict[str, str], ConversionStats]: + tensors = load_file(str(src), device="cpu") + out: dict[str, torch.Tensor] = {} + stats = ConversionStats() + for name, tensor in tensors.items(): + kind = classify_weight(name, tensor) + stats = stats.add(kind) + if dry_run or kind == "copied": + out[name] = tensor + continue + max_val = 63 if kind == "mining" else 127 + quantized, scale = quantize_channelwise( + tensor, + max_val=max_val, + device=device, + chunk_rows=chunk_rows, + ) + out[name] = quantized + out[f"{name.removesuffix('.weight')}.weight_scale"] = scale + + if not dry_run: + dst.parent.mkdir(parents=True, exist_ok=True) + save_file(out, str(dst)) + total_size = sum(tensor.nbytes for tensor in out.values()) + weight_map = {key: dst.name for key in out} + return total_size, weight_map, stats + + +def write_index(output_dir: Path, total_size: int, weight_map: dict[str, str]) -> None: + index = { + "metadata": {"total_size": total_size}, + "weight_map": dict(sorted(weight_map.items())), + } + (output_dir / "model.safetensors.index.json").write_text( + json.dumps(index, indent=2, sort_keys=True) + "\n" + ) + + +def convert_checkpoint( + source: str, + output_dir: Path, + *, + hf_token_env: str, + device: str, + chunk_rows: int, + dry_run: bool, +) -> ConversionStats: + token = os.environ.get(hf_token_env) + source_dir = resolve_source(source, token) + output_dir.mkdir(parents=True, exist_ok=True) + + safetensors_files = sorted(source_dir.glob("*.safetensors")) + if not safetensors_files: + raise FileNotFoundError(f"no .safetensors files found in {source_dir}") + + if not dry_run: + copy_metadata_files(source_dir, output_dir) + patch_config(output_dir) + patch_processor_metadata(output_dir) + + total_size = 0 + weight_map: dict[str, str] = {} + stats = ConversionStats() + for src in safetensors_files: + dst = output_dir / src.name + file_size, file_map, file_stats = convert_safetensors_file( + src, + dst, + device=device, + chunk_rows=chunk_rows, + dry_run=dry_run, + ) + total_size += file_size + weight_map.update(file_map) + stats = ConversionStats( + copied=stats.copied + file_stats.copied, + mining=stats.mining + file_stats.mining, + non_mining=stats.non_mining + file_stats.non_mining, + ) + print( + f"{src.name}: copied={file_stats.copied} " + f"mining={file_stats.mining} non_mining={file_stats.non_mining}" + ) + + if not dry_run: + write_index(output_dir, total_size, weight_map) + return stats + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source", help="Local model directory or Hugging Face model id") + parser.add_argument("output_dir", type=Path) + parser.add_argument("--hf-token-env", default="HF_TOKEN") + parser.add_argument("--device", default="cpu", help="cpu or cuda") + parser.add_argument("--chunk-rows", type=int, default=4096) + parser.add_argument("--dry-run", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + stats = convert_checkpoint( + args.source, + args.output_dir, + hf_token_env=args.hf_token_env, + device=args.device, + chunk_rows=args.chunk_rows, + dry_run=args.dry_run, + ) + print( + "total: " + f"copied={stats.copied} mining={stats.mining} non_mining={stats.non_mining}" + ) + if args.dry_run: + print("dry run only; no checkpoint was written") + + +if __name__ == "__main__": + main() diff --git a/src/openjarvis/cli/mine_cmd.py b/src/openjarvis/cli/mine_cmd.py index f5e1c53a..127700cd 100644 --- a/src/openjarvis/cli/mine_cmd.py +++ b/src/openjarvis/cli/mine_cmd.py @@ -166,7 +166,7 @@ def _artifact_has_weights(paths: set[str]) -> bool: def _artifact_quant_method(model_payload: Any) -> str: if not isinstance(model_payload, dict): return "" - config = model_payload.get("config") + config = model_payload.get("config", model_payload) if not isinstance(config, dict): return "" quant_config = config.get("quantization_config") @@ -178,7 +178,7 @@ def _artifact_quant_method(model_payload: Any) -> str: def _artifact_architectures(model_payload: Any) -> list[str]: if not isinstance(model_payload, dict): return [] - config = model_payload.get("config") + config = model_payload.get("config", model_payload) if not isinstance(config, dict): return [] architectures = config.get("architectures") @@ -187,6 +187,19 @@ def _artifact_architectures(model_payload: Any) -> list[str]: return [str(item) for item in architectures if isinstance(item, str)] +def _local_artifact_payload_and_paths(path: Path) -> tuple[dict[str, Any], set[str]]: + config_path = path / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"missing {config_path}") + payload = json.loads(config_path.read_text()) + paths = { + str(item.relative_to(path)) + for item in path.rglob("*") + if item.is_file() and ".git" not in item.parts + } + return payload, paths + + def _pid_alive(pid: int | None) -> bool: if not pid: return False @@ -359,6 +372,18 @@ def doctor() -> None: prompt="env var holding pearld password", ) @click.option("--model", default=DEFAULT_PEARL_MODEL) +@click.option( + "--local-model-path", + type=click.Path(file_okay=False, path_type=Path), + default=None, + help="Local converted Pearl checkpoint directory to mount into the container.", +) +@click.option( + "--vllm-arg", + "vllm_args", + multiple=True, + help="Extra argument to pass through to `vllm serve`; repeat as needed.", +) @click.option("--image", default=PEARL_IMAGE_TAG) @click.option( "--cuda-visible-devices", @@ -380,6 +405,8 @@ def init( pearld_user: str, pearld_password_env: str, model: str, + local_model_path: Path | None, + vllm_args: tuple[str, ...], image: str, cuda_visible_devices: str, gateway_host: str, @@ -398,16 +425,21 @@ def init( if selected_provider == "vllm-pearl": hw = _detect_hardware() - cap = detect_for_engine_model( - hw=hw, - engine_id="vllm", - model=model, - provider_id="vllm-pearl", - ) - if not cap.supported: + if local_model_path is None: + cap = detect_for_engine_model( + hw=hw, + engine_id="vllm", + model=model, + provider_id="vllm-pearl", + ) + if not cap.supported: + raise click.ClickException( + f"vllm-pearl not supported on this host: {cap.reason}\n" + "See `jarvis mine models` and `jarvis mine doctor` for details." + ) + elif not local_model_path.exists(): raise click.ClickException( - f"vllm-pearl not supported on this host: {cap.reason}\n" - "See `jarvis mine models` and `jarvis mine doctor` for details." + f"local Pearl model path does not exist: {local_model_path}" ) model_spec = get_pearl_model_spec(model) max_model_len = ( @@ -463,6 +495,12 @@ pearld_rpc_user = "{pearld_user}" pearld_rpc_password_env = "{pearld_password_env}" hf_token_env = "HF_TOKEN" """ + if local_model_path is not None: + path = local_model_path.expanduser().resolve() + section += f'local_model_path = "{path}"\n' + if vllm_args: + encoded_args = ", ".join(json.dumps(arg) for arg in vllm_args) + section += f"vllm_args = [{encoded_args}]\n" if selected_provider != "vllm-pearl": section = f""" [mining] @@ -801,13 +839,17 @@ def inspect_model( click.echo(f"model: {model}") results: list[bool] = [] records: list[dict[str, Any]] = [] + local_path = Path(model) + is_local = local_path.exists() def record(name: str, ok: bool, info: str) -> None: records.append({"name": name, "ok": ok, "info": info}) _validation_row(results, name, ok, info) - spec = get_pearl_model_spec(model) - if spec is None: + spec = None if is_local else get_pearl_model_spec(model) + if is_local: + record("Model registry", True, "local artifact (not registered)") + elif spec is None: planned = pearl_variant_for_base_model(model) info = f"raw model; planned Pearl variant {planned}" if planned else "unknown" record("Model registry", False, info) @@ -819,20 +861,27 @@ def inspect_model( token = os.environ.get(hf_token_env, "") model_payload: Any = None paths: set[str] = set() - try: - model_payload = _hf_json( - _hf_model_path(model, ""), - token=token, - timeout=timeout, - ) - record("HF model", True, "accessible") - except urllib.error.HTTPError as exc: - info = f"HTTP {exc.code}" - if exc.code in {401, 403} and not token: - info += f"; set ${hf_token_env} for gated/private models" - record("HF model", False, info) - except Exception as exc: # noqa: BLE001 - record("HF model", False, str(exc).splitlines()[0]) + if is_local: + try: + model_payload, paths = _local_artifact_payload_and_paths(local_path) + record("Local artifact", True, str(local_path)) + except Exception as exc: # noqa: BLE001 + record("Local artifact", False, str(exc).splitlines()[0]) + else: + try: + model_payload = _hf_json( + _hf_model_path(model, ""), + token=token, + timeout=timeout, + ) + record("HF model", True, "accessible") + except urllib.error.HTTPError as exc: + info = f"HTTP {exc.code}" + if exc.code in {401, 403} and not token: + info += f"; set ${hf_token_env} for gated/private models" + record("HF model", False, info) + except Exception as exc: # noqa: BLE001 + record("HF model", False, str(exc).splitlines()[0]) if model_payload is not None: quant_method = _artifact_quant_method(model_payload) @@ -842,16 +891,19 @@ def inspect_model( quant_method or "missing quantization_config.quant_method", ) - try: - tree_payload = _hf_json( - _hf_model_path(model, "/tree/main?recursive=false"), - token=token, - timeout=timeout, - ) - paths = _artifact_file_paths(tree_payload) - record("HF file list", bool(paths), f"{len(paths)} files") - except Exception as exc: # noqa: BLE001 - record("HF file list", False, str(exc).splitlines()[0]) + if is_local: + record("File list", bool(paths), f"{len(paths)} files") + else: + try: + tree_payload = _hf_json( + _hf_model_path(model, "/tree/main?recursive=false"), + token=token, + timeout=timeout, + ) + paths = _artifact_file_paths(tree_payload) + record("HF file list", bool(paths), f"{len(paths)} files") + except Exception as exc: # noqa: BLE001 + record("HF file list", False, str(exc).splitlines()[0]) if paths: has_config = "config.json" in paths diff --git a/src/openjarvis/mining/_docker.py b/src/openjarvis/mining/_docker.py index e288a997..ff91f02a 100644 --- a/src/openjarvis/mining/_docker.py +++ b/src/openjarvis/mining/_docker.py @@ -9,6 +9,7 @@ from __future__ import annotations import os import re +import shlex import subprocess from pathlib import Path from typing import TYPE_CHECKING, Any @@ -24,6 +25,7 @@ from openjarvis.mining._constants import ( ) CONTAINER_NAME = "openjarvis-pearl-miner" +LOCAL_MODEL_BIND_PATH = "/models/openjarvis-local-pearl-model" _SECRET_LOG_PATTERNS = ( (re.compile(r"(rpc_password:\s*)\S+", re.IGNORECASE), r"\1[REDACTED]"), @@ -265,13 +267,29 @@ class PearlDockerLauncher: hf_token_env = extra.get("hf_token_env", "HF_TOKEN") hf_token = os.environ.get(hf_token_env, "") - model = extra.get("model", "pearl-ai/Llama-3.3-70B-Instruct-pearl") + model = str(extra.get("model", "pearl-ai/Llama-3.3-70B-Instruct-pearl")) + local_model_path = extra.get("local_model_path") vllm_port = int(extra.get("vllm_port", 8000)) gpu_mem = float(extra.get("gpu_memory_utilization", 0.96)) max_len = int(extra.get("max_model_len", 8192)) + model_arg = model + local_model_host_path: Path | None = None + if local_model_path: + local_model_host_path = Path(str(local_model_path)).expanduser().resolve() + if not local_model_host_path.exists(): + raise ConfigurationError( + f"local Pearl model path does not exist: {local_model_host_path}" + ) + if not local_model_host_path.is_dir(): + raise ConfigurationError( + "local Pearl model path must be a directory: " + f"{local_model_host_path}" + ) + model_arg = LOCAL_MODEL_BIND_PATH + command = [ - model, + model_arg, "--host", "0.0.0.0", "--port", @@ -282,6 +300,9 @@ class PearlDockerLauncher: "--max-model-len", str(max_len), ] + if local_model_host_path is not None: + command.extend(["--served-model-name", model]) + command.extend(_extra_vllm_args(extra.get("vllm_args", []))) environment = { "PEARLD_RPC_URL": extra.get("pearld_rpc_url", "http://localhost:44107"), @@ -296,7 +317,12 @@ class PearlDockerLauncher: cuda_devices = str(extra.get("cuda_visible_devices", "")).strip() device_ids = [part.strip() for part in cuda_devices.split(",") if part.strip()] if device_ids: - environment["CUDA_VISIBLE_DEVICES"] = ",".join(device_ids) + # Docker's device request exposes the selected host GPUs inside + # the container as a compact 0..N-1 list. Keep host IDs for the + # NVIDIA runtime and use container-local IDs for CUDA/vLLM. + environment["CUDA_VISIBLE_DEVICES"] = ",".join( + str(idx) for idx, _ in enumerate(device_ids) + ) environment["NVIDIA_VISIBLE_DEVICES"] = ",".join(device_ids) # Dynamic import so tests don't need the real `docker` package shape. @@ -316,6 +342,11 @@ class PearlDockerLauncher: volumes = { str(hf_cache): {"bind": "/root/.cache/huggingface", "mode": "rw"}, } + if local_model_host_path is not None: + volumes[str(local_model_host_path)] = { + "bind": LOCAL_MODEL_BIND_PATH, + "mode": "ro", + } # Sanitize wrap so an APIError from the daemon doesn't surface the # full container spec (which contains the resolved password) in a @@ -418,3 +449,13 @@ def _redact_container_logs(text: str) -> str: for pattern, replacement in _SECRET_LOG_PATTERNS: redacted = pattern.sub(replacement, redacted) return redacted + + +def _extra_vllm_args(value: Any) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + return shlex.split(value) + if isinstance(value, (list, tuple)): + return [str(item) for item in value] + raise ConfigurationError("[mining.extra].vllm_args must be a string or list") diff --git a/tests/cli/test_mine_cmd.py b/tests/cli/test_mine_cmd.py index 5f7c56d9..44adf7ac 100644 --- a/tests/cli/test_mine_cmd.py +++ b/tests/cli/test_mine_cmd.py @@ -89,6 +89,30 @@ def test_mine_inspect_model_gemma4_requires_processor_metadata(monkeypatch) -> N assert "missing preprocessor_config.json" in result.output +def test_mine_inspect_model_accepts_local_artifact(tmp_path: Path) -> None: + artifact = tmp_path / "qwen-pearl" + artifact.mkdir() + (artifact / "config.json").write_text( + json.dumps( + { + "architectures": ["Qwen3_5ForConditionalGeneration"], + "quantization_config": {"quant_method": "pearl"}, + } + ) + ) + (artifact / "tokenizer.json").write_text("{}") + (artifact / "model.safetensors").write_text("placeholder") + + result = CliRunner().invoke( + cli, + ["mine", "inspect-model", "--model", str(artifact)], + ) + + assert result.exit_code == 0 + assert "local artifact" in result.output + assert "Artifact inspection passed" in result.output + + def test_mine_init_writes_mining_config(tmp_path: Path) -> None: config_path = tmp_path / "config.toml" @@ -170,6 +194,58 @@ def test_mine_init_writes_cuda_visible_devices_for_vllm( assert 'cuda_visible_devices = "0,1"' in content +def test_mine_init_writes_local_model_path_for_vllm( + tmp_path: Path, +) -> None: + config_path = tmp_path / "config.toml" + local_model = tmp_path / "converted-qwen" + local_model.mkdir() + + with ( + patch("openjarvis.cli.mine_cmd._detect_hardware"), + patch( + "openjarvis.cli.mine_cmd.check_docker_available", + return_value=(True, ""), + ), + patch("openjarvis.cli.mine_cmd.check_disk_free", return_value=(True, "")), + patch("openjarvis.cli.mine_cmd._docker_from_env", return_value=MagicMock()), + patch( + "openjarvis.cli.mine_cmd.PearlDockerLauncher.ensure_image", + return_value="openjarvis/pearl-miner:master", + ), + ): + result = CliRunner().invoke( + cli, + [ + "mine", + "init", + "--provider", + "vllm-pearl", + "--wallet-address", + "prl1qtestingaddr", + "--pearld-rpc-url", + "http://127.0.0.1:44107", + "--pearld-rpc-user", + "rpcuser", + "--pearld-rpc-password-env", + "TEST_PEARLD_PASSWORD", + "--model", + "pearl-ai/Qwen3.5-9B-pearl", + "--local-model-path", + str(local_model), + "--vllm-arg=--language-model-only", + "--vllm-arg=--skip-mm-profiling", + ], + env={"OPENJARVIS_CONFIG": str(config_path)}, + ) + + assert result.exit_code == 0 + content = config_path.read_text() + assert 'model = "pearl-ai/Qwen3.5-9B-pearl"' in content + assert f'local_model_path = "{local_model.resolve()}"' in content + assert 'vllm_args = ["--language-model-only", "--skip-mm-profiling"]' in content + + def test_mine_start_uses_configured_provider(tmp_path: Path, monkeypatch) -> None: config_path = tmp_path / "config.toml" config_path.write_text( diff --git a/tests/mining/test_docker.py b/tests/mining/test_docker.py index 69e0dc77..7c7c3323 100644 --- a/tests/mining/test_docker.py +++ b/tests/mining/test_docker.py @@ -190,6 +190,69 @@ def test_launcher_start_calls_run_with_expected_kwargs(_env_password): assert kwargs["device_requests"][0].device_ids == ["0", "1"] +def test_launcher_start_maps_host_gpu_ids_to_container_local_cuda_ids( + _env_password, +) -> None: + from openjarvis.mining._docker import PearlDockerLauncher + from openjarvis.mining._stubs import MiningConfig, SoloTarget + + fake = MagicMock() + fake.containers.run.return_value = MagicMock(id="cid-1", status="running") + launcher = PearlDockerLauncher(client=fake) + cfg = MiningConfig( + provider="vllm-pearl", + wallet_address="prl1qaaa", + submit_target=SoloTarget(pearld_rpc_url="http://localhost:44107"), + extra={ + "pearld_rpc_password_env": "PEARLD_RPC_PASSWORD", + "cuda_visible_devices": "1", + }, + ) + + launcher.start(cfg, image="openjarvis/pearl-miner:main") + + kwargs = fake.containers.run.call_args.kwargs + assert kwargs["environment"]["NVIDIA_VISIBLE_DEVICES"] == "1" + assert kwargs["environment"]["CUDA_VISIBLE_DEVICES"] == "0" + if kwargs["device_requests"] is not None: + assert kwargs["device_requests"][0].device_ids == ["1"] + + +def test_launcher_start_mounts_local_model_path(_env_password, tmp_path): + from openjarvis.mining._docker import LOCAL_MODEL_BIND_PATH, PearlDockerLauncher + from openjarvis.mining._stubs import MiningConfig, SoloTarget + + local_model = tmp_path / "qwen-pearl" + local_model.mkdir() + fake = MagicMock() + fake.containers.run.return_value = MagicMock(id="cid-1", status="running") + launcher = PearlDockerLauncher(client=fake) + cfg = MiningConfig( + provider="vllm-pearl", + wallet_address="prl1qaaa", + submit_target=SoloTarget(pearld_rpc_url="http://localhost:44107"), + extra={ + "model": "pearl-ai/Qwen3.5-9B-pearl", + "local_model_path": str(local_model), + "pearld_rpc_password_env": "PEARLD_RPC_PASSWORD", + "vllm_args": ["--language-model-only", "--skip-mm-profiling"], + }, + ) + + launcher.start(cfg, image="openjarvis/pearl-miner:main") + + kwargs = fake.containers.run.call_args.kwargs + assert kwargs["command"][0] == LOCAL_MODEL_BIND_PATH + assert "--served-model-name" in kwargs["command"] + assert "pearl-ai/Qwen3.5-9B-pearl" in kwargs["command"] + assert "--language-model-only" in kwargs["command"] + assert "--skip-mm-profiling" in kwargs["command"] + assert kwargs["volumes"][str(local_model.resolve())] == { + "bind": LOCAL_MODEL_BIND_PATH, + "mode": "ro", + } + + def test_launcher_stop_calls_container_stop_and_remove(): from openjarvis.mining._docker import PearlDockerLauncher diff --git a/tests/scripts/test_pearl_model_converter.py b/tests/scripts/test_pearl_model_converter.py new file mode 100644 index 00000000..6eff4cb3 --- /dev/null +++ b/tests/scripts/test_pearl_model_converter.py @@ -0,0 +1,116 @@ +"""Tests for the experimental Pearl model converter script.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +safetensors_torch = pytest.importorskip("safetensors.torch") +load_file = safetensors_torch.load_file +save_file = safetensors_torch.save_file + + +def _load_converter(): + path = Path(__file__).parents[2] / "scripts" / "mining" / "pearl_model_converter.py" + spec = importlib.util.spec_from_file_location("pearl_model_converter", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_converter_quantizes_linear_weights_and_writes_pearl_config( + tmp_path: Path, +) -> None: + converter = _load_converter() + src = tmp_path / "src" + out = tmp_path / "out" + src.mkdir() + (src / "config.json").write_text( + json.dumps({"architectures": ["TinyForCausalLM"], "model_type": "tiny"}) + ) + (src / "tokenizer_config.json").write_text("{}") + save_file( + { + "model.embed_tokens.weight": torch.ones(4, 4, dtype=torch.bfloat16), + "model.layers.0.self_attn.q_proj.weight": torch.arange( + 16, dtype=torch.float32 + ).reshape(4, 4), + "model.layers.0.self_attn.o_proj.weight": torch.arange( + 16, dtype=torch.float32 + ).reshape(4, 4), + "model.layers.0.input_layernorm.weight": torch.ones(4), + }, + src / "model.safetensors", + ) + + stats = converter.convert_checkpoint( + str(src), + out, + hf_token_env="HF_TOKEN", + device="cpu", + chunk_rows=2, + dry_run=False, + ) + + assert stats.non_mining == 1 + assert stats.mining == 1 + tensors = load_file(out / "model.safetensors") + assert tensors["model.layers.0.self_attn.q_proj.weight"].dtype == torch.int8 + assert tensors["model.layers.0.self_attn.o_proj.weight"].dtype == torch.int8 + assert tensors["model.layers.0.self_attn.q_proj.weight_scale"].shape == (4, 1) + assert tensors["model.layers.0.self_attn.o_proj.weight_scale"].shape == (4, 1) + assert tensors["model.embed_tokens.weight"].dtype == torch.bfloat16 + + config = json.loads((out / "config.json").read_text()) + assert config["quantization_config"]["quant_method"] == "pearl" + assert ( + config["quantization_config"]["config_groups"]["group_1"]["weights"]["num_bits"] + == 7 + ) + assert "re:.*visual.*" in config["quantization_config"]["ignore"] + index = json.loads((out / "model.safetensors.index.json").read_text()) + assert "model.layers.0.self_attn.o_proj.weight_scale" in index["weight_map"] + + +def test_converter_adds_gemma4_preprocessor_compat_file( + tmp_path: Path, +) -> None: + converter = _load_converter() + src = tmp_path / "src" + out = tmp_path / "out" + src.mkdir() + (src / "config.json").write_text( + json.dumps( + { + "architectures": ["Gemma4ForConditionalGeneration"], + "model_type": "gemma4", + } + ) + ) + (src / "processor_config.json").write_text('{"processor_class":"Gemma4Processor"}') + (src / "tokenizer_config.json").write_text("{}") + save_file( + { + "language_model.model.layers.0.self_attn.q_proj.weight": torch.arange( + 16, dtype=torch.float32 + ).reshape(4, 4), + }, + src / "model.safetensors", + ) + + converter.convert_checkpoint( + str(src), + out, + hf_token_env="HF_TOKEN", + device="cpu", + chunk_rows=2, + dry_run=False, + ) + + assert (out / "processor_config.json").exists() + assert (out / "preprocessor_config.json").exists()