diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 07493c4..c547780 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -35,6 +35,18 @@ jobs: - name: Type check with MyPy run: mypy src/rm_abstract --ignore-missing-imports + continue-on-error: true # existing type errors are being fixed module by module + + # Internal addresses, credentials and keys must never reach this public repository + sensitive: + name: Sensitive-data scan + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - run: python tools/check_sensitive.py . # Unit tests (multiple Python versions) test: @@ -64,7 +76,10 @@ jobs: run: pytest tests/unit -v --tb=short - name: ai-rm conformance (rm-core protocol v0.1) - run: python -m rm_core.conformance airm.accelerator_providers rm-abstract + run: python -m rm_core.conformance + + - name: Node daemon dry-run without devices + run: python -m rm_abstract.airm.daemon --provider sim --once --dry-run - name: Run tests with coverage run: pytest tests/ -v --cov=rm_abstract --cov-report=xml --cov-report=term diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..435f04d --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,33 @@ +# Changelog + +## 0.4.0 — 2026-10-19 + +Repositioned as the **L1 accelerator abstraction** of [ai-rm](https://github.com/ai-computing/ai-rm). + +### Added +- ai-rm protocol v0.1 adapter `rm_abstract.airm`: `AcceleratorProvider` `rm-abstract` + (entry point `airm.accelerator_providers`), discovery through read-only vendor tools with + `HardwareDetector` as fallback, datasheet `peak_tflops` for RNGD and ATOM. +- Node daemon `python -m rm_abstract.airm.daemon`: `AcceleratorProfile` per device, node labels + `airm.ai-computing.io/.`, `airm_accelerators` on `/metrics`. +- Helm chart `deploy/helm/rm-abstract-node` (DaemonSet, RBAC). +- `ai-rm.yaml` component manifest; dependency on rm-core v0.1.1. +- `tools/check_sensitive.py` and a CI job that scans every push. +- README in English with a Korean translation (`README.ko.md`). + +### Changed +- Package metadata: description, keywords, authors and URLs now point to `ai-computing`. +- Ruff configuration moved to `[tool.ruff.lint]`; code formatted with ruff and black. +- CI runs the full `rm_core.conformance` and a daemon dry-run; mypy is reported but non-blocking + while existing type errors are fixed. +- `test_all_examples.sh` → `scripts/run_examples.sh` (no hard-coded path); + `test_multi_engine.py` → `scripts/manual/multi_engine_check.py`. + +### Moved +- `CLEANUP_SUMMARY.md`, `docs/FINAL_TEST_REPORT.md`, `docs/TEST_RESULTS.md` and + `docs/PRESENTATION.md` → `docs/archive/` (point-in-time, not maintained). + +## 0.3.0 + +vLLM-compatible `LLM` API on GPU and CPU; vLLM, Triton and TorchServe engines with automatic server +management; plugin registry for resources and engines. diff --git a/README.ko.md b/README.ko.md new file mode 100644 index 0000000..fcf8e4d --- /dev/null +++ b/README.ko.md @@ -0,0 +1,243 @@ +
+ +

rm_abstract_layer

+ +

이종 AI 인프라를 위한 가속기 추상화

+ +

+ English · 한국어 +

+ +

+ CI + ai-rm L1 계층 + 프로토콜 v0.1 + Python 3.9+ + 라이선스: BSD-3-Clause +

+ +

+ 빠른 시작 · + ai-rm 연동 · + vLLM 호환 API · + 구조 · + 문서 +

+ +
+ +--- + +**rm_abstract_layer**는 [ai-rm](https://github.com/ai-computing/ai-rm)의 **L1 가속기 추상화** 구성요소입니다. 각 노드의 +GPU와 NPU(FuriosaAI RNGD, Rebellions ATOM)를 탐지해 [rm-core](https://github.com/ai-computing/rm-core)의 +`Accelerator` 모델로 표현하고, Kubernetes에 `AcceleratorProfile` 객체와 노드 레이블로 게시합니다. 배치(llm-d EPP)와 +관측(lmd-top)이 이 정보를 읽습니다. + +같은 코드를 GPU나 CPU에서 실행하고 서빙 엔진(vLLM, Triton, TorchServe)을 바꿔 쓸 수 있는 **vLLM 호환 Python API**도 +제공합니다. + + + + + + + +
+🔎 장치 탐지
+읽기 전용 벤더 도구(nvidia-smi, furiosa-smi, rbln-stat)를 먼저 쓰고, 없으면 SDK +기반으로 탐지합니다. 장치 자체 번호, 모델, 메모리, 데이터시트 최대 TFLOPS. +
+☸️ Kubernetes 게시
+노드 데몬(DaemonSet)이 장치마다 AcceleratorProfile을 apply하고, 노드에 +airm.ai-computing.io/<vendor>.<model> 레이블을 달고 airm_accelerators를 냅니다. +
+🔁 하나의 API, 여러 백엔드
+from vllm import LLM을 from rm_abstract import LLM으로 바꾸고, 장치와 엔진은 인자나 환경 +변수로 고릅니다. +
+ +## 🚀 빠른 시작 + +```bash +git clone https://github.com/ai-computing/rm_abstract_layer.git && cd rm_abstract_layer +pip install -r requirements/requirements.txt +pip install -e . + +python -m rm_abstract.airm.daemon --once --dry-run # 이 노드의 장치를 k8s 객체로 출력 +python -m rm_abstract.airm.daemon --provider sim --once --dry-run # 장치 없이 시연 +python -m rm_core.conformance airm.accelerator_providers rm-abstract # 계약 검사 (PASS / SKIP / FAIL) +``` + +ai-rm으로 설치할 때는 [ai-rm](https://github.com/ai-computing/ai-rm) 저장소에서 `make setup COMPONENTS=rm_abstract_layer`를 +실행합니다(절차는 [`ai-rm.yaml`](ai-rm.yaml)에 있습니다). + +## ☸️ ai-rm 장치 탐지 + +### `AcceleratorProvider` 어댑터 + +`rm_abstract.airm`은 rm-core v0.1 `AcceleratorProvider` 계약을 구현하고 `airm.accelerator_providers` entry-point +그룹에 `rm-abstract`로 등록됩니다. + +```python +from rm_core import registry + +provider = registry.load("airm.accelerator_providers", "rm-abstract") +for acc in provider.discover(): + print(acc.id, acc.vendor, acc.model, acc.memory_bytes, acc.peak_tflops) +# node-a/furiosa:0 furiosa RNGD ... {'bf16': 256.0, 'fp8': 512.0} +``` + +| 필드 | 값 | +|---|---| +| `id` | `<노드>/:`, index는 장치 자체 번호 (`npu3` → `furiosa:3`) | +| 노드 이름 | `AIRM_NODE_NAME` (DaemonSet에서 `spec.nodeName`으로 설정), 없으면 hostname | +| `peak_tflops` | 벤더 데이터시트에 값이 있는 모델(RNGD, ATOM)만, `peak_source`에 출처 기록. 그 외는 미상 | +| CPU | 보고하지 않음. 가속기가 없는 노드는 `[]` | + +### 노드 데몬 + +`python -m rm_abstract.airm.daemon`은 주기적으로 탐지해 + +- 장치마다 `AcceleratorProfile`(`airm.ai-computing.io/v1alpha1`)을 apply하고, `AIRM_REMOVE_AFTER`회 연속으로 보이지 + 않은 장치의 프로파일은 삭제합니다. +- 노드 레이블 `airm.ai-computing.io/.=<개수>`, `airm.ai-computing.io/managed=true`를 패치합니다. +- `/metrics`에 `airm_accelerators{vendor,model}`을 냅니다. + +| 변수 | 기본값 | 의미 | +|---|---|---| +| `AIRM_NODE_NAME` | hostname | id와 레이블에 쓰는 노드 이름 | +| `AIRM_PROVIDER` | `rm-abstract` | provider entry point (시연은 `sim`) | +| `AIRM_INTERVAL_S` | `60` | 탐지 주기 | +| `AIRM_METRICS_PORT` | `9400` | `/metrics` 포트, `0`이면 끔 | +| `AIRM_REMOVE_AFTER` | `3` | 프로파일 삭제 전 연속 미탐지 횟수 (도구 일시 실패로 깜박이지 않도록) | + +클러스터에 apply하려면 Kubernetes 클라이언트가 필요합니다: `pip install -e ".[k8s]"`. + +### Helm chart + +[`deploy/helm/rm-abstract-node`](deploy/helm/rm-abstract-node)는 데몬을 DaemonSet으로 설치하고 ServiceAccount와 +ClusterRole(`acceleratorprofiles` CRUD, `nodes` get/patch)을 만듭니다. `AcceleratorProfile` CRD는 +[rm-core](https://github.com/ai-computing/rm-core/tree/v0.1.1/k8s)에 있고 llm-d-setup / ai-rm umbrella가 설치합니다. +이미지는 사이트별로 `image.repository` / `image.tag`로 지정합니다. + +```bash +helm template rm-abstract-node deploy/helm/rm-abstract-node --set image.repository=/rm-abstract-node +``` + +## 🔁 vLLM 호환 API + +```python +# from vllm import LLM, SamplingParams +from rm_abstract import LLM, SamplingParams + +llm = LLM(model="gpt2") # 자동: GPU > CPU +llm = LLM(model="gpt2", device="cpu") # 또는 RM_DEVICE=cpu +llm = LLM(model="gpt2", engine="triton") # 또는 RM_ENGINE=triton +outputs = llm.generate(["Hello, I am"], SamplingParams(max_tokens=50)) +print(outputs[0].outputs[0].text) +``` + +
+실행 중 제어, 엔진, 오류 처리 +
+ +```python +import rm_abstract + +rm_abstract.init(device="auto") +print(rm_abstract.get_available_resources()) # {'nvidia': True, 'cpu': True} +rm_abstract.switch_device("cpu") +``` + +| 엔진 | 자동으로 해 주는 일 | +|---|---| +| vLLM (기본) | — | +| Triton | HuggingFace → ONNX 변환, 모델 저장소 구성, Docker 컨테이너 시작·정리 | +| TorchServe | MAR 패키징, 서버 시작, 모델 등록, 정리 | + +```python +from rm_abstract.exceptions import BackendNotAvailableError + +try: + llm = LLM(model="gpt2", device="nvidia") +except BackendNotAvailableError: + llm = LLM(model="gpt2", device="cpu") +``` + +설치 extras: `.[gpu]`, `.[triton]`, `.[torchserve]`, `.[serving]`, `.[all]` ([requirements 안내](docs/README_REQUIREMENTS.md)). +자세한 내용은 [빠른 시작](docs/QUICKSTART.md), [서빙 엔진](docs/SERVING_ENGINES.md)을 참고하세요. +
+ +### 지원 현황 + +| 장치 | 탐지·`AcceleratorProfile` | `LLM` 추론 | +|---|:-:|:-:| +| NVIDIA GPU | ✅ | ✅ vLLM · Triton · TorchServe | +| CPU | — *(가속기로 보고하지 않음)* | ✅ transformers · TorchServe | +| FuriosaAI RNGD | ✅ | 🧪 플러그인 있음, 미검증 | +| Rebellions ATOM | ✅ | 🧪 플러그인 있음, 미검증 | + +## 🏗 구조 + +```mermaid +flowchart LR + subgraph node["노드"] + tools["nvidia-smi · furiosa-smi · rbln-stat"] --> prov["rm_abstract.airm
AcceleratorProvider"] + det["HardwareDetector
(SDK 대체 경로)"] --> prov + prov --> d["노드 데몬"] + end + d -- "AcceleratorProfile · 노드 레이블" --> k8s[("Kubernetes API")] + d -. "airm_accelerators" .-> prom[("Prometheus")] + k8s --> epp["llm-d EPP / 배치"] + k8s --> top["lmd-top"] +``` + +```text +src/rm_abstract/ +├── airm/ # ai-rm 프로토콜 v0.1: AcceleratorProvider, 노드 데몬 +├── resources/ # 장치 플러그인: gpu/nvidia, npu/furiosa, npu/rebellions, cpu +├── engines/ # 서빙 엔진: vLLM, Triton, TorchServe, 모델 변환 +├── common/ # 하드웨어 탐지, 입출력 어댑터 +├── compat/ # vLLM / transformers 호환 계층 +├── registry/ # 플러그인 레지스트리와 탐색 +├── core/ # 설정, DeviceFlowController +└── llm.py # vLLM 호환 LLM 프록시 +``` + +## 🧪 개발 + +```bash +pip install -e ".[dev]" +pytest -m "not gpu and not slow" tests/unit # 가속기 없이 도는 빠른 테스트 +ruff check src tests && black --check src tests +python tools/check_sensitive.py . # 내부 주소·토큰·키 검사 +bash scripts/run_examples.sh # 예제 전체 실행 (모델·장치 필요) +``` + +안내: [기여](docs/CONTRIBUTING.md) · [장치 추가](docs/ADDING_RESOURCES.md) · [엔진 추가](docs/ADDING_ENGINES.md). + +## 🗺 로드맵 + +- [x] rm-core v0.1 `AcceleratorProvider`, 노드 데몬, Helm chart +- [ ] `rm-abstract-node` 컨테이너 이미지 배포 +- [ ] `LLM` API를 통한 RNGD / ATOM 추론 검증 +- [ ] Optimus Prime용 엔진 백엔드 어댑터 (ai-rm [엔진 통합 설계](https://github.com/ai-computing/ai-rm/blob/master/docs/ko/engine-integration.md)) +- [ ] 장치 모델 확장: 파티션, 인터커넥트 토폴로지, 실시간 사용률 + +## 📚 문서 + +| | | +|---|---| +| **사용** | [빠른 시작](docs/QUICKSTART.md) · [서빙 엔진](docs/SERVING_ENGINES.md) · [Requirements](docs/README_REQUIREMENTS.md) | +| **이해** | [아키텍처](docs/ARCHITECTURE.md) · [프로젝트 구조](docs/PROJECT_STRUCTURE.md) · [변경 이력](CHANGELOG.md) | +| **확장** | [기여](docs/CONTRIBUTING.md) · [자원 추가](docs/ADDING_RESOURCES.md) · [엔진 추가](docs/ADDING_ENGINES.md) | +| **ai-rm** | [Umbrella](https://github.com/ai-computing/ai-rm) · [rm-core 프로토콜](https://github.com/ai-computing/rm-core/blob/v0.1.1/PROTOCOL.md) · [구성요소 카드](https://github.com/ai-computing/ai-rm/blob/master/docs/components/rm_abstract_layer.md) | + +예전 테스트 보고서와 발표 자료: [docs/archive](docs/archive/README.md). 상세 문서는 영어로 작성되어 있습니다. + +## 📄 라이선스 + +BSD-3-Clause. NPU 백엔드는 각 벤더 SDK(Rebellions RBLN SDK, FuriosaAI SDK)가 필요하며 해당 라이선스를 따릅니다. 이 +프로젝트는 SDK를 재배포하지 않습니다. 서드파티 라이선스 목록은 [영문 README](README.md#-license)에 있습니다. + +
ai-rm 구성요소 · ai-computing 팀 · ETRI
diff --git a/README.md b/README.md index b016327..4c155c2 100644 --- a/README.md +++ b/README.md @@ -1,415 +1,265 @@ -# RM Abstract Layer +
-> **GPU / NPU / CPU Heterogeneous Accelerator Unified LLM Inference Abstraction Layer** +

rm_abstract_layer

-[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/) -[![License: BSD-3-Clause](https://img.shields.io/badge/License-BSD--3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause) +

Accelerator abstraction for heterogeneous AI infrastructure

---- - -## Overview +

+ English · 한국어 +

-RM Abstract Layer is an abstraction layer that enables LLM inference on various hardware accelerators (GPU, CPU) with **minimal code modification**. +

+ CI + ai-rm layer L1 + Protocol v0.1 + Python 3.9+ + License: BSD-3-Clause +

-### Key Features +

+ Quick start · + ai-rm integration · + vLLM-compatible API · + Architecture · + Docs +

-- **Minimal Code Change**: Just change `from vllm import LLM` to `from rm_abstract import LLM` -- **Device Transparency**: vLLM-compatible API works on GPU and CPU -- **Runtime Switching**: Real-time switching between GPU and CPU -- **Auto Server Management**: Triton/TorchServe auto-start with model conversion -- **Plugin Architecture**: Easily add new resources and engines via plugin registry -- **Multiple Serving Engines**: vLLM, Triton, TorchServe support -- **Hardware Detection**: Automatic detection and priority-based selection (GPU > CPU) +
--- -## ai-rm - -[ai-rm](https://github.com/ai-computing/ai-rm)의 **L1 가속기 추상화** 구성요소입니다(`components/l1-accel/rm_abstract_layer`). [rm-core protocol v0.1](https://github.com/ai-computing/rm-core/blob/v0.1.1/PROTOCOL.md)을 구현합니다. - -- **역할**: 장치 탐지와 장치 능력 모델(`resources/base.py`의 `ResourceInfo`, `ResourceCapability`), 벤더별 자원 플러그인(`resources/gpu`, `resources/npu/furiosa.py`, `resources/npu/rebellions.py`, `resources/cpu`), 서빙 엔진 플러그인(vLLM, Triton, TorchServe). -- **adapter** (`rm_abstract.airm`): `AcceleratorProvider` `rm-abstract` (entry point `airm.accelerator_providers`). - - 장치별로 벤더 관리 도구(`nvidia-smi`, `furiosa-smi`, `rbln-stat`, 읽기 전용)를 먼저 쓰고, 도구가 없으면 기존 `HardwareDetector`로 탐지해 `rm_core.Accelerator`로 변환합니다. CPU는 보고하지 않으며 장치가 없으면 `[]`입니다. - - `id`는 `<노드>/:`이고 index는 장치 자체 번호(`npu3` → `furiosa:3`)입니다. 노드 이름은 `AIRM_NODE_NAME`, 없으면 hostname입니다. - - `peak_tflops`는 벤더 데이터시트 값이 있는 모델(RNGD, ATOM)만 채우고 `peak_source`에 출처를 적습니다. NVIDIA GPU는 비워 둡니다(미상). - - 생성자 인자: `node` (선택, 기본 `AIRM_NODE_NAME`/hostname). -- **노드 데몬**: `python -m rm_abstract.airm.daemon` — 주기적으로 탐지해 장치마다 `AcceleratorProfile`을 apply하고(연속으로 사라진 장치의 프로파일은 삭제), 노드 레이블 `airm.ai-computing.io/.=<개수>`, `airm.ai-computing.io/managed=true`를 패치합니다. `/metrics`에 `airm_accelerators{vendor,model}`을 냅니다. - ```bash - python -m rm_abstract.airm.daemon --once --dry-run # 만들 객체만 출력 (클러스터 접근 없음) - python -m rm_abstract.airm.daemon --provider sim --once --dry-run # 장치 없이 시연 - pip install -e ".[k8s]" # 실제 apply에는 kubernetes 클라이언트 필요 - ``` - 환경 변수: `AIRM_NODE_NAME`, `AIRM_PROVIDER`(기본 `rm-abstract`), `AIRM_INTERVAL_S`(기본 60), `AIRM_METRICS_PORT`(기본 9400, 0이면 끔), `AIRM_REMOVE_AFTER`(기본 3: 연속 3회 탐지되지 않은 장치만 삭제 — 도구 일시 실패로 프로파일이 깜박이지 않도록). -- **k8s + llm-d**: Helm chart `deploy/helm/rm-abstract-node/` (DaemonSet + ServiceAccount/ClusterRole: `acceleratorprofiles` CRUD, `nodes` get/patch). AcceleratorProfile CRD는 rm-core(`k8s/acceleratorprofile.crd.yaml`)에 있고 ai-rm umbrella chart/llm-d-setup이 설치합니다. 이미지 기본값은 `ghcr.io/ai-computing/rm-abstract-node:`이고, 사이트 override(`image.repository`/`image.tag`)로 바꿉니다. 배치 정책(EPP)과 lmd-top은 이 프로파일·레이블을 읽습니다. -- **설치·테스트**: `ai-rm.yaml` 참고 (`pip install -e .`, `pytest -m "not gpu and not slow" tests/unit`, `python -m rm_core.conformance airm.accelerator_providers rm-abstract`). -- **지표**: `accelerators` (rm-core `METRICS`). -- **관련 성과지표**: 이종 자원 수(장치 등록·탐지). -- **`cluster.local.env` 변수**: 없음. - -## Quick Start - -### Installation +**rm_abstract_layer** is the **L1 accelerator abstraction** of [ai-rm](https://github.com/ai-computing/ai-rm). It +discovers GPUs and NPUs (FuriosaAI RNGD, Rebellions ATOM) on each node, describes them through the +[rm-core](https://github.com/ai-computing/rm-core) `Accelerator` model and publishes them to Kubernetes as +`AcceleratorProfile` objects and node labels, which placement (llm-d EPP) and observation (lmd-top) read. + +It also provides a **vLLM-compatible Python API** that runs the same code on GPU or CPU with switchable serving +engines (vLLM, Triton, TorchServe). + + + + + + + +
+🔎 Discover devices
+Read-only vendor tools first (nvidia-smi, furiosa-smi, rbln-stat), SDK-based +detection as fallback. Device's own index, model, memory and datasheet peak TFLOPS. +
+☸️ Publish to Kubernetes
+A node daemon (DaemonSet) applies one AcceleratorProfile per device, labels nodes +airm.ai-computing.io/<vendor>.<model> and exports airm_accelerators. +
+🔁 One API, many backends
+Change from vllm import LLM to from rm_abstract import LLM; pick the device and engine by +argument or environment variable. +
+ +## 🚀 Quick start ```bash -# Basic installation (CPU with transformers) +git clone https://github.com/ai-computing/rm_abstract_layer.git && cd rm_abstract_layer pip install -r requirements/requirements.txt pip install -e . -# GPU support (vLLM) -pip install -r requirements/requirements-gpu.txt -pip install -e ".[gpu]" +python -m rm_abstract.airm.daemon --once --dry-run # devices on this node, as k8s objects +python -m rm_abstract.airm.daemon --provider sim --once --dry-run # same, with simulated devices +python -m rm_core.conformance airm.accelerator_providers rm-abstract # contract check (PASS / SKIP / FAIL) +``` + +Installing through ai-rm: `make setup COMPONENTS=rm_abstract_layer` in the +[ai-rm](https://github.com/ai-computing/ai-rm) umbrella (steps come from [`ai-rm.yaml`](ai-rm.yaml)). + +## ☸️ Device discovery for ai-rm -# Triton engine support (ONNX conversion) -pip install -r requirements/requirements-triton.txt -pip install -e ".[triton]" +### `AcceleratorProvider` adapter -# TorchServe engine support -pip install -r requirements/requirements-torchserve.txt -pip install -e ".[torchserve]" +`rm_abstract.airm` implements the rm-core v0.1 `AcceleratorProvider` contract and registers it as +`rm-abstract` in the `airm.accelerator_providers` entry-point group: -# All serving engines -pip install -r requirements/requirements-serving.txt -pip install -e ".[serving]" +```python +from rm_core import registry -# Full installation (everything) -pip install -r requirements/requirements-all.txt -pip install -e ".[all]" +provider = registry.load("airm.accelerator_providers", "rm-abstract") +for acc in provider.discover(): + print(acc.id, acc.vendor, acc.model, acc.memory_bytes, acc.peak_tflops) +# node-a/furiosa:0 furiosa RNGD ... {'bf16': 256.0, 'fp8': 512.0} ``` -### Basic Usage +| Field | Value | +|---|---| +| `id` | `/:`; index is the device's own number (`npu3` → `furiosa:3`) | +| node name | `AIRM_NODE_NAME` (set from `spec.nodeName` in the DaemonSet), else the hostname | +| `peak_tflops` | Only where the vendor datasheet states it (RNGD, ATOM), with `peak_source`; unknown otherwise | +| CPUs | Not reported; a node without accelerators yields `[]` | -```python -# Before (vLLM only - GPU required) -# from vllm import LLM, SamplingParams +### Node daemon -# After (rm_abstract - works on GPU and CPU) -from rm_abstract import LLM, SamplingParams +`python -m rm_abstract.airm.daemon` discovers periodically and -# Rest of your code remains the same! -llm = LLM(model="gpt2") -outputs = llm.generate(["Hello, I am"], SamplingParams(max_tokens=50)) -print(outputs[0].outputs[0].text) -``` +- applies one `AcceleratorProfile` (`airm.ai-computing.io/v1alpha1`) per device, deleting profiles of devices that + stayed missing for `AIRM_REMOVE_AFTER` consecutive rounds, +- patches node labels `airm.ai-computing.io/.=` and `airm.ai-computing.io/managed=true`, +- serves `airm_accelerators{vendor,model}` on `/metrics`. -### Device Selection +| Variable | Default | Meaning | +|---|---|---| +| `AIRM_NODE_NAME` | hostname | Node name used in ids and labels | +| `AIRM_PROVIDER` | `rm-abstract` | Provider entry point (`sim` for a demo) | +| `AIRM_INTERVAL_S` | `60` | Discovery interval | +| `AIRM_METRICS_PORT` | `9400` | `/metrics` port, `0` disables | +| `AIRM_REMOVE_AFTER` | `3` | Missed rounds before a profile is deleted (avoids flapping on a transient tool failure) | -```python -from rm_abstract import LLM, SamplingParams +Applying to a cluster needs the Kubernetes client: `pip install -e ".[k8s]"`. -# Auto-detect best device (GPU > CPU) -llm = LLM(model="gpt2") +### Helm chart -# Explicit device selection -llm = LLM(model="gpt2", device="nvidia") # NVIDIA GPU -llm = LLM(model="gpt2", device="cpu") # CPU +[`deploy/helm/rm-abstract-node`](deploy/helm/rm-abstract-node) installs the daemon as a DaemonSet with a +ServiceAccount and ClusterRole (`acceleratorprofiles` CRUD, `nodes` get/patch). The `AcceleratorProfile` CRD +ships with [rm-core](https://github.com/ai-computing/rm-core/tree/v0.1.1/k8s) and is installed by llm-d-setup / the +ai-rm umbrella. Set the image for your site with `image.repository` / `image.tag`. -# Environment variable control -# RM_DEVICE=nvidia python your_script.py -# RM_DEVICE=cpu python your_script.py +```bash +helm template rm-abstract-node deploy/helm/rm-abstract-node --set image.repository=/rm-abstract-node ``` -### Engine Selection +## 🔁 vLLM-compatible API ```python +# from vllm import LLM, SamplingParams from rm_abstract import LLM, SamplingParams -# vLLM engine (default) -llm = LLM(model="gpt2", engine="vllm") - -# Triton Inference Server (auto-starts Docker, converts to ONNX) -llm = LLM(model="gpt2", engine="triton") - -# TorchServe (auto-starts server, creates MAR package) -llm = LLM(model="gpt2", engine="torchserve") - -# Environment variable control -# RM_ENGINE=triton python your_script.py +llm = LLM(model="gpt2") # auto: GPU > CPU +llm = LLM(model="gpt2", device="cpu") # or RM_DEVICE=cpu +llm = LLM(model="gpt2", engine="triton") # or RM_ENGINE=triton +outputs = llm.generate(["Hello, I am"], SamplingParams(max_tokens=50)) +print(outputs[0].outputs[0].text) ``` -### Classic API Usage +
+Runtime control, engines and errors +
```python import rm_abstract -# Initialize (auto-select optimal device: GPU > CPU) rm_abstract.init(device="auto") - -# Check available resources -resources = rm_abstract.get_available_resources() -print(resources) -# {'nvidia': True, 'cpu': True} - -# Get device info -info = rm_abstract.get_device_info() -print(f"Current device: {info['resource_type']}:{info['device_id']}") - -# Runtime device switching +print(rm_abstract.get_available_resources()) # {'nvidia': True, 'cpu': True} rm_abstract.switch_device("cpu") -rm_abstract.switch_device("nvidia") # If available ``` ---- - -## Supported Environment - -### Hardware Resources - -| Device | Status | Backend | -|--------|--------|---------| -| NVIDIA GPU | ✅ Supported | vLLM / Triton / TorchServe | -| CPU | ✅ Supported | transformers / TorchServe | -| Rebellions ATOM NPU | 🔧 Planned | vLLM-RBLN / Optimum-RBLN | -| FuriosaAI RNGD NPU | 🔧 Planned | FuriosaAI SDK | - -### Serving Engines - -| Engine | Status | Auto Management | Description | -|--------|--------|-----------------|-------------| -| vLLM | ✅ Supported | - | High-performance LLM serving | -| Triton | ✅ Supported | Docker auto-start, ONNX conversion | Multi-model serving | -| TorchServe | ✅ Supported | Server auto-start, MAR packaging | PyTorch native serving | - ---- - -## Architecture - -```text -src/rm_abstract/ -├── __init__.py # Main interface (init, switch_device, LLM, SamplingParams) -├── llm.py # LLM proxy class (vLLM-compatible API) -├── exceptions.py # Exception classes -├── auto.py # Auto import hooking (use() function) -├── run.py # CLI runner -├── __main__.py # Module entry point -├── airm/ # ai-rm protocol v0.1 adapter (AcceleratorProvider, node daemon) -├── common/ # Shared utilities -│ ├── input_adapter.py # Unified input format handling -│ ├── output_adapter.py # Unified output format conversion -│ ├── model_converter.py # Model format conversion -│ └── hardware_detector.py # Hardware detection -├── compat/ # Compatibility modules -│ ├── vllm_compat.py # vLLM compatibility layer -│ └── transformers_compat.py # Transformers compatibility layer -├── resources/ # Hardware-specific implementations -│ ├── base.py # Resource ABC -│ ├── gpu/nvidia.py # NVIDIA GPU (vLLM) -│ ├── npu/rebellions.py # Rebellions ATOM NPU -│ ├── npu/furiosa.py # FuriosaAI RNGD NPU -│ └── cpu/default.py # CPU fallback -├── engines/ # Serving engines -│ ├── base.py # ServingEngine base class -│ ├── vllm.py # vLLM engine -│ ├── triton.py # Triton engine (auto-start, ONNX conversion) -│ ├── torchserve.py # TorchServe engine (auto-start, MAR packaging) -│ └── model_converter.py # HuggingFace to ONNX/MAR conversion -├── registry/ # Plugin system -│ ├── plugin_registry.py # Centralized registry -│ └── discovery.py # Auto-discovery -└── core/ # Core modules - ├── config.py # Configuration - └── controller.py # DeviceFlowController -``` - ---- - -## Auto Server Management - -### Triton Inference Server - -When using `engine="triton"`, rm_abstract will automatically: - -1. Convert HuggingFace model to ONNX format -2. Create Triton model repository with proper config.pbtxt -3. Start Docker container with Triton server -4. Clean up on exit +| Engine | Managed for you | +|---|---| +| vLLM (default) | — | +| Triton | HuggingFace → ONNX conversion, model repository, Docker container start/stop | +| TorchServe | MAR packaging, server start, model registration, cleanup | ```python -from rm_abstract import LLM, SamplingParams +from rm_abstract.exceptions import BackendNotAvailableError -# Triton with auto-management (requires Docker) -llm = LLM(model="gpt2", engine="triton", device="nvidia") -outputs = llm.generate(["Hello"], SamplingParams(max_tokens=50)) +try: + llm = LLM(model="gpt2", device="nvidia") +except BackendNotAvailableError: + llm = LLM(model="gpt2", device="cpu") ``` -### TorchServe - -When using `engine="torchserve"`, rm_abstract will automatically: - -1. Convert HuggingFace model to MAR (Model Archive) format -2. Create TorchServe configuration -3. Start TorchServe server -4. Register and load the model -5. Clean up on exit - -```python -from rm_abstract import LLM, SamplingParams - -# TorchServe with auto-management (requires torchserve installed) -llm = LLM(model="gpt2", engine="torchserve", device="cpu") -outputs = llm.generate(["Hello"], SamplingParams(max_tokens=50)) +Installation extras: `.[gpu]`, `.[triton]`, `.[torchserve]`, `.[serving]`, `.[all]` (see +[requirements guide](docs/README_REQUIREMENTS.md)). More in [quick start](docs/QUICKSTART.md) and +[serving engines](docs/SERVING_ENGINES.md). +
+ +### Support matrix + +| Device | Discovery & `AcceleratorProfile` | Inference through `LLM` | +|---|:-:|:-:| +| NVIDIA GPU | ✅ | ✅ vLLM · Triton · TorchServe | +| CPU | — *(not reported as an accelerator)* | ✅ transformers · TorchServe | +| FuriosaAI RNGD | ✅ | 🧪 plugin present, not validated | +| Rebellions ATOM | ✅ | 🧪 plugin present, not validated | + +## 🏗 Architecture + +```mermaid +flowchart LR + subgraph node["Node"] + tools["nvidia-smi · furiosa-smi · rbln-stat"] --> prov["rm_abstract.airm
AcceleratorProvider"] + det["HardwareDetector
(SDK fallback)"] --> prov + prov --> d["node daemon"] + end + d -- "AcceleratorProfile · node labels" --> k8s[("Kubernetes API")] + d -. "airm_accelerators" .-> prom[("Prometheus")] + k8s --> epp["llm-d EPP / placement"] + k8s --> top["lmd-top"] ``` ---- - -## Examples - -All 5 examples are tested and validated (✅ 100% pass rate): - ```text -examples/ -├── 01_quick_start.py # Simplest usage -├── 02_device_selection.py # Device control (CPU/GPU/NPU) -├── 03_auto_import_hook.py # Auto import hooking -├── 04_serving_engines.py # Test all serving engines -└── 05_engine_comparison.py # Performance comparison +src/rm_abstract/ +├── airm/ # ai-rm protocol v0.1: AcceleratorProvider, node daemon +├── resources/ # device plugins: gpu/nvidia, npu/furiosa, npu/rebellions, cpu +├── engines/ # serving engines: vLLM, Triton, TorchServe, model conversion +├── common/ # hardware detection, input/output adapters +├── compat/ # vLLM / transformers compatibility layers +├── registry/ # plugin registry and discovery +├── core/ # configuration, DeviceFlowController +└── llm.py # vLLM-compatible LLM proxy ``` -### Run Examples +## 🧪 Development ```bash -# Quick start -PYTHONPATH=./src python examples/01_quick_start.py - -# Device selection -PYTHONPATH=./src python examples/02_device_selection.py - -# Auto import hook -PYTHONPATH=./src python examples/03_auto_import_hook.py - -# Test all serving engines -PYTHONPATH=./src python examples/04_serving_engines.py - -# Compare engine performance -PYTHONPATH=./src python examples/05_engine_comparison.py - -# Run all examples with automated test -bash test_all_examples.sh +pip install -e ".[dev]" +pytest -m "not gpu and not slow" tests/unit # quick tests, no accelerator needed +ruff check src tests && black --check src tests +python tools/check_sensitive.py . # internal addresses, tokens, keys +bash scripts/run_examples.sh # run every example (needs models / devices) ``` ---- +Guides: [contributing](docs/CONTRIBUTING.md) · [adding a device](docs/ADDING_RESOURCES.md) · +[adding an engine](docs/ADDING_ENGINES.md). -## Testing +## 🗺 Roadmap -```bash -# Run all tests -pytest tests/ -v +- [x] rm-core v0.1 `AcceleratorProvider`, node daemon, Helm chart +- [ ] Published `rm-abstract-node` container image +- [ ] Validated RNGD / ATOM inference through the `LLM` API +- [ ] Engine-facing backend adapters for Optimus Prime (see the ai-rm [engine plan](https://github.com/ai-computing/ai-rm/blob/master/docs/engine-integration.md)) +- [ ] Wider device model: partitions, interconnect topology, live utilization -# Run specific test files -pytest tests/test_core.py -v -pytest tests/unit/test_llm.py -v -pytest tests/unit/test_registry.py -v -``` +## 📚 Documentation ---- +| | | +|---|---| +| **Use** | [Quick start](docs/QUICKSTART.md) · [Serving engines](docs/SERVING_ENGINES.md) · [Requirements](docs/README_REQUIREMENTS.md) | +| **Understand** | [Architecture](docs/ARCHITECTURE.md) · [Project structure](docs/PROJECT_STRUCTURE.md) · [Changelog](CHANGELOG.md) | +| **Extend** | [Contributing](docs/CONTRIBUTING.md) · [Adding resources](docs/ADDING_RESOURCES.md) · [Adding engines](docs/ADDING_ENGINES.md) | +| **ai-rm** | [Umbrella](https://github.com/ai-computing/ai-rm) · [rm-core protocol](https://github.com/ai-computing/rm-core/blob/v0.1.1/PROTOCOL.md) · [Component card](https://github.com/ai-computing/ai-rm/blob/master/docs/components/rm_abstract_layer.md) | -## Error Handling +Older test reports and slides: [docs/archive](docs/archive/README.md). -```python -from rm_abstract.exceptions import ( - RMAbstractError, # Base exception - NotInitializedError, # Not initialized - BackendNotAvailableError, # Backend unavailable - InvalidDeviceError, # Invalid device spec -) +## 📄 License -try: - from rm_abstract import LLM - llm = LLM(model="gpt2", device="nvidia") -except BackendNotAvailableError as e: - print(f"GPU not available: {e}") - llm = LLM(model="gpt2", device="cpu") # Fallback to CPU -``` +BSD-3-Clause. NPU backends need the vendors' own SDKs (Rebellions RBLN SDK, FuriosaAI SDK) under their +license terms; this project does not redistribute them. ---- +
+Third-party licenses +
-## Device Options - -```python -from rm_abstract import LLM - -# Auto-select (GPU > CPU priority) -llm = LLM(model="gpt2", device="auto") - -# Specific GPU -llm = LLM(model="gpt2", device="nvidia") -llm = LLM(model="gpt2", device="gpu") # alias -llm = LLM(model="gpt2", device="cuda") # alias - -# CPU -llm = LLM(model="gpt2", device="cpu") - -# NPU (Planned - not yet supported) -# llm = LLM(model="gpt2", device="rbln") # Rebellions NPU -# llm = LLM(model="gpt2", device="furiosa") # FuriosaAI NPU -``` - ---- - -## Documentation - -### Getting Started -| Document | Description | -|----------|-------------| -| [docs/QUICKSTART.md](docs/QUICKSTART.md) | 5-minute quick start guide | -| [docs/SERVING_ENGINES.md](docs/SERVING_ENGINES.md) | Complete guide for all serving engines | -| [docs/README_REQUIREMENTS.md](docs/README_REQUIREMENTS.md) | Installation and dependencies guide | - -### Technical Documentation -| Document | Description | -|----------|-------------| -| [docs/ARCHITECTURE.md](docs/ARCHITECTURE.md) | System architecture and design | -| [docs/PROJECT_STRUCTURE.md](docs/PROJECT_STRUCTURE.md) | Project layout and organization | -| [docs/FINAL_TEST_REPORT.md](docs/FINAL_TEST_REPORT.md) | Test results and validation | -| [CLEANUP_SUMMARY.md](CLEANUP_SUMMARY.md) | Project cleanup summary | - -### Development -| Document | Description | -|----------|-------------| -| [docs/CONTRIBUTING.md](docs/CONTRIBUTING.md) | Development and contribution guide | -| [docs/ADDING_RESOURCES.md](docs/ADDING_RESOURCES.md) | Guide for adding new hardware resources | -| [docs/ADDING_ENGINES.md](docs/ADDING_ENGINES.md) | Guide for adding new serving engines | - -### Presentation -| Document | Description | -|----------|-------------| -| [docs/PRESENTATION.md](docs/PRESENTATION.md) | Project presentation slides (Korean) | - ---- - -## License - -BSD-3-Clause License - -### Third-Party Licenses - -This project integrates with various third-party libraries and SDKs: - -#### Open Source Dependencies (Permissive Licenses) - -| Package | License | Usage | -|---------|---------|-------| -| [PyTorch](https://github.com/pytorch/pytorch) | BSD-3-Clause | Deep learning framework | -| [vLLM](https://github.com/vllm-project/vllm) | Apache-2.0 | High-performance LLM serving | -| [Triton Inference Server](https://github.com/triton-inference-server/server) | BSD-3-Clause | Multi-model serving | -| [TorchServe](https://github.com/pytorch/serve) | Apache-2.0 | PyTorch model serving | -| [HuggingFace Transformers](https://github.com/huggingface/transformers) | Apache-2.0 | Model hub and inference | -| [vllm-rbln](https://github.com/rebellions-sw/vllm-rbln) | Apache-2.0 | vLLM plugin for RBLN NPU | -| [optimum-rbln](https://github.com/rebellions-sw/optimum-rbln) | Apache-2.0 | HuggingFace integration for RBLN | -| [FuriosaAI SDK](https://github.com/furiosa-ai/furiosa-sdk) | Apache-2.0 | FuriosaAI Python SDK | - -#### Proprietary NPU SDKs - -The following NPU backends require proprietary SDKs with separate license agreements: +| Package | License | +|---|---| +| [PyTorch](https://github.com/pytorch/pytorch) | BSD-3-Clause | +| [vLLM](https://github.com/vllm-project/vllm) | Apache-2.0 | +| [Triton Inference Server](https://github.com/triton-inference-server/server) | BSD-3-Clause | +| [TorchServe](https://github.com/pytorch/serve) | Apache-2.0 | +| [HuggingFace Transformers](https://github.com/huggingface/transformers) | Apache-2.0 | +| [vllm-rbln](https://github.com/rebellions-sw/vllm-rbln) · [optimum-rbln](https://github.com/rebellions-sw/optimum-rbln) | Apache-2.0 | +| [FuriosaAI SDK](https://github.com/furiosa-ai/furiosa-sdk) | Apache-2.0 | | NPU | SDK | License | Access | -|-----|-----|---------|--------| -| **Rebellions ATOM** | RBLN SDK (rebel-compiler, drivers, runtime) | Proprietary | [Request Access](https://rebellions.ai/request-form-rbln-sdk/) | -| **FuriosaAI RNGD** | FuriosaAI SDK (kernel driver, firmware, runtime) | Proprietary (EULA) | [Early Access Program](https://furiosa.ai/) | - -> **Note**: Users are responsible for complying with the respective SDK license terms when using NPU backends. This project only provides abstraction interfaces and does not redistribute any proprietary SDK components. - ---- - -## Contributing +|---|---|---|---| +| Rebellions ATOM | RBLN SDK (compiler, driver, runtime) | Proprietary | [Request access](https://rebellions.ai/request-form-rbln-sdk/) | +| FuriosaAI RNGD | FuriosaAI SDK (driver, firmware, runtime) | Proprietary (EULA) | [furiosa.ai](https://furiosa.ai/) | +
-Contributions are welcome! See [docs/CONTRIBUTING.md](docs/CONTRIBUTING.md) for details. +
Part of ai-rm · ai-computing team · ETRI
diff --git a/deploy/helm/rm-abstract-node/Chart.yaml b/deploy/helm/rm-abstract-node/Chart.yaml index 877695b..ddc2aa3 100644 --- a/deploy/helm/rm-abstract-node/Chart.yaml +++ b/deploy/helm/rm-abstract-node/Chart.yaml @@ -2,5 +2,5 @@ apiVersion: v2 name: rm-abstract-node description: ai-rm L1 node daemon - publishes each node's accelerators as AcceleratorProfiles and node labels type: application -version: 0.3.0 -appVersion: "0.3.0" +version: 0.4.0 +appVersion: "0.4.0" diff --git a/docs/PROJECT_STRUCTURE.md b/docs/PROJECT_STRUCTURE.md index a3b651e..ec09168 100644 --- a/docs/PROJECT_STRUCTURE.md +++ b/docs/PROJECT_STRUCTURE.md @@ -8,12 +8,11 @@ rm_abstract_layer/ │ ├── QUICKSTART.md # Quick start guide │ ├── SERVING_ENGINES.md # Serving engines guide │ ├── README_REQUIREMENTS.md # Requirements installation -│ ├── FINAL_TEST_REPORT.md # Latest test results ✅ -│ ├── TEST_RESULTS.md # Initial test findings │ ├── ARCHITECTURE.md # System architecture │ ├── CONTRIBUTING.md # Development guide │ ├── ADDING_RESOURCES.md # Adding new hardware -│ └── ADDING_ENGINES.md # Adding new engines +│ ├── ADDING_ENGINES.md # Adding new engines +│ └── archive/ # 0.2–0.3 test reports and slides (not maintained) │ ├── examples/ # 🎯 Examples (cleaned up!) │ ├── 01_quick_start.py # Simplest usage @@ -33,6 +32,7 @@ rm_abstract_layer/ │ ├── src/rm_abstract/ # 🔧 Source code │ ├── __init__.py # Main API exports +│ ├── airm/ # ai-rm protocol v0.1: AcceleratorProvider, node daemon │ ├── auto.py # Auto import hooking ⭐ │ ├── llm.py # LLM proxy class │ ├── exceptions.py # Exception classes @@ -72,7 +72,11 @@ rm_abstract_layer/ │ ├── integration/ # Integration tests │ └── conftest.py # Pytest configuration │ -├── test_multi_engine.py # Multi-engine test script +├── deploy/helm/rm-abstract-node/ # Helm chart for the node daemon (DaemonSet) +├── scripts/ # setup scripts, run_examples.sh, manual/multi_engine_check.py +├── tools/check_sensitive.py # internal addresses / tokens / keys scan +├── ai-rm.yaml # ai-rm component manifest +├── CHANGELOG.md ├── pyproject.toml # Project metadata └── README.md # Main README ``` @@ -109,7 +113,7 @@ rm_abstract_layer/ | `docs/QUICKSTART.md` | New users (5-min guide) | | `docs/SERVING_ENGINES.md` | All users (engine details) | | `docs/README_REQUIREMENTS.md` | All users (installation) | -| `docs/FINAL_TEST_REPORT.md` | All users (test results) | +| `docs/archive/FINAL_TEST_REPORT.md` | All users (test results) | | `docs/ARCHITECTURE.md` | Developers | | `docs/CONTRIBUTING.md` | Contributors | @@ -136,7 +140,7 @@ rm_abstract_layer/ - `QUICKSTART.md` - `SERVING_ENGINES.md` - `README_REQUIREMENTS.md` - - `FINAL_TEST_REPORT.md` + - `archive/FINAL_TEST_REPORT.md` - `PROJECT_STRUCTURE.md` (this file) ### ✅ Improved @@ -163,13 +167,13 @@ rm_abstract_layer/ ### For Testers -1. **Test Report**: `docs/FINAL_TEST_REPORT.md` -2. **Test Script**: `test_multi_engine.py` +1. **Test Report**: `docs/archive/FINAL_TEST_REPORT.md` +2. **Test Script**: `scripts/manual/multi_engine_check.py` 3. **Engine Tests**: `examples/04_serving_engines.py` ## Next Steps -See `docs/FINAL_TEST_REPORT.md` for: +See `docs/archive/FINAL_TEST_REPORT.md` for: - Test results (all engines ✅) - Known issues - Recommendations diff --git a/docs/QUICKSTART.md b/docs/QUICKSTART.md index a06182f..89a018e 100644 --- a/docs/QUICKSTART.md +++ b/docs/QUICKSTART.md @@ -243,8 +243,8 @@ from vllm import LLM ## Get Help -- **Issues**: https://github.com/gspark-etri/rm_abstract_layer/issues -- **Discussions**: https://github.com/gspark-etri/rm_abstract_layer/discussions +- **Issues**: https://github.com/ai-computing/rm_abstract_layer/issues +- **Discussions**: https://github.com/ai-computing/rm_abstract_layer/discussions ## Verified Environments diff --git a/CLEANUP_SUMMARY.md b/docs/archive/CLEANUP_SUMMARY.md similarity index 100% rename from CLEANUP_SUMMARY.md rename to docs/archive/CLEANUP_SUMMARY.md diff --git a/docs/FINAL_TEST_REPORT.md b/docs/archive/FINAL_TEST_REPORT.md similarity index 100% rename from docs/FINAL_TEST_REPORT.md rename to docs/archive/FINAL_TEST_REPORT.md diff --git a/docs/PRESENTATION.md b/docs/archive/PRESENTATION.md similarity index 100% rename from docs/PRESENTATION.md rename to docs/archive/PRESENTATION.md diff --git a/docs/archive/README.md b/docs/archive/README.md new file mode 100644 index 0000000..6b5e6fb --- /dev/null +++ b/docs/archive/README.md @@ -0,0 +1,12 @@ +# Archive + +Point-in-time reports from the 0.2–0.3 development cycle, kept for reference. They describe the +project as it was then (vLLM-compatible API on GPU/CPU) and are not maintained. Paths and commands +inside them may no longer match the repository; see the main [README](../../README.md) for the +current state. + +| File | Contents | +|---|---| +| `CLEANUP_SUMMARY.md` | Repository cleanup summary (0.2) | +| `FINAL_TEST_REPORT.md`, `TEST_RESULTS.md` | Example and engine test reports (0.2) | +| `PRESENTATION.md` | Project presentation slides (Korean) | diff --git a/docs/TEST_RESULTS.md b/docs/archive/TEST_RESULTS.md similarity index 100% rename from docs/TEST_RESULTS.md rename to docs/archive/TEST_RESULTS.md diff --git a/pyproject.toml b/pyproject.toml index ee87359..c3aae8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,15 +4,15 @@ build-backend = "setuptools.build_meta" [project] name = "rm-abstract" -version = "0.3.0" -description = "LLM Heterogeneous Resource Orchestrator - GPU/NPU/PIM Meta-Serving Layer" +version = "0.4.0" +description = "Accelerator abstraction for heterogeneous AI infrastructure: GPU/NPU/CPU discovery, AcceleratorProfile publishing and a vLLM-compatible API (ai-rm L1)" readme = "README.md" license = {text = "BSD-3-Clause"} requires-python = ">=3.9" authors = [ - {name = "RM Abstract Team", email = "contact@rmabstract.dev"} + {name = "ai-computing team, ETRI"} ] -keywords = ["llm", "gpu", "npu", "pim", "inference", "serving", "orchestration", "heterogeneous"] +keywords = ["llm", "gpu", "npu", "rngd", "atom", "accelerator", "kubernetes", "llm-d", "ai-rm", "heterogeneous"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", @@ -126,10 +126,10 @@ all = [ ] [project.urls] -Homepage = "https://github.com/gspark-etri/rm_abstract_layer" -Documentation = "https://github.com/gspark-etri/rm_abstract_layer#readme" -Repository = "https://github.com/gspark-etri/rm_abstract_layer.git" -Issues = "https://github.com/gspark-etri/rm_abstract_layer/issues" +Homepage = "https://github.com/ai-computing/rm_abstract_layer" +Documentation = "https://github.com/ai-computing/rm_abstract_layer#readme" +Repository = "https://github.com/ai-computing/rm_abstract_layer.git" +Issues = "https://github.com/ai-computing/rm_abstract_layer/issues" [tool.setuptools] package-dir = {"" = "src"} @@ -159,6 +159,8 @@ extend-exclude = ''' [tool.ruff] line-length = 88 target-version = "py39" + +[tool.ruff.lint] select = [ "E", # pycodestyle errors "W", # pycodestyle warnings @@ -169,12 +171,15 @@ select = [ "UP", # pyupgrade ] ignore = [ - "E501", # line too long (handled by black) - "B008", # do not perform function calls in argument defaults - "C901", # too complex + "E501", # line too long (handled by black) + "B008", # do not perform function calls in argument defaults + "C901", # too complex + "UP006", # typing.List -> list: style only, kept as-is for now + "UP035", # deprecated typing imports: style only, kept as-is for now + "B904", # raise ... from err: to be fixed module by module ] -[tool.ruff.per-file-ignores] +[tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401"] [tool.mypy] diff --git a/test_multi_engine.py b/scripts/manual/multi_engine_check.py similarity index 98% rename from test_multi_engine.py rename to scripts/manual/multi_engine_check.py index c684ee5..4e627e1 100644 --- a/test_multi_engine.py +++ b/scripts/manual/multi_engine_check.py @@ -5,7 +5,7 @@ import os # Add src to path -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) from rm_abstract import LLM, SamplingParams from rm_abstract.auto import use diff --git a/test_all_examples.sh b/scripts/run_examples.sh similarity index 97% rename from test_all_examples.sh rename to scripts/run_examples.sh index 1afbdf1..847eb1e 100755 --- a/test_all_examples.sh +++ b/scripts/run_examples.sh @@ -6,7 +6,7 @@ echo "TESTING ALL EXAMPLES" echo "================================================================================" echo "" -cd /home/gspark/rm_abstract_layer +cd "$(dirname "$0")/.." total=0 passed=0 diff --git a/src/rm_abstract/__init__.py b/src/rm_abstract/__init__.py index 62de7fd..152db33 100644 --- a/src/rm_abstract/__init__.py +++ b/src/rm_abstract/__init__.py @@ -19,26 +19,26 @@ - registry/: Centralized plugin registration system """ -from typing import Optional, Dict, Any import os +from typing import Any, Dict, Optional -from .core.controller import DeviceFlowController +from .common import HardwareDetector, InputAdapter, ModelConverter, OutputAdapter from .core.config import Config +from .core.controller import DeviceFlowController from .exceptions import ( - RMAbstractError, - NotInitializedError, BackendNotAvailableError, InvalidDeviceError, + NotInitializedError, + RMAbstractError, ) # Main LLM interface (recommended) -from .llm import LLM, SamplingParams, RequestOutput, CompletionOutput +from .llm import LLM, CompletionOutput, RequestOutput, SamplingParams # Modular components -from .registry import PluginRegistry, PluginType, get_registry, discover_all -from .common import InputAdapter, OutputAdapter, ModelConverter, HardwareDetector +from .registry import PluginRegistry, PluginType, discover_all, get_registry -__version__ = "0.3.0" +__version__ = "0.4.0" __all__ = [ # Main LLM interface (recommended) "LLM", @@ -115,7 +115,9 @@ def init( _initialized = True if verbose: - print(f"[RM Abstract] Initialized with device: {_global_controller.device_name}") + print( + f"[RM Abstract] Initialized with device: {_global_controller.device_name}" + ) return _global_controller diff --git a/src/rm_abstract/__main__.py b/src/rm_abstract/__main__.py index 2336bc8..887b044 100644 --- a/src/rm_abstract/__main__.py +++ b/src/rm_abstract/__main__.py @@ -7,13 +7,14 @@ python -m rm_abstract list-devices """ -import sys import argparse +import sys def cmd_run(args): """Run a script with hooks enabled.""" from rm_abstract.run import main + sys.argv = ["rm_abstract.run"] + args main() @@ -54,7 +55,9 @@ def main(): # run command run_parser = subparsers.add_parser("run", help="Run a script with hooks enabled") - run_parser.add_argument("script_args", nargs=argparse.REMAINDER, help="Script and arguments") + run_parser.add_argument( + "script_args", nargs=argparse.REMAINDER, help="Script and arguments" + ) # info command subparsers.add_parser("info", help="Show RM Abstract information") diff --git a/src/rm_abstract/auto.py b/src/rm_abstract/auto.py index 575c8b9..4282702 100644 --- a/src/rm_abstract/auto.py +++ b/src/rm_abstract/auto.py @@ -17,10 +17,10 @@ engine: vllm, triton, torchserve """ -import sys -import os import logging -from importlib.abc import MetaPathFinder, Loader +import os +import sys +from importlib.abc import Loader, MetaPathFinder from importlib.machinery import ModuleSpec from types import ModuleType from typing import Optional @@ -101,7 +101,9 @@ def use( disable_hooks() enable_hooks() - logger.info(f"RM Abstract configured: device={_config['device']}, engine={_config['engine']}, hook_vllm={_config['hook_vllm']}") + logger.info( + f"RM Abstract configured: device={_config['device']}, engine={_config['engine']}, hook_vllm={_config['hook_vllm']}" + ) def get_config() -> dict: @@ -119,6 +121,7 @@ def __init__(self): def __getattr__(self, name): from rm_abstract.compat import vllm_compat + return getattr(vllm_compat, name) @@ -139,7 +142,10 @@ def _get_original(self): sys.meta_path.remove(finder) try: import importlib - self._original_transformers = importlib.import_module("transformers") + + self._original_transformers = importlib.import_module( + "transformers" + ) finally: sys.meta_path.insert(0, finder) break @@ -148,6 +154,7 @@ def _get_original(self): def __getattr__(self, name): if name in ("AutoModelForCausalLM", "AutoTokenizer"): from rm_abstract.compat import transformers_compat + return getattr(transformers_compat, name) original = self._get_original() diff --git a/src/rm_abstract/common/__init__.py b/src/rm_abstract/common/__init__.py index 7364bdb..06b2239 100644 --- a/src/rm_abstract/common/__init__.py +++ b/src/rm_abstract/common/__init__.py @@ -8,10 +8,10 @@ - HardwareDetector: Hardware detection utilities """ +from .hardware_detector import HardwareDetector, HardwareInfo from .input_adapter import InputAdapter, InputFormat +from .model_converter import ConversionConfig, ModelConverter from .output_adapter import OutputAdapter, OutputFormat -from .model_converter import ModelConverter, ConversionConfig -from .hardware_detector import HardwareDetector, HardwareInfo __all__ = [ "InputAdapter", diff --git a/src/rm_abstract/common/hardware_detector.py b/src/rm_abstract/common/hardware_detector.py index c3a1c65..1c6e5ba 100644 --- a/src/rm_abstract/common/hardware_detector.py +++ b/src/rm_abstract/common/hardware_detector.py @@ -5,10 +5,10 @@ Detects available devices (GPU, NPU, CPU) across different vendors. """ +import logging from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Optional, Callable -import logging +from typing import Any, Callable, Dict, List, Optional logger = logging.getLogger(__name__) @@ -383,8 +383,8 @@ def _detect_furiosa_npu(self) -> List[HardwareInfo]: def _detect_cpu(self) -> List[HardwareInfo]: """Detect CPU (always available)""" - import platform import os + import platform return [ HardwareInfo( diff --git a/src/rm_abstract/common/input_adapter.py b/src/rm_abstract/common/input_adapter.py index b916124..0833257 100644 --- a/src/rm_abstract/common/input_adapter.py +++ b/src/rm_abstract/common/input_adapter.py @@ -5,10 +5,10 @@ Eliminates duplicated input handling logic across GPU, NPU, and CPU backends. """ +import logging from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, List, Optional, Union -import logging +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) @@ -242,7 +242,9 @@ def _adapt_token_dict( original_format=InputFormat.TOKEN_IDS_DICT, ) - def _adapt_token_ids(self, token_ids: Any, tokenizer: Optional[Any]) -> AdaptedInput: + def _adapt_token_ids( + self, token_ids: Any, tokenizer: Optional[Any] + ) -> AdaptedInput: """Adapt token ID tensor or list""" # Try to decode to text prompts prompts = self._decode_token_ids(token_ids, tokenizer) diff --git a/src/rm_abstract/common/model_converter.py b/src/rm_abstract/common/model_converter.py index 9c8d754..4bbd902 100644 --- a/src/rm_abstract/common/model_converter.py +++ b/src/rm_abstract/common/model_converter.py @@ -5,12 +5,12 @@ Handles ONNX conversion and other format transformations. """ -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional, Tuple -from pathlib import Path -import logging import hashlib +import logging import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) @@ -136,9 +136,7 @@ def to_onnx( logger.info(f"ONNX model saved to: {output_path}") return output_path - def _create_dummy_input( - self, model: Any, config: ConversionConfig - ) -> Any: + def _create_dummy_input(self, model: Any, config: ConversionConfig) -> Any: """ Create dummy input for ONNX export. diff --git a/src/rm_abstract/common/output_adapter.py b/src/rm_abstract/common/output_adapter.py index ae7e96e..e14b506 100644 --- a/src/rm_abstract/common/output_adapter.py +++ b/src/rm_abstract/common/output_adapter.py @@ -5,10 +5,10 @@ Eliminates duplicated output handling logic across backends. """ +import logging from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, List, Optional, Union -import logging +from typing import Any, List, Optional logger = logging.getLogger(__name__) @@ -119,9 +119,7 @@ def detect_format(self, outputs: Any) -> OutputFormat: return OutputFormat.TEXT - def adapt( - self, outputs: Any, tokenizer: Optional[Any] = None - ) -> AdaptedOutput: + def adapt(self, outputs: Any, tokenizer: Optional[Any] = None) -> AdaptedOutput: """ Adapt outputs to standardized format. diff --git a/src/rm_abstract/compat/__init__.py b/src/rm_abstract/compat/__init__.py index 849f9dc..3a48ed8 100644 --- a/src/rm_abstract/compat/__init__.py +++ b/src/rm_abstract/compat/__init__.py @@ -5,8 +5,8 @@ that automatically route to the appropriate backend based on device selection. """ -from .vllm_compat import LLM, SamplingParams, RequestOutput from .transformers_compat import AutoModelForCausalLM, AutoTokenizer +from .vllm_compat import LLM, RequestOutput, SamplingParams __all__ = [ "LLM", diff --git a/src/rm_abstract/compat/transformers_compat.py b/src/rm_abstract/compat/transformers_compat.py index 54185f8..3029fbb 100644 --- a/src/rm_abstract/compat/transformers_compat.py +++ b/src/rm_abstract/compat/transformers_compat.py @@ -15,8 +15,8 @@ tokenizer = AutoTokenizer.from_pretrained("gpt2") """ -import os import logging +import os from typing import Any, Dict, Optional, Union logger = logging.getLogger(__name__) @@ -38,18 +38,21 @@ def _detect_device() -> str: """Detect the best available device on this host (NPU > GPU > CPU).""" try: import rebel # noqa: F401 + return "rbln" except ImportError: pass try: import furiosa_sdk # noqa: F401 + return "furiosa" except ImportError: pass try: import torch + if torch.cuda.is_available(): return "nvidia" except ImportError: @@ -80,7 +83,7 @@ def from_pretrained( trust_remote_code: bool = False, torch_dtype: Any = None, device_map: Optional[Union[str, Dict]] = None, - **kwargs + **kwargs, ) -> Any: """ Load a pretrained model with automatic backend selection. @@ -121,10 +124,7 @@ def _load_rbln(cls, model_name: str, trust_remote_code: bool, **kwargs) -> Any: from optimum.rbln import RBLNAutoModelForCausalLM model = RBLNAutoModelForCausalLM.from_pretrained( - model_name, - export=True, - trust_remote_code=trust_remote_code, - **kwargs + model_name, export=True, trust_remote_code=trust_remote_code, **kwargs ) logger.info(f"Loaded model with optimum-rbln: {model_name}") return model @@ -140,9 +140,7 @@ def _load_furiosa(cls, model_name: str, trust_remote_code: bool, **kwargs) -> An from furiosa_llm import LLM as FuriosaLLM model = FuriosaLLM.from_pretrained( - model_name, - trust_remote_code=trust_remote_code, - **kwargs + model_name, trust_remote_code=trust_remote_code, **kwargs ) logger.info(f"Loaded model with FuriosaAI: {model_name}") return model @@ -158,6 +156,7 @@ def _get_hf_automodel(cls): import sys from rm_abstract.auto import disable_hooks, enable_hooks + disable_hooks() try: @@ -175,7 +174,7 @@ def _load_gpu( trust_remote_code: bool, torch_dtype: Any, device_map: Optional[Union[str, Dict]], - **kwargs + **kwargs, ) -> Any: """Load model on GPU using transformers.""" import torch @@ -193,18 +192,14 @@ def _load_gpu( trust_remote_code=trust_remote_code, torch_dtype=torch_dtype, device_map=device_map, - **kwargs + **kwargs, ) logger.info(f"Loaded model on GPU: {model_name}") return model @classmethod def _load_cpu( - cls, - model_name: str, - trust_remote_code: bool, - torch_dtype: Any = None, - **kwargs + cls, model_name: str, trust_remote_code: bool, torch_dtype: Any = None, **kwargs ) -> Any: """Load model on CPU using transformers.""" import torch @@ -219,7 +214,7 @@ def _load_cpu( trust_remote_code=trust_remote_code, torch_dtype=torch_dtype, device_map="cpu", - **kwargs + **kwargs, ) logger.info(f"Loaded model on CPU: {model_name}") return model @@ -240,6 +235,7 @@ def _get_hf_autotokenizer(cls): import sys from rm_abstract.auto import disable_hooks, enable_hooks + disable_hooks() try: diff --git a/src/rm_abstract/compat/vllm_compat.py b/src/rm_abstract/compat/vllm_compat.py index a52d70f..188e486 100644 --- a/src/rm_abstract/compat/vllm_compat.py +++ b/src/rm_abstract/compat/vllm_compat.py @@ -15,10 +15,10 @@ output = llm.generate("Hello", SamplingParams()) """ -import os import logging -from typing import Any, Dict, List, Optional, Union +import os from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Union logger = logging.getLogger(__name__) @@ -38,6 +38,7 @@ def _is_engine_available(engine: str) -> bool: if engine == "vllm": try: import builtins + builtins.__import__("vllm", globals(), locals(), [], 0) return True except ImportError: @@ -45,6 +46,7 @@ def _is_engine_available(engine: str) -> bool: elif engine == "triton": try: import subprocess + result = subprocess.run( ["docker", "--version"], capture_output=True, @@ -56,6 +58,7 @@ def _is_engine_available(engine: str) -> bool: elif engine == "torchserve": try: import subprocess + result = subprocess.run( ["torchserve", "--version"], capture_output=True, @@ -74,6 +77,7 @@ class SamplingParams: Works across all backends with automatic parameter translation. """ + n: int = 1 best_of: Optional[int] = None presence_penalty: float = 0.0 @@ -110,6 +114,7 @@ def to_dict(self) -> Dict[str, Any]: @dataclass class CompletionOutput: """Single completion output.""" + index: int text: str token_ids: List[int] = field(default_factory=list) @@ -120,6 +125,7 @@ class CompletionOutput: @dataclass class RequestOutput: """vLLM-compatible request output.""" + request_id: str prompt: str prompt_token_ids: List[int] @@ -168,7 +174,7 @@ def __init__( tensor_parallel_size: int = 1, dtype: str = "auto", trust_remote_code: bool = False, - **kwargs + **kwargs, ): """ Initialize LLM. @@ -229,18 +235,21 @@ def _resolve_device(self, device: str) -> str: # Auto-detection priority: NPU > GPU > CPU try: import rebel # noqa: F401 + return "rbln" except ImportError: pass try: import furiosa_sdk # noqa: F401 + return "furiosa" except ImportError: pass try: import torch + if torch.cuda.is_available(): return "nvidia" except ImportError: @@ -255,6 +264,7 @@ def _init_vllm_backend(self): # Directly use real vllm module # We access it through __builtins__ import to bypass any hooks import builtins + real_import = builtins.__import__ vllm_module = real_import("vllm", globals(), locals(), [], 0) @@ -267,7 +277,7 @@ def _init_vllm_backend(self): tensor_parallel_size=self.tensor_parallel_size, dtype=self.dtype, trust_remote_code=self.trust_remote_code, - **self.extra_kwargs + **self.extra_kwargs, ) self._sampling_params_cls = VLLMSamplingParams logger.info(f"Initialized vLLM backend for model: {self.model_name}") @@ -282,6 +292,7 @@ def _init_rbln_backend(self): try: # Try vllm-rbln first - use builtins import to bypass hooks import builtins + real_import = builtins.__import__ vllm_module = real_import("vllm", globals(), locals(), [], 0) @@ -292,7 +303,7 @@ def _init_rbln_backend(self): tokenizer=self.tokenizer_name, device="rbln", trust_remote_code=self.trust_remote_code, - **self.extra_kwargs + **self.extra_kwargs, ) logger.info(f"Initialized vLLM-RBLN backend for model: {self.model_name}") @@ -308,7 +319,9 @@ def _init_rbln_backend(self): trust_remote_code=self.trust_remote_code, ) self._tokenizer = AutoTokenizer.from_pretrained(self.model_name) - logger.info(f"Initialized optimum-rbln backend for model: {self.model_name}") + logger.info( + f"Initialized optimum-rbln backend for model: {self.model_name}" + ) except ImportError: logger.warning("RBLN backend not available, falling back to CPU") @@ -323,7 +336,7 @@ def _init_furiosa_backend(self): self._backend_llm = FuriosaLLM.from_pretrained( self.model_name, trust_remote_code=self.trust_remote_code, - **self.extra_kwargs + **self.extra_kwargs, ) logger.info(f"Initialized FuriosaAI backend for model: {self.model_name}") @@ -335,8 +348,12 @@ def _init_triton_backend(self, device: str): """Initialize Triton Inference Server backend.""" self._backend = "triton" try: - from rm_abstract.engines.triton import TritonEngine, TritonConfig - from rm_abstract.engines.base import ServingConfig, ServingEngineType, DeviceTarget + from rm_abstract.engines.base import ( + DeviceTarget, + ServingConfig, + ServingEngineType, + ) + from rm_abstract.engines.triton import TritonEngine # Map device string to DeviceTarget device_map = { @@ -368,8 +385,12 @@ def _init_torchserve_backend(self, device: str): """Initialize TorchServe backend.""" self._backend = "torchserve" try: + from rm_abstract.engines.base import ( + DeviceTarget, + ServingConfig, + ServingEngineType, + ) from rm_abstract.engines.torchserve import TorchServeEngine - from rm_abstract.engines.base import ServingConfig, ServingEngineType, DeviceTarget # Map device string to DeviceTarget device_map = { @@ -393,7 +414,9 @@ def _init_torchserve_backend(self, device: str): logger.info(f"Initialized TorchServe backend for model: {self.model_name}") except ImportError as e: - logger.warning(f"TorchServe backend not available ({e}), falling back to CPU") + logger.warning( + f"TorchServe backend not available ({e}), falling back to CPU" + ) self._init_cpu_backend() def _init_cpu_backend(self): @@ -406,6 +429,7 @@ def _init_cpu_backend(self): # Temporarily bypass hooks from rm_abstract.auto import disable_hooks, enable_hooks + disable_hooks() try: @@ -447,7 +471,7 @@ def generate( self, prompts: Union[str, List[str]], sampling_params: Optional[SamplingParams] = None, - **kwargs + **kwargs, ) -> List[RequestOutput]: """ Generate completions for the given prompts. @@ -469,16 +493,15 @@ def generate( # Route to appropriate generation method based on backend if self._backend in ("triton", "torchserve"): return self._generate_engine_style(prompts, sampling_params, **kwargs) - elif self._backend in ("vllm", "rbln") and hasattr(self._backend_llm, "generate"): + elif self._backend in ("vllm", "rbln") and hasattr( + self._backend_llm, "generate" + ): return self._generate_vllm_style(prompts, sampling_params, **kwargs) else: return self._generate_transformers_style(prompts, sampling_params, **kwargs) def _generate_vllm_style( - self, - prompts: List[str], - sampling_params: SamplingParams, - **kwargs + self, prompts: List[str], sampling_params: SamplingParams, **kwargs ) -> List[RequestOutput]: """Generate using vLLM-style API.""" # Convert to vLLM SamplingParams if available @@ -486,7 +509,9 @@ def _generate_vllm_style( vllm_params = self._sampling_params_cls(**sampling_params.to_dict()) outputs = self._backend_llm.generate(prompts, vllm_params, **kwargs) else: - outputs = self._backend_llm.generate(prompts, sampling_params.to_dict(), **kwargs) + outputs = self._backend_llm.generate( + prompts, sampling_params.to_dict(), **kwargs + ) # Convert to our RequestOutput format results = [] @@ -502,28 +527,29 @@ def _generate_vllm_style( ) for j, o in enumerate(output.outputs) ] - results.append(RequestOutput( - request_id=getattr(output, "request_id", str(i)), - prompt=prompts[i], - prompt_token_ids=getattr(output, "prompt_token_ids", []), - outputs=completions, - )) + results.append( + RequestOutput( + request_id=getattr(output, "request_id", str(i)), + prompt=prompts[i], + prompt_token_ids=getattr(output, "prompt_token_ids", []), + outputs=completions, + ) + ) else: # Simple string output - results.append(RequestOutput( - request_id=str(i), - prompt=prompts[i], - prompt_token_ids=[], - outputs=[CompletionOutput(index=0, text=str(output))], - )) + results.append( + RequestOutput( + request_id=str(i), + prompt=prompts[i], + prompt_token_ids=[], + outputs=[CompletionOutput(index=0, text=str(output))], + ) + ) return results def _generate_transformers_style( - self, - prompts: List[str], - sampling_params: SamplingParams, - **kwargs + self, prompts: List[str], sampling_params: SamplingParams, **kwargs ) -> List[RequestOutput]: """Generate using transformers-style API.""" import torch @@ -541,39 +567,40 @@ def _generate_transformers_style( top_k=sampling_params.top_k if sampling_params.top_k > 0 else None, do_sample=sampling_params.temperature > 0, num_return_sequences=sampling_params.n, - **kwargs + **kwargs, ) # Decode outputs completions = [] for j, ids in enumerate(output_ids): # Remove prompt tokens - new_ids = ids[inputs.input_ids.shape[1]:] + new_ids = ids[inputs.input_ids.shape[1] :] text = self._tokenizer.decode( new_ids, skip_special_tokens=sampling_params.skip_special_tokens, ) - completions.append(CompletionOutput( - index=j, - text=text, - token_ids=new_ids.tolist(), - finish_reason="stop", - )) - - results.append(RequestOutput( - request_id=str(i), - prompt=prompt, - prompt_token_ids=inputs.input_ids[0].tolist(), - outputs=completions, - )) + completions.append( + CompletionOutput( + index=j, + text=text, + token_ids=new_ids.tolist(), + finish_reason="stop", + ) + ) + + results.append( + RequestOutput( + request_id=str(i), + prompt=prompt, + prompt_token_ids=inputs.input_ids[0].tolist(), + outputs=completions, + ) + ) return results def _generate_engine_style( - self, - prompts: List[str], - sampling_params: SamplingParams, - **kwargs + self, prompts: List[str], sampling_params: SamplingParams, **kwargs ) -> List[RequestOutput]: """Generate using ServingEngine (Triton/TorchServe) API.""" if self._engine_instance is None: @@ -587,7 +614,7 @@ def _generate_engine_style( temperature=sampling_params.temperature, top_p=sampling_params.top_p, max_tokens=sampling_params.max_tokens, - **kwargs + **kwargs, ) # Convert to RequestOutput format @@ -598,17 +625,21 @@ def _generate_engine_style( else: text = str(output) - results.append(RequestOutput( - request_id=str(i), - prompt=prompt, - prompt_token_ids=[], - outputs=[CompletionOutput( - index=0, - text=text, - token_ids=[], - finish_reason="stop", - )], - )) + results.append( + RequestOutput( + request_id=str(i), + prompt=prompt, + prompt_token_ids=[], + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=[], + finish_reason="stop", + ) + ], + ) + ) return results diff --git a/src/rm_abstract/core/__init__.py b/src/rm_abstract/core/__init__.py index 7892a83..e753930 100644 --- a/src/rm_abstract/core/__init__.py +++ b/src/rm_abstract/core/__init__.py @@ -1,7 +1,7 @@ """Core module - Core abstraction interfaces""" -from .controller import DeviceFlowController from .config import Config +from .controller import DeviceFlowController __all__ = [ "DeviceFlowController", diff --git a/src/rm_abstract/core/config.py b/src/rm_abstract/core/config.py index ce2b724..cd731d8 100644 --- a/src/rm_abstract/core/config.py +++ b/src/rm_abstract/core/config.py @@ -2,10 +2,10 @@ Configuration management module """ +import os from dataclasses import dataclass, field -from typing import Any, Dict, Optional from pathlib import Path -import os +from typing import Any, Dict, Optional @dataclass diff --git a/src/rm_abstract/core/controller.py b/src/rm_abstract/core/controller.py index bed8482..facf51a 100644 --- a/src/rm_abstract/core/controller.py +++ b/src/rm_abstract/core/controller.py @@ -5,12 +5,12 @@ Uses the plugin registry system for resource discovery. """ -from typing import Any, Dict, Optional, Type import logging +from typing import Any, Dict, Optional, Type +from ..registry import PluginType, discover_resources, get_registry +from ..resources.base import Resource from .config import Config -from ..resources.base import Resource, ResourceInfo -from ..registry import PluginRegistry, PluginType, get_registry, discover_resources logger = logging.getLogger(__name__) diff --git a/src/rm_abstract/engines/__init__.py b/src/rm_abstract/engines/__init__.py index 6c46c48..b7bf8fd 100644 --- a/src/rm_abstract/engines/__init__.py +++ b/src/rm_abstract/engines/__init__.py @@ -12,14 +12,14 @@ """ from .base import ( + DeviceTarget, + ServingConfig, ServingEngine, ServingEngineType, - ServingConfig, - DeviceTarget, ) -from .vllm import VLLMEngine -from .triton import TritonEngine from .torchserve import TorchServeEngine +from .triton import TritonEngine +from .vllm import VLLMEngine __all__ = [ "ServingEngine", diff --git a/src/rm_abstract/engines/base.py b/src/rm_abstract/engines/base.py index c6a3f8b..61f34d4 100644 --- a/src/rm_abstract/engines/base.py +++ b/src/rm_abstract/engines/base.py @@ -4,11 +4,11 @@ Abstract base class for all serving engines (vLLM, Triton, TorchServe, etc.) """ +import logging from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Optional, Union -import logging +from typing import Any, Dict, List, Optional logger = logging.getLogger(__name__) diff --git a/src/rm_abstract/engines/model_converter.py b/src/rm_abstract/engines/model_converter.py index 90615f0..a3e27f1 100644 --- a/src/rm_abstract/engines/model_converter.py +++ b/src/rm_abstract/engines/model_converter.py @@ -6,12 +6,12 @@ - MAR (Model Archive) for TorchServe """ +import logging import os import shutil import tempfile -import logging from pathlib import Path -from typing import Optional, Dict, Any +from typing import Optional logger = logging.getLogger(__name__) @@ -72,14 +72,14 @@ def to_onnx( output_path.mkdir(parents=True, exist_ok=True) try: - from transformers import AutoModelForCausalLM, AutoTokenizer import torch + from transformers import AutoModelForCausalLM, AutoTokenizer print(f"[ONNX] Converting {model_name} to ONNX format...", flush=True) logger.info(f"Converting {model_name} to ONNX...") # Load model and tokenizer - print(f"[ONNX] Loading model and tokenizer...", flush=True) + print("[ONNX] Loading model and tokenizer...", flush=True) tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) model.eval() @@ -102,7 +102,7 @@ def to_onnx( dummy_input = {k: v.cuda() for k, v in dummy_input.items()} # Export to ONNX - print(f"[ONNX] Exporting to ONNX (this may take a while)...", flush=True) + print("[ONNX] Exporting to ONNX (this may take a while)...", flush=True) torch.onnx.export( model, (dummy_input["input_ids"], dummy_input["attention_mask"]), @@ -163,8 +163,8 @@ def to_mar( return str(mar_file) try: + import torch # noqa: F401 from transformers import AutoModelForCausalLM, AutoTokenizer - import torch print(f"[MAR] Converting {model_name} to MAR format...", flush=True) logger.info(f"Converting {model_name} to MAR format...") @@ -174,14 +174,14 @@ def to_mar( tmp_path = Path(tmp_dir) # Save model and tokenizer - print(f"[MAR] Loading model and tokenizer...", flush=True) + print("[MAR] Loading model and tokenizer...", flush=True) tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) model_dir = tmp_path / "model" model_dir.mkdir() - print(f"[MAR] Saving model files...", flush=True) + print("[MAR] Saving model files...", flush=True) tokenizer.save_pretrained(str(model_dir)) model.save_pretrained(str(model_dir)) @@ -190,7 +190,7 @@ def to_mar( handler_file.write_text(self._get_handler_code(handler)) # Create MAR using torch-model-archiver - print(f"[MAR] Creating MAR archive...", flush=True) + print("[MAR] Creating MAR archive...", flush=True) self._create_mar_archive( model_name=safe_name, model_dir=str(model_dir), @@ -295,12 +295,18 @@ def _create_mar_archive( cmd = [ "torch-model-archiver", - "--model-name", model_name, - "--version", "1.0", - "--serialized-file", os.path.join(model_dir, "pytorch_model.bin"), - "--handler", handler_file, - "--extra-files", f"{model_dir}/config.json,{model_dir}/tokenizer.json,{model_dir}/vocab.json", - "--export-path", output_path, + "--model-name", + model_name, + "--version", + "1.0", + "--serialized-file", + os.path.join(model_dir, "pytorch_model.bin"), + "--handler", + handler_file, + "--extra-files", + f"{model_dir}/config.json,{model_dir}/tokenizer.json,{model_dir}/vocab.json", + "--export-path", + output_path, "--force", ] @@ -312,7 +318,15 @@ def _create_mar_archive( cmd[cmd.index("--serialized-file") + 1] = str(model_files[0]) else: # No serialized file, use handler-only mode - cmd = [c for c in cmd if c not in ["--serialized-file", cmd[cmd.index("--serialized-file") + 1]]] + cmd = [ + c + for c in cmd + if c + not in [ + "--serialized-file", + cmd[cmd.index("--serialized-file") + 1], + ] + ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: @@ -352,7 +366,7 @@ def create_triton_config( shutil.copy(onnx_src, version_dir / "model.onnx") # Create config.pbtxt - config_content = f'''name: "{safe_name}" + config_content = f"""name: "{safe_name}" platform: "{platform}" max_batch_size: {max_batch_size} @@ -388,7 +402,7 @@ def create_triton_config( preferred_batch_size: [1, 2, 4, 8] max_queue_delay_microseconds: 100 }} -''' +""" config_file = model_dir / "config.pbtxt" config_file.write_text(config_content) diff --git a/src/rm_abstract/engines/torchserve.py b/src/rm_abstract/engines/torchserve.py index 9b971d7..8e98242 100644 --- a/src/rm_abstract/engines/torchserve.py +++ b/src/rm_abstract/engines/torchserve.py @@ -5,17 +5,16 @@ Supports automatic MAR packaging and server management. """ -import os import atexit -from typing import Any, Dict, List, Optional -from dataclasses import dataclass import logging +import os import subprocess import time +from dataclasses import dataclass from pathlib import Path -import tempfile +from typing import Any, Dict, List, Optional -from .base import ServingEngine, ServingEngineType, ServingConfig, DeviceTarget +from .base import DeviceTarget, ServingConfig, ServingEngine, ServingEngineType logger = logging.getLogger(__name__) @@ -101,6 +100,7 @@ def is_server_running(cls, port: int = 8080) -> bool: """Check if TorchServe server is already running""" try: import requests + response = requests.get(f"http://localhost:{port}/ping", timeout=2) return response.status_code == 200 except Exception: @@ -154,6 +154,7 @@ def _prepare_model(self, model_name: str) -> None: # Copy MAR to model store import shutil + safe_name = model_name.replace("/", "_") dest_mar = self._model_store_path / f"{safe_name}.mar" if Path(mar_path).exists() and not dest_mar.exists(): @@ -164,6 +165,7 @@ def _prepare_model(self, model_name: str) -> None: # Load tokenizer for inference from transformers import AutoTokenizer + self._tokenizer = AutoTokenizer.from_pretrained(model_name) if self._tokenizer.pad_token is None: self._tokenizer.pad_token = self._tokenizer.eos_token @@ -190,8 +192,10 @@ def start_server(self) -> None: cmd = [ "torchserve", "--start", - "--model-store", model_store, - "--ts-config", str(self._config_file), + "--model-store", + model_store, + "--ts-config", + str(self._config_file), "--ncs", # No config snapshot ] @@ -285,7 +289,7 @@ def _wait_for_ready(self, timeout: int = 30) -> None: url = f"http://localhost:{self.ts_config.inference_port}/ping" start_time = time.time() - print(f"[TorchServe] Starting server...", flush=True) + print("[TorchServe] Starting server...", flush=True) check_interval = 1 dots = 0 @@ -304,10 +308,14 @@ def _wait_for_ready(self, timeout: int = 30) -> None: # Progress indicator dots = (dots + 1) % 4 progress = "." * (dots + 1) - print(f"\r[TorchServe] Waiting for server{progress.ljust(4)} ({elapsed}/{timeout}s)", end="", flush=True) + print( + f"\r[TorchServe] Waiting for server{progress.ljust(4)} ({elapsed}/{timeout}s)", + end="", + flush=True, + ) time.sleep(check_interval) - print(f"\n[TorchServe] Timeout!", flush=True) + print("\n[TorchServe] Timeout!", flush=True) raise TimeoutError(f"TorchServe did not become ready within {timeout}s") def _register_model(self) -> None: @@ -389,7 +397,9 @@ def infer(self, inputs: Any, **kwargs) -> Any: import requests safe_name = self.config.model_name.replace("/", "_") - url = f"http://localhost:{self.ts_config.inference_port}/predictions/{safe_name}" + url = ( + f"http://localhost:{self.ts_config.inference_port}/predictions/{safe_name}" + ) # Prepare input if isinstance(inputs, str): @@ -400,10 +410,12 @@ def infer(self, inputs: Any, **kwargs) -> Any: data = {"data": inputs} # Add generation parameters - data.update({ - "max_tokens": kwargs.get("max_tokens", 100), - "temperature": kwargs.get("temperature", 0.7), - }) + data.update( + { + "max_tokens": kwargs.get("max_tokens", 100), + "temperature": kwargs.get("temperature", 0.7), + } + ) try: response = requests.post(url, json=data, timeout=60) @@ -424,6 +436,7 @@ def get_metrics(self) -> Dict[str, Any]: if self._is_running: try: import requests + response = requests.get( f"http://localhost:{self.ts_config.metrics_port}/metrics", timeout=5, diff --git a/src/rm_abstract/engines/triton.py b/src/rm_abstract/engines/triton.py index ac82228..4e6d34a 100644 --- a/src/rm_abstract/engines/triton.py +++ b/src/rm_abstract/engines/triton.py @@ -5,16 +5,16 @@ Supports automatic model conversion and Docker container management. """ -import os import atexit -from typing import Any, Dict, List, Optional -from dataclasses import dataclass, field import logging +import os import subprocess import time +from dataclasses import dataclass from pathlib import Path +from typing import Any, Dict, List, Optional -from .base import ServingEngine, ServingEngineType, ServingConfig, DeviceTarget +from .base import DeviceTarget, ServingConfig, ServingEngine, ServingEngineType logger = logging.getLogger(__name__) @@ -102,9 +102,9 @@ def is_server_running(cls, port: int = 8000) -> bool: """Check if Triton server is already running""" try: import requests + response = requests.get( - f"http://localhost:{port}/v2/health/ready", - timeout=2 + f"http://localhost:{port}/v2/health/ready", timeout=2 ) return response.status_code == 200 except Exception: @@ -146,7 +146,7 @@ def _prepare_model(self, model_name: str) -> None: logger.info(f"Converting {model_name} to ONNX...") onnx_path = converter.to_onnx( model_name, - device="cuda" if self.config.device == DeviceTarget.GPU else "cpu" + device="cuda" if self.config.device == DeviceTarget.GPU else "cpu", ) # Create model repository @@ -165,6 +165,7 @@ def _prepare_model(self, model_name: str) -> None: # Load tokenizer for inference from transformers import AutoTokenizer + self._tokenizer = AutoTokenizer.from_pretrained(model_name) if self._tokenizer.pad_token is None: self._tokenizer.pad_token = self._tokenizer.eos_token @@ -185,11 +186,17 @@ def start_server(self) -> None: # Build Docker run command cmd = [ - "docker", "run", "-d", - "--name", self.triton_config.container_name, - "-p", f"{self.triton_config.http_port}:8000", - "-p", f"{self.triton_config.grpc_port}:8001", - "-p", f"{self.triton_config.metrics_port}:8002", + "docker", + "run", + "-d", + "--name", + self.triton_config.container_name, + "-p", + f"{self.triton_config.http_port}:8000", + "-p", + f"{self.triton_config.grpc_port}:8001", + "-p", + f"{self.triton_config.metrics_port}:8002", "--shm-size=1g", ] @@ -199,13 +206,16 @@ def start_server(self) -> None: # Add model repository volume model_repo = self.triton_config.model_repository or "/models" - cmd.extend([ - "-v", f"{model_repo}:/models", - self.triton_config.docker_image, - "tritonserver", - "--model-repository=/models", - "--strict-model-config=false", # Allow auto-completion - ]) + cmd.extend( + [ + "-v", + f"{model_repo}:/models", + self.triton_config.docker_image, + "tritonserver", + "--model-repository=/models", + "--strict-model-config=false", # Allow auto-completion + ] + ) try: # Remove existing container if exists @@ -268,13 +278,13 @@ def _cleanup(self) -> None: def _wait_for_ready(self, timeout: int = 60) -> None: """Wait for Triton to be ready""" + import requests - import sys url = f"http://localhost:{self.triton_config.http_port}/v2/health/ready" start_time = time.time() - print(f"[Triton] Starting Docker container...", flush=True) + print("[Triton] Starting Docker container...", flush=True) check_interval = 2 dots = 0 @@ -295,7 +305,7 @@ def _wait_for_ready(self, timeout: int = 60) -> None: capture_output=True, text=True, ) - print(f"\n[Triton] Container failed!", flush=True) + print("\n[Triton] Container failed!", flush=True) raise RuntimeError(f"Triton container exited: {logs.stderr}") try: @@ -310,10 +320,14 @@ def _wait_for_ready(self, timeout: int = 60) -> None: # Progress indicator dots = (dots + 1) % 4 progress = "." * (dots + 1) - print(f"\r[Triton] Waiting for server{progress.ljust(4)} ({elapsed}/{timeout}s)", end="", flush=True) + print( + f"\r[Triton] Waiting for server{progress.ljust(4)} ({elapsed}/{timeout}s)", + end="", + flush=True, + ) time.sleep(check_interval) - print(f"\n[Triton] Timeout!", flush=True) + print("\n[Triton] Timeout!", flush=True) raise TimeoutError(f"Triton server did not become ready within {timeout}s") def infer(self, inputs: Any, **kwargs) -> Any: @@ -331,8 +345,8 @@ def infer(self, inputs: Any, **kwargs) -> Any: raise RuntimeError("Triton server not running. Call load_model() first.") try: - import tritonclient.http as httpclient import numpy as np + import tritonclient.http as httpclient except ImportError: raise ImportError( "tritonclient not installed. " @@ -350,6 +364,7 @@ def infer(self, inputs: Any, **kwargs) -> Any: if isinstance(inputs, str): if self._tokenizer is None: from transformers import AutoTokenizer + self._tokenizer = AutoTokenizer.from_pretrained(self.config.model_name) if self._tokenizer.pad_token is None: self._tokenizer.pad_token = self._tokenizer.eos_token @@ -409,8 +424,7 @@ def infer(self, inputs: Any, **kwargs) -> Any: # Decode output generated_text = self._tokenizer.decode( - generated_ids[len(tokens["input_ids"][0]):], - skip_special_tokens=True + generated_ids[len(tokens["input_ids"][0]) :], skip_special_tokens=True ) return {"generated_text": generated_text} @@ -426,6 +440,7 @@ def get_metrics(self) -> Dict[str, Any]: if self._is_running: try: import requests + response = requests.get( f"http://localhost:{self.triton_config.metrics_port}/metrics", timeout=5, diff --git a/src/rm_abstract/engines/vllm.py b/src/rm_abstract/engines/vllm.py index ba93b32..276c1da 100644 --- a/src/rm_abstract/engines/vllm.py +++ b/src/rm_abstract/engines/vllm.py @@ -5,10 +5,10 @@ Supports GPU (CUDA) and NPU (RBLN) backends. """ -from typing import Any, Dict, List, Optional, Union import logging +from typing import Any, List, Union -from .base import ServingEngine, ServingEngineType, ServingConfig, DeviceTarget +from .base import DeviceTarget, ServingConfig, ServingEngine, ServingEngineType logger = logging.getLogger(__name__) @@ -45,7 +45,7 @@ def __init__(self, config: ServingConfig): def is_available(cls) -> bool: """Check if vLLM is available""" try: - import vllm + import vllm # noqa: F401 return True except ImportError: @@ -67,12 +67,12 @@ def supported_devices(cls) -> List[DeviceTarget]: # Check RBLN NPU try: - import vllm_rbln + import vllm_rbln # noqa: F401 devices.append(DeviceTarget.NPU_RBLN) except ImportError: try: - import rebel + import rebel # noqa: F401 devices.append(DeviceTarget.NPU_RBLN) except ImportError: @@ -103,12 +103,11 @@ def load_model(self, model_name_or_path: str, **kwargs) -> Any: # Handle RBLN NPU device if self.config.device == DeviceTarget.NPU_RBLN: try: - import vllm_rbln + import vllm_rbln # noqa: F401 llm_kwargs["device"] = "rbln" except ImportError: try: - from vllm.config import DeviceConfig llm_kwargs["device"] = "rbln" except Exception: diff --git a/src/rm_abstract/exceptions.py b/src/rm_abstract/exceptions.py index 9621ee9..0b75435 100644 --- a/src/rm_abstract/exceptions.py +++ b/src/rm_abstract/exceptions.py @@ -4,11 +4,12 @@ Provides clear, actionable error messages for common failure scenarios. """ -from typing import Optional, List +from typing import List, Optional class RMAbstractError(Exception): """Base exception for all RM Abstract Layer errors""" + pass @@ -16,17 +17,20 @@ class RMAbstractError(Exception): # Initialization Errors # ============================================================ + class InitializationError(RMAbstractError): """Error during initialization""" + pass class NotInitializedError(RMAbstractError): """RM Abstract Layer not initialized""" - + def __init__(self, message: Optional[str] = None): super().__init__( - message or "RM Abstract Layer not initialized. Call rm_abstract.init() first." + message + or "RM Abstract Layer not initialized. Call rm_abstract.init() first." ) @@ -34,44 +38,52 @@ def __init__(self, message: Optional[str] = None): # Backend Errors # ============================================================ + class BackendError(RMAbstractError): """Base class for backend-related errors""" + pass class BackendNotAvailableError(BackendError): """Requested backend is not available""" - - def __init__(self, backend_name: str, reason: Optional[str] = None, install_hint: Optional[str] = None): + + def __init__( + self, + backend_name: str, + reason: Optional[str] = None, + install_hint: Optional[str] = None, + ): self.backend_name = backend_name self.reason = reason self.install_hint = install_hint - + message = f"Backend '{backend_name}' is not available" if reason: message += f": {reason}" if install_hint: message += f"\nInstall with: {install_hint}" - + super().__init__(message) class BackendInitError(BackendError): """Failed to initialize backend""" + pass class DeviceNotFoundError(BackendError): """Specified device not found""" - + def __init__(self, device: str, available_devices: Optional[List[str]] = None): self.device = device self.available_devices = available_devices - + message = f"Device '{device}' not found" if available_devices: message += f". Available devices: {available_devices}" - + super().__init__(message) @@ -79,36 +91,37 @@ def __init__(self, device: str, available_devices: Optional[List[str]] = None): # Model Errors # ============================================================ + class ModelError(RMAbstractError): """Base class for model-related errors""" + pass class ModelLoadError(ModelError): """Failed to load model""" - + def __init__(self, model_name: str, reason: Optional[str] = None): self.model_name = model_name self.reason = reason - + message = f"Failed to load model '{model_name}'" if reason: message += f": {reason}" - + super().__init__(message) class ModelNotLoadedError(ModelError): """Model not loaded""" - + def __init__(self, message: Optional[str] = None): - super().__init__( - message or "Model not loaded. Call load_model() first." - ) + super().__init__(message or "Model not loaded. Call load_model() first.") class InferenceError(ModelError): """Error during inference""" + pass @@ -116,48 +129,50 @@ class InferenceError(ModelError): # Serving Errors # ============================================================ + class ServingError(RMAbstractError): """Base class for serving-related errors""" + pass class ServerStartError(ServingError): """Failed to start server""" - + def __init__(self, engine_name: str, reason: Optional[str] = None): self.engine_name = engine_name self.reason = reason - + message = f"Failed to start {engine_name} server" if reason: message += f": {reason}" - + super().__init__(message) class ServerNotRunningError(ServingError): """Server is not running""" - + def __init__(self, engine_name: Optional[str] = None): message = "Server is not running" if engine_name: message = f"{engine_name} server is not running" message += ". Call start_server() first." - + super().__init__(message) class ServerConnectionError(ServingError): """Cannot connect to server""" - + def __init__(self, url: str, reason: Optional[str] = None): self.url = url self.reason = reason - + message = f"Cannot connect to server at {url}" if reason: message += f": {reason}" - + super().__init__(message) @@ -165,23 +180,25 @@ def __init__(self, url: str, reason: Optional[str] = None): # Configuration Errors # ============================================================ + class ConfigurationError(RMAbstractError): """Invalid configuration""" + pass class InvalidDeviceError(ConfigurationError): """Invalid device specification""" - + def __init__(self, device: str, valid_formats: Optional[List[str]] = None): self.device = device - + message = f"Invalid device specification: '{device}'" if valid_formats: message += f". Valid formats: {valid_formats}" else: message += ". Valid formats: 'auto', 'gpu:0', 'cpu', 'rbln:0'" - + super().__init__(message) @@ -189,34 +206,36 @@ def __init__(self, device: str, valid_formats: Optional[List[str]] = None): # Docker Errors (for Triton) # ============================================================ + class DockerError(RMAbstractError): """Base class for Docker-related errors""" + pass class DockerNotAvailableError(DockerError): """Docker is not available""" - + def __init__(self, reason: Optional[str] = None): message = "Docker is not available" if reason: message += f": {reason}" message += "\nEnsure Docker is installed and running, and you have permission to use it." message += "\nTry: sudo usermod -aG docker $USER && newgrp docker" - + super().__init__(message) class DockerImageNotFoundError(DockerError): """Docker image not found""" - + def __init__(self, image: str, build_hint: Optional[str] = None): self.image = image - + message = f"Docker image '{image}' not found" if build_hint: message += f"\nBuild with: {build_hint}" - + super().__init__(message) @@ -224,21 +243,22 @@ def __init__(self, image: str, build_hint: Optional[str] = None): # Dependency Errors # ============================================================ + class DependencyError(RMAbstractError): """Missing or incompatible dependency""" + pass class PackageNotInstalledError(DependencyError): """Required package not installed""" - + def __init__(self, package: str, install_cmd: Optional[str] = None): self.package = package self.install_cmd = install_cmd - + message = f"Required package '{package}' is not installed" if install_cmd: message += f"\nInstall with: {install_cmd}" - - super().__init__(message) + super().__init__(message) diff --git a/src/rm_abstract/llm.py b/src/rm_abstract/llm.py index 93e9343..34d1ae4 100644 --- a/src/rm_abstract/llm.py +++ b/src/rm_abstract/llm.py @@ -17,8 +17,8 @@ outputs = llm.generate(["Hello, I am"], SamplingParams(max_tokens=50)) """ -import os import logging +import os from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Union @@ -272,7 +272,9 @@ def _init_vllm_gpu_backend(self): self._init_transformers_backend("nvidia") except Exception as e: # Fallback to transformers on GPU memory errors, etc. - logger.warning(f"vLLM initialization failed ({e}), falling back to transformers GPU") + logger.warning( + f"vLLM initialization failed ({e}), falling back to transformers GPU" + ) self._init_transformers_backend("nvidia") def _init_vllm_rbln_backend(self): @@ -329,8 +331,8 @@ def _init_furiosa_backend(self): def _init_triton_backend(self, device: str): """Initialize Triton Inference Server backend with auto-start.""" try: - from .engines.triton import TritonEngine, TritonConfig - from .engines.base import ServingConfig, ServingEngineType, DeviceTarget + from .engines.base import DeviceTarget, ServingConfig, ServingEngineType + from .engines.triton import TritonConfig, TritonEngine # Check if Docker is available if not TritonEngine.is_available(): @@ -364,20 +366,26 @@ def _init_triton_backend(self, device: str): self._backend_llm = TritonEngine(config, triton_config) # load_model will auto-convert to ONNX and start Docker container self._backend_llm.load_model(self.model_name) - logger.info(f"Initialized Triton backend with auto-start: {self.model_name}") + logger.info( + f"Initialized Triton backend with auto-start: {self.model_name}" + ) except ImportError as e: - logger.warning(f"Triton backend not available ({e}), falling back to transformers") + logger.warning( + f"Triton backend not available ({e}), falling back to transformers" + ) self._init_transformers_backend(device) except Exception as e: - logger.warning(f"Failed to start Triton ({e}), falling back to transformers") + logger.warning( + f"Failed to start Triton ({e}), falling back to transformers" + ) self._init_transformers_backend(device) def _init_torchserve_backend(self, device: str): """Initialize TorchServe backend with auto-start.""" try: - from .engines.torchserve import TorchServeEngine, TorchServeConfig - from .engines.base import ServingConfig, ServingEngineType, DeviceTarget + from .engines.base import DeviceTarget, ServingConfig, ServingEngineType + from .engines.torchserve import TorchServeConfig, TorchServeEngine # Check if TorchServe is installed if not TorchServeEngine.is_available(): @@ -410,7 +418,9 @@ def _init_torchserve_backend(self, device: str): self._backend_llm = TorchServeEngine(config, ts_config) # load_model will auto-package MAR and start server self._backend_llm.load_model(self.model_name) - logger.info(f"Initialized TorchServe backend with auto-start: {self.model_name}") + logger.info( + f"Initialized TorchServe backend with auto-start: {self.model_name}" + ) except ImportError as e: logger.warning( @@ -418,7 +428,9 @@ def _init_torchserve_backend(self, device: str): ) self._init_transformers_backend(device) except Exception as e: - logger.warning(f"Failed to start TorchServe ({e}), falling back to transformers") + logger.warning( + f"Failed to start TorchServe ({e}), falling back to transformers" + ) self._init_transformers_backend(device) def _init_cpu_backend(self): @@ -428,8 +440,8 @@ def _init_cpu_backend(self): def _init_transformers_backend(self, device: str): """Initialize transformers backend for CPU or GPU.""" try: - from transformers import AutoModelForCausalLM, AutoTokenizer import torch + from transformers import AutoModelForCausalLM, AutoTokenizer # Determine if using GPU or CPU use_gpu = device in ("nvidia", "gpu", "cuda") and torch.cuda.is_available() diff --git a/src/rm_abstract/registry/__init__.py b/src/rm_abstract/registry/__init__.py index 242c05e..03746de 100644 --- a/src/rm_abstract/registry/__init__.py +++ b/src/rm_abstract/registry/__init__.py @@ -9,15 +9,15 @@ - discovery.py: Auto-discovery of available plugins """ +from .discovery import discover_all, discover_engines, discover_resources from .plugin_registry import ( - PluginRegistry, PluginInfo, + PluginRegistry, PluginType, + get_plugin, get_registry, register_plugin, - get_plugin, ) -from .discovery import discover_resources, discover_engines, discover_all __all__ = [ "PluginRegistry", diff --git a/src/rm_abstract/registry/discovery.py b/src/rm_abstract/registry/discovery.py index 3dc36ae..9ca7ca2 100644 --- a/src/rm_abstract/registry/discovery.py +++ b/src/rm_abstract/registry/discovery.py @@ -4,12 +4,12 @@ Auto-discovery and registration of available plugins. """ -from typing import List, Optional import logging +from typing import List, Optional from .plugin_registry import ( - PluginRegistry, PluginInfo, + PluginRegistry, PluginType, get_registry, ) diff --git a/src/rm_abstract/registry/plugin_registry.py b/src/rm_abstract/registry/plugin_registry.py index 92adeea..9bbd954 100644 --- a/src/rm_abstract/registry/plugin_registry.py +++ b/src/rm_abstract/registry/plugin_registry.py @@ -5,10 +5,10 @@ Eliminates hardcoded backend/engine names throughout the codebase. """ +import logging from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, Dict, List, Optional, Type -import logging logger = logging.getLogger(__name__) diff --git a/src/rm_abstract/resources/__init__.py b/src/rm_abstract/resources/__init__.py index 1d18ab9..136d19e 100644 --- a/src/rm_abstract/resources/__init__.py +++ b/src/rm_abstract/resources/__init__.py @@ -11,7 +11,7 @@ - cpu/: CPU fallback resource """ -from .base import Resource, ResourceInfo, ResourceCapability +from .base import Resource, ResourceCapability, ResourceInfo __all__ = [ "Resource", diff --git a/src/rm_abstract/resources/base.py b/src/rm_abstract/resources/base.py index b41b0cc..d954031 100644 --- a/src/rm_abstract/resources/base.py +++ b/src/rm_abstract/resources/base.py @@ -4,11 +4,11 @@ Defines the common interface for all hardware resources. """ +import logging from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Optional, Type -import logging +from typing import Any, Dict, Optional logger = logging.getLogger(__name__) @@ -129,9 +129,7 @@ def initialize(self) -> None: ... @abstractmethod - def prepare_model( - self, model: Any, config: Optional[Dict[str, Any]] = None - ) -> Any: + def prepare_model(self, model: Any, config: Optional[Dict[str, Any]] = None) -> Any: """ Prepare model for execution on this resource. diff --git a/src/rm_abstract/resources/cpu/default.py b/src/rm_abstract/resources/cpu/default.py index e38234e..295601e 100644 --- a/src/rm_abstract/resources/cpu/default.py +++ b/src/rm_abstract/resources/cpu/default.py @@ -4,13 +4,13 @@ Fallback resource using PyTorch CPU inference. """ -from typing import Any, Dict, Optional -import platform -import os import logging +import os +import platform +from typing import Any, Dict, Optional -from ..base import Resource, ResourceType, ResourceInfo, ResourceCapability from ...common.input_adapter import InputAdapter +from ..base import Resource, ResourceCapability, ResourceInfo, ResourceType logger = logging.getLogger(__name__) @@ -43,7 +43,7 @@ def __init__(self, device_id: int = 0, **kwargs): def is_available() -> bool: """CPU is always available if PyTorch is installed""" try: - import torch + import torch # noqa: F401 return True except ImportError: @@ -57,9 +57,7 @@ def initialize(self) -> None: self._initialized = True logger.info("CPU resource initialized") - def prepare_model( - self, model: Any, config: Optional[Dict[str, Any]] = None - ) -> Any: + def prepare_model(self, model: Any, config: Optional[Dict[str, Any]] = None) -> Any: """ Prepare model for CPU inference. @@ -70,7 +68,6 @@ def prepare_model( Returns: Model on CPU in eval mode """ - import torch if hasattr(model, "to"): model = model.to("cpu") @@ -128,9 +125,7 @@ def execute(self, model: Any, inputs: Any, **kwargs) -> Any: else: return model(exec_inputs, **exec_kwargs) - def _prepare_execution( - self, inputs: Any, kwargs: Dict[str, Any] - ) -> tuple: + def _prepare_execution(self, inputs: Any, kwargs: Dict[str, Any]) -> tuple: """Prepare inputs for execution""" import torch diff --git a/src/rm_abstract/resources/gpu/nvidia.py b/src/rm_abstract/resources/gpu/nvidia.py index 212fbe2..1e4774e 100644 --- a/src/rm_abstract/resources/gpu/nvidia.py +++ b/src/rm_abstract/resources/gpu/nvidia.py @@ -4,12 +4,12 @@ GPU resource implementation using vLLM for high-performance inference. """ -from typing import Any, Dict, Optional import logging +from typing import Any, Dict, Optional -from ..base import Resource, ResourceType, ResourceInfo, ResourceCapability from ...common.input_adapter import InputAdapter from ...common.output_adapter import OutputAdapter +from ..base import Resource, ResourceCapability, ResourceInfo, ResourceType logger = logging.getLogger(__name__) @@ -51,7 +51,7 @@ def is_available() -> bool: return False # Check vLLM import - import vllm + import vllm # noqa: F401 return True except ImportError: @@ -76,9 +76,7 @@ def initialize(self) -> None: logger.error(f"Failed to initialize NVIDIA GPU: {e}") raise - def prepare_model( - self, model: Any, config: Optional[Dict[str, Any]] = None - ) -> Any: + def prepare_model(self, model: Any, config: Optional[Dict[str, Any]] = None) -> Any: """ Prepare model using vLLM. diff --git a/src/rm_abstract/resources/npu/__init__.py b/src/rm_abstract/resources/npu/__init__.py index da87122..009eb20 100644 --- a/src/rm_abstract/resources/npu/__init__.py +++ b/src/rm_abstract/resources/npu/__init__.py @@ -5,7 +5,7 @@ Each vendor has its own file for clear separation. """ -from .rebellions import RebellionsATOM from .furiosa import FuriosaRNGD +from .rebellions import RebellionsATOM __all__ = ["RebellionsATOM", "FuriosaRNGD"] diff --git a/src/rm_abstract/resources/npu/furiosa.py b/src/rm_abstract/resources/npu/furiosa.py index acfec6a..ef03f85 100644 --- a/src/rm_abstract/resources/npu/furiosa.py +++ b/src/rm_abstract/resources/npu/furiosa.py @@ -4,15 +4,15 @@ NPU resource implementation for FuriosaAI RNGD chips. """ -from typing import Any, Dict, Optional -from pathlib import Path -import os import logging +import os +from pathlib import Path +from typing import Any, Dict, Optional -from ..base import NPUResource, ResourceType, ResourceInfo, ResourceCapability from ...common.input_adapter import InputAdapter +from ...common.model_converter import ConversionConfig, ModelConverter from ...common.output_adapter import OutputAdapter -from ...common.model_converter import ModelConverter, ConversionConfig +from ..base import NPUResource, ResourceCapability, ResourceInfo, ResourceType logger = logging.getLogger(__name__) @@ -102,9 +102,7 @@ def initialize(self) -> None: logger.error(f"Failed to initialize FuriosaAI RNGD: {e}") raise - def prepare_model( - self, model: Any, config: Optional[Dict[str, Any]] = None - ) -> Any: + def prepare_model(self, model: Any, config: Optional[Dict[str, Any]] = None) -> Any: """ Prepare model for Furiosa NPU. @@ -124,11 +122,19 @@ def prepare_model( compile_config = {**self.compile_options, **config} # Generate cache key - cache_key = self._model_converter.get_cache_key(model, ConversionConfig(**{ - k: v for k, v in compile_config.items() - if k in ["batch_size", "seq_length", "opset_version"] - })) - cache_path = os.path.join(self.cache_dir, f"{cache_key}.{self.compiled_model_extension}") + cache_key = self._model_converter.get_cache_key( + model, + ConversionConfig( + **{ + k: v + for k, v in compile_config.items() + if k in ["batch_size", "seq_length", "opset_version"] + } + ), + ) + cache_path = os.path.join( + self.cache_dir, f"{cache_key}.{self.compiled_model_extension}" + ) # Check cache if os.path.exists(cache_path): @@ -219,15 +225,13 @@ def execute(self, model: Any, inputs: Any, **kwargs) -> Any: def _execute_forward(self, model: Any, inputs: Any, **kwargs) -> Any: """Execute single forward pass on NPU""" from furiosa import runtime - import numpy as np # Convert inputs to numpy if hasattr(inputs, "numpy"): inputs = inputs.numpy() elif isinstance(inputs, dict): inputs = { - k: v.numpy() if hasattr(v, "numpy") else v - for k, v in inputs.items() + k: v.numpy() if hasattr(v, "numpy") else v for k, v in inputs.items() } # Create runner if needed diff --git a/src/rm_abstract/resources/npu/rebellions.py b/src/rm_abstract/resources/npu/rebellions.py index 15e8b0b..e6965cd 100644 --- a/src/rm_abstract/resources/npu/rebellions.py +++ b/src/rm_abstract/resources/npu/rebellions.py @@ -5,16 +5,16 @@ Supports two modes: vLLM-RBLN and Optimum-RBLN. """ +import logging +import os from enum import Enum -from typing import Any, Dict, Literal, Optional from pathlib import Path -import os -import logging +from typing import Any, Dict, Literal, Optional -from ..base import NPUResource, ResourceType, ResourceInfo, ResourceCapability from ...common.input_adapter import InputAdapter +from ...common.model_converter import ModelConverter from ...common.output_adapter import OutputAdapter -from ...common.model_converter import ModelConverter, ConversionConfig +from ..base import NPUResource, ResourceCapability, ResourceInfo, ResourceType logger = logging.getLogger(__name__) @@ -98,7 +98,7 @@ def is_available() -> bool: def _check_vllm_rbln_static() -> bool: """Check if vLLM-RBLN is available (static version)""" try: - import vllm_rbln + import vllm_rbln # noqa: F401 return True except ImportError: @@ -106,7 +106,7 @@ def _check_vllm_rbln_static() -> bool: # Check rebel SDK + vLLM try: - import rebel + import rebel # noqa: F401 import vllm return True @@ -115,8 +115,8 @@ def _check_vllm_rbln_static() -> bool: # Check rbln SDK + vLLM try: - import rbln - import vllm + import rbln # noqa: F401 + import vllm # noqa: F401 return True except ImportError: @@ -128,12 +128,12 @@ def _check_vllm_rbln_static() -> bool: def _check_optimum_rbln_static() -> bool: """Check if Optimum-RBLN is available (static version)""" try: - from optimum.rbln import RBLNModelForCausalLM + from optimum.rbln import RBLNModelForCausalLM # noqa: F401 return True except ImportError: try: - import optimum_rbln + import optimum_rbln # noqa: F401 return True except ImportError: @@ -200,7 +200,10 @@ def _determine_mode(self) -> RBLNMode: else: if self._requested_mode == RBLNMode.VLLM and not self._check_vllm_rbln(): raise RuntimeError("vLLM-RBLN not available") - if self._requested_mode == RBLNMode.OPTIMUM and not self._check_optimum_rbln(): + if ( + self._requested_mode == RBLNMode.OPTIMUM + and not self._check_optimum_rbln() + ): raise RuntimeError("Optimum-RBLN not available") return self._requested_mode @@ -220,9 +223,7 @@ def initialize(self) -> None: logger.error(f"Failed to initialize Rebellions ATOM: {e}") raise - def prepare_model( - self, model: Any, config: Optional[Dict[str, Any]] = None - ) -> Any: + def prepare_model(self, model: Any, config: Optional[Dict[str, Any]] = None) -> Any: """Prepare model for RBLN NPU""" if not self._initialized: self.initialize() @@ -261,12 +262,11 @@ def _prepare_vllm(self, model: Any, config: Dict[str, Any]) -> Any: # Add RBLN device try: - import vllm_rbln + import vllm_rbln # noqa: F401 llm_kwargs["device"] = "rbln" except ImportError: try: - from vllm.config import DeviceConfig llm_kwargs["device"] = "rbln" except Exception: @@ -389,7 +389,9 @@ def _execute_optimum(self, model: Any, inputs: Any, **kwargs) -> Any: temperature = kwargs.get("temperature", 1.0) top_p = kwargs.get("top_p", 1.0) - input_dict = inputs if isinstance(inputs, dict) else {"input_ids": inputs.input_ids} + input_dict = ( + inputs if isinstance(inputs, dict) else {"input_ids": inputs.input_ids} + ) outputs = model.generate( **input_dict, max_new_tokens=max_new_tokens, @@ -400,7 +402,11 @@ def _execute_optimum(self, model: Any, inputs: Any, **kwargs) -> Any: return outputs else: with torch.no_grad(): - input_dict = inputs if isinstance(inputs, dict) else {"input_ids": inputs.input_ids} + input_dict = ( + inputs + if isinstance(inputs, dict) + else {"input_ids": inputs.input_ids} + ) return model(**input_dict) def compile_model(self, model: Any, **kwargs) -> Any: diff --git a/src/rm_abstract/run.py b/src/rm_abstract/run.py index e8a192e..ac79c5c 100644 --- a/src/rm_abstract/run.py +++ b/src/rm_abstract/run.py @@ -21,10 +21,10 @@ python -m rm_abstract.run --device=rbln serve.py --port 8000 """ -import sys -import os import argparse +import os import runpy +import sys def main(): @@ -38,13 +38,15 @@ def main(): help="Python script to run", ) parser.add_argument( - "--device", "-d", + "--device", + "-d", default="auto", choices=["auto", "cpu", "nvidia", "gpu", "rbln", "furiosa"], help="Device selection (default: auto)", ) parser.add_argument( - "--engine", "-e", + "--engine", + "-e", default="vllm", choices=["vllm", "triton", "torchserve"], help="Serving engine (default: vllm)", @@ -60,7 +62,8 @@ def main(): help="Disable transformers import hooking", ) parser.add_argument( - "--verbose", "-v", + "--verbose", + "-v", action="store_true", help="Enable verbose logging", ) @@ -80,6 +83,7 @@ def main(): # Setup logging if args.verbose: import logging + logging.basicConfig( level=logging.DEBUG, format="[RM Abstract] %(levelname)s: %(message)s", @@ -92,7 +96,9 @@ def main(): sys.argv = [args.script] + script_args # Run the script - print(f"[RM Abstract] Running {args.script} with device={args.device}, engine={args.engine}") + print( + f"[RM Abstract] Running {args.script} with device={args.device}, engine={args.engine}" + ) try: runpy.run_path(args.script, run_name="__main__") diff --git a/tests/conftest.py b/tests/conftest.py index 6150d42..db3e383 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,10 +2,11 @@ Pytest configuration and fixtures """ -import pytest import sys from pathlib import Path +import pytest + # Add src to path sys.path.insert(0, str(Path(__file__).parent.parent / "src")) diff --git a/tests/test_core.py b/tests/test_core.py index 8f11774..e46acab 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -2,8 +2,6 @@ Core functionality tests """ -import pytest - class TestInitialization: """Test rm_abstract initialization""" diff --git a/tests/unit/test_compat.py b/tests/unit/test_compat.py index 63c6a10..68a52b2 100644 --- a/tests/unit/test_compat.py +++ b/tests/unit/test_compat.py @@ -1,8 +1,7 @@ """Tests for compatibility layer (legacy compat module).""" import os -import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import patch class TestCompatSamplingParams: @@ -52,7 +51,7 @@ class TestCompatRequestOutput: def test_request_output_creation(self): """Test creating RequestOutput.""" - from rm_abstract.compat.vllm_compat import RequestOutput, CompletionOutput + from rm_abstract.compat.vllm_compat import CompletionOutput, RequestOutput output = RequestOutput( request_id="0", @@ -104,8 +103,11 @@ def test_device_from_env(self): from rm_abstract.compat import transformers_compat as tc # RM_DEVICE wins over host detection, so the result must not depend on the host. - with patch.dict(os.environ, {"RM_DEVICE": "cpu"}), patch.object( - tc, "_detect_device", side_effect=AssertionError("must not auto-detect") + with ( + patch.dict(os.environ, {"RM_DEVICE": "cpu"}), + patch.object( + tc, "_detect_device", side_effect=AssertionError("must not auto-detect") + ), ): assert tc._resolve_device(tc._get_device()) == "cpu" @@ -119,6 +121,7 @@ def test_auto_device_uses_detection(self): def test_explicit_device(self): """Test explicit device selection.""" from rm_abstract.compat.transformers_compat import _resolve_device + assert _resolve_device("rbln") == "rbln" assert _resolve_device("nvidia:0") == "nvidia" @@ -145,10 +148,10 @@ def test_compat_init_exports(self): """Test compat __init__ exports.""" from rm_abstract.compat import ( LLM, - SamplingParams, - RequestOutput, AutoModelForCausalLM, AutoTokenizer, + RequestOutput, + SamplingParams, ) assert LLM is not None diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 502e4ed..95b4f81 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,6 +1,5 @@ """Config 클래스 테스트""" -import pytest from rm_abstract.core.config import Config diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py index 85808ff..ab3b187 100644 --- a/tests/unit/test_exceptions.py +++ b/tests/unit/test_exceptions.py @@ -3,35 +3,35 @@ """ import pytest + from rm_abstract.exceptions import ( - RMAbstractError, - InitializationError, - NotInitializedError, BackendError, - BackendNotAvailableError, BackendInitError, + BackendNotAvailableError, + ConfigurationError, + DependencyError, DeviceNotFoundError, + DockerError, + DockerNotAvailableError, + InferenceError, + InitializationError, + InvalidDeviceError, ModelError, ModelLoadError, ModelNotLoadedError, - InferenceError, - ServingError, - ServerStartError, - ServerNotRunningError, - ServerConnectionError, - ConfigurationError, - InvalidDeviceError, - DockerError, - DockerNotAvailableError, - DockerImageNotFoundError, - DependencyError, + NotInitializedError, PackageNotInstalledError, + RMAbstractError, + ServerConnectionError, + ServerNotRunningError, + ServerStartError, + ServingError, ) class TestExceptionHierarchy: """Test exception inheritance hierarchy""" - + def test_base_exception(self): """RMAbstractError is base for all exceptions""" assert issubclass(InitializationError, RMAbstractError) @@ -41,19 +41,19 @@ def test_base_exception(self): assert issubclass(ConfigurationError, RMAbstractError) assert issubclass(DockerError, RMAbstractError) assert issubclass(DependencyError, RMAbstractError) - + def test_backend_subclasses(self): """Backend error subclasses""" assert issubclass(BackendNotAvailableError, BackendError) assert issubclass(BackendInitError, BackendError) assert issubclass(DeviceNotFoundError, BackendError) - + def test_model_subclasses(self): """Model error subclasses""" assert issubclass(ModelLoadError, ModelError) assert issubclass(ModelNotLoadedError, ModelError) assert issubclass(InferenceError, ModelError) - + def test_serving_subclasses(self): """Serving error subclasses""" assert issubclass(ServerStartError, ServingError) @@ -63,12 +63,12 @@ def test_serving_subclasses(self): class TestNotInitializedError: """Test NotInitializedError""" - + def test_default_message(self): err = NotInitializedError() assert "not initialized" in str(err).lower() assert "init()" in str(err) - + def test_custom_message(self): err = NotInitializedError("Custom message") assert str(err) == "Custom message" @@ -76,17 +76,17 @@ def test_custom_message(self): class TestBackendNotAvailableError: """Test BackendNotAvailableError""" - + def test_basic(self): err = BackendNotAvailableError("gpu") assert "gpu" in str(err) assert err.backend_name == "gpu" - + def test_with_reason(self): err = BackendNotAvailableError("gpu", reason="CUDA not found") assert "CUDA not found" in str(err) assert err.reason == "CUDA not found" - + def test_with_install_hint(self): err = BackendNotAvailableError("gpu", install_hint="pip install torch") assert "pip install torch" in str(err) @@ -95,12 +95,12 @@ def test_with_install_hint(self): class TestDeviceNotFoundError: """Test DeviceNotFoundError""" - + def test_basic(self): err = DeviceNotFoundError("gpu:5") assert "gpu:5" in str(err) assert err.device == "gpu:5" - + def test_with_available(self): err = DeviceNotFoundError("gpu:5", available_devices=["gpu:0", "gpu:1"]) assert "gpu:0" in str(err) @@ -109,12 +109,12 @@ def test_with_available(self): class TestModelLoadError: """Test ModelLoadError""" - + def test_basic(self): err = ModelLoadError("gpt2") assert "gpt2" in str(err) assert err.model_name == "gpt2" - + def test_with_reason(self): err = ModelLoadError("gpt2", reason="Out of memory") assert "Out of memory" in str(err) @@ -122,12 +122,12 @@ def test_with_reason(self): class TestServerStartError: """Test ServerStartError""" - + def test_basic(self): err = ServerStartError("Triton") assert "Triton" in str(err) assert err.engine_name == "Triton" - + def test_with_reason(self): err = ServerStartError("Triton", reason="Port already in use") assert "Port already in use" in str(err) @@ -135,12 +135,12 @@ def test_with_reason(self): class TestInvalidDeviceError: """Test InvalidDeviceError""" - + def test_basic(self): err = InvalidDeviceError("invalid:device") assert "invalid:device" in str(err) assert "Valid formats" in str(err) - + def test_with_valid_formats(self): err = InvalidDeviceError("bad", valid_formats=["gpu:0", "cpu"]) assert "gpu:0" in str(err) @@ -148,12 +148,12 @@ def test_with_valid_formats(self): class TestDockerNotAvailableError: """Test DockerNotAvailableError""" - + def test_default(self): err = DockerNotAvailableError() assert "Docker" in str(err) assert "usermod" in str(err) # Contains fix hint - + def test_with_reason(self): err = DockerNotAvailableError(reason="Permission denied") assert "Permission denied" in str(err) @@ -161,12 +161,12 @@ def test_with_reason(self): class TestPackageNotInstalledError: """Test PackageNotInstalledError""" - + def test_basic(self): err = PackageNotInstalledError("vllm") assert "vllm" in str(err) assert err.package == "vllm" - + def test_with_install_cmd(self): err = PackageNotInstalledError("vllm", install_cmd="pip install vllm") assert "pip install vllm" in str(err) @@ -174,7 +174,7 @@ def test_with_install_cmd(self): class TestExceptionCatching: """Test catching exceptions with base classes""" - + def test_catch_with_base(self): """All exceptions can be caught with RMAbstractError""" exceptions_to_test = [ @@ -186,24 +186,23 @@ def test_catch_with_base(self): DockerNotAvailableError(), PackageNotInstalledError("test"), ] - + for exc in exceptions_to_test: with pytest.raises(RMAbstractError): raise exc - + def test_catch_backend_errors(self): """Backend errors can be caught with BackendError""" with pytest.raises(BackendError): raise BackendNotAvailableError("gpu") - + with pytest.raises(BackendError): raise DeviceNotFoundError("gpu:5") - + def test_catch_model_errors(self): """Model errors can be caught with ModelError""" with pytest.raises(ModelError): raise ModelLoadError("gpt2") - + with pytest.raises(ModelError): raise ModelNotLoadedError() - diff --git a/tests/unit/test_llm.py b/tests/unit/test_llm.py index 975da6d..315ee4c 100644 --- a/tests/unit/test_llm.py +++ b/tests/unit/test_llm.py @@ -1,8 +1,7 @@ """Tests for the new LLM proxy class.""" import os -import pytest -from unittest.mock import patch, MagicMock +from unittest.mock import patch class TestSamplingParams: @@ -52,7 +51,7 @@ class TestRequestOutput: def test_request_output_creation(self): """Test creating RequestOutput.""" - from rm_abstract import RequestOutput, CompletionOutput + from rm_abstract import CompletionOutput, RequestOutput output = RequestOutput( request_id="0", @@ -193,19 +192,23 @@ class TestLLMModuleExports: def test_llm_exported(self): """Test LLM is exported from rm_abstract.""" from rm_abstract import LLM + assert LLM is not None def test_sampling_params_exported(self): """Test SamplingParams is exported.""" from rm_abstract import SamplingParams + assert SamplingParams is not None def test_request_output_exported(self): """Test RequestOutput is exported.""" from rm_abstract import RequestOutput + assert RequestOutput is not None def test_completion_output_exported(self): """Test CompletionOutput is exported.""" from rm_abstract import CompletionOutput + assert CompletionOutput is not None diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py index e97c12c..55a3f6e 100644 --- a/tests/unit/test_registry.py +++ b/tests/unit/test_registry.py @@ -1,13 +1,16 @@ """Plugin registry tests""" -import pytest from rm_abstract.registry import ( - PluginRegistry, PluginInfo, PluginType, get_registry, ) -from rm_abstract.resources.base import Resource, ResourceType, ResourceInfo, ResourceCapability +from rm_abstract.resources.base import ( + Resource, + ResourceCapability, + ResourceInfo, + ResourceType, +) class MockResource(Resource): diff --git a/tests/unit/test_resource.py b/tests/unit/test_resource.py index 7583584..3ca1860 100644 --- a/tests/unit/test_resource.py +++ b/tests/unit/test_resource.py @@ -1,7 +1,11 @@ """Resource base class tests""" -import pytest -from rm_abstract.resources.base import Resource, ResourceType, ResourceInfo, ResourceCapability +from rm_abstract.resources.base import ( + Resource, + ResourceCapability, + ResourceInfo, + ResourceType, +) class MockResource(Resource): diff --git a/tools/check_sensitive.py b/tools/check_sensitive.py new file mode 100644 index 0000000..c6457c7 --- /dev/null +++ b/tools/check_sensitive.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Find data that must not be pushed to GitHub: internal IPs, credentials, tokens, keys. + + python3 tools/check_sensitive.py [REPO_DIR ...] + +Scans tracked and staged text files (git ls-files) of each repo and exits 1 if anything is found. +Allow a known-safe line by adding a regex to `.sensitive-allow` at the repo root (one per line), +e.g. a fake token in a unit test. Python 3.6+, stdlib only. +""" +import os +import re +import subprocess +import sys + +PATTERNS = [ + ("internal IPv4", re.compile(r"(?