Metadata-Version: 2.4
Name: autodraft-sd
Version: 0.1.4
Summary: Speculative decoding engine with local and remote target model execution
Author-email: MobiTree Authors <rlaehdgus021818@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/PJChoi1/MobiTree
Project-URL: Repository, https://github.com/PJChoi1/MobiTree
Project-URL: Issues, https://github.com/PJChoi1/MobiTree/issues
Keywords: speculative-decoding,llm,inference,mobitree
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: torch
Requires-Dist: transformers
Requires-Dist: accelerate
Requires-Dist: fschat[llm_judge]
Requires-Dist: shortuuid
Requires-Dist: tqdm
Requires-Dist: datasets
Requires-Dist: numpy
Requires-Dist: psutil
Requires-Dist: nvidia-ml-py
Requires-Dist: safetensors
Provides-Extra: quant
Requires-Dist: bitsandbytes; extra == "quant"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"

# MobiTree 실행 가이드 (환경 구축 + UI 사용)

이 문서는 **처음 설치하는 사람도 복사/붙여넣기만으로 실행**할 수 있게 작성했습니다.

- 대상: Ubuntu/Linux + NVIDIA GPU 환경
- 범위: 드라이버/CUDA 확인, `venv` 구성, 의존성 설치, Target 실행, UI 실행, UI 사용법

---

## 1) 사전 준비

### 1-1. 필수 소프트웨어

- NVIDIA Driver 설치
- Python 3.10 이상
- `git`
- (선택) `tmux` 또는 `screen` (서버 장시간 실행 시 편리)

### 1-2. 드라이버/CUDA 확인 (복붙)

```bash
nvidia-smi
python3 --version
```

정상이라면:

- `nvidia-smi`가 GPU/Driver 정보를 출력
- Python 버전이 출력

참고:

- 이 저장소는 `torch==2.7.0+cu128` 기준 `requirements.txt`를 사용합니다.
- 시스템에 CUDA Toolkit(`nvcc`)이 없어도, PyTorch wheel + NVIDIA driver 조합으로 동작할 수 있습니다.
- `requirements.txt`에 PyTorch cu128 인덱스가 포함되어 있으므로 `pip install -r requirements.txt`만 실행하면 됩니다.

---

## 2) 프로젝트 받기

```bash
git clone git@github.com:PJChoi1/MobiTree.git MobiTree
cd MobiTree
git checkout seperate
```

이미 코드가 있다면 이 단계는 생략하세요.

---

## 3) 가상환경(venv) + 의존성 설치

아래 블록을 **한 번에 그대로 실행**하세요.

```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
pip install -r requirements.txt
```

설치 확인:

```bash
python -c "import torch; print('torch:', torch.__version__); print('cuda:', torch.cuda.is_available())"
```

`cuda: True`가 나오면 GPU 인식이 정상입니다.

---

## 4) 실행 방법 (가장 많이 쓰는 방법)

아래 순서대로 진행하세요. (`터미널 A`: UI)

### 4-1. 1단계: UI 실행

터미널 A:

```bash
source .venv/bin/activate
python3 -m uvicorn chat_ui.main:app --host 0.0.0.0 --port 8000
```

브라우저 접속:

- [http://localhost:8000](http://localhost:8000)

### 4-2. 2단계: UI 접속 확인

브라우저에서 UI가 열리면 상단 `Server` 영역과 `Add Server` 버튼이 보이는지 확인하세요.

### 4-3. 3단계: UI에서 Server 추가

브라우저 UI에서:

1. `Add Server` 클릭
2. `Server Name` 입력 (예: `icnl-server`)
3. `IP Address` 입력 (예: `163.152.163.152`)
4. `Port` 입력 (예: `26001`)
5. `Add Server` 클릭

---

## 5) UI 사용법 (처음 사용자용)

### 5-1. 서버/모델/Quantization 선택

이미 등록된 서버를 기준으로 아래 순서로 선택하세요.

1. `Server` 드롭다운에서 사용할 서버 선택
2. `Server Model`에서 target 모델 선택
3. `Draft Model`에서 draft 모델 선택
4. `Server Q`/`Draft Q`에서 quantization(4bit/8bit) 선택
5. `Start`를 눌러 런타임 시작 후 질문 전송

### 5-2. 주요 버튼/설정값 의미

- `Start`: 채팅 런타임 시작
- `Stop`: 채팅 런타임 중지 + target 모델 unload 요청
- `Profile LLM`: reference cache/profile 갱신
- `Token Source Coloring`: 토큰 출처별 색상 표시
- `Algorithm`: 디코딩 전략 선택 (`MobiTree`, `Server-Only`, `Server-Only-AR`, `OPT-Tree`, `Fixed-tree`)
- `Mode`: 실행 모드 선택 (`Chat`은 대화, `Benchmark`는 평가 중심)
- `Max New Tokens`: 한 번 응답에서 생성할 최대 토큰 수
- `Dataset`: Benchmark 모드에서 사용할 평가 데이터셋 선택

### 5-3. UI에서 add한 server name으로 Target 실행

핵심: 4-3에서 추가한 `Server Name`과 Target 실행 시 `--server-name` 값을 **같게** 맞추세요.

예를 들어 4-3에서 `Server Name=icnl-server`로 추가했다면, 터미널 B에서:

```bash
source .venv/bin/activate
./run_target.sh --host 0.0.0.0 --port 26001 --device-map auto --load-in-8bit --lazy-load --enable-auto-target-profile --server-name icnl-server
```

설명:

- `--host`/`--port`: 4-3에서 입력한 `IP Address`/`Port`와 동일해야 연결됩니다.
- `--server-name`: 4-3에서 입력한 `Server Name`과 동일하게 사용하세요.
- `--lazy-load`: 최초 접속 시 모델 로드
- `--enable-auto-target-profile`: target profile 없으면 자동 생성

---

## 6) 캐시/프로파일 파일 위치

- Target profile: `data/profile/profile_target_<server-name>_<model>_tq-<quant>.json`
- Reference cache: `data/reference/ref_<server-name>_<base>_<device>_<draft>_tq-<...>_dq-<...>_mt_bench_<metric>_<mode>_*.json`

파일명이 맞아야 재사용됩니다.

---

## 7) `autodraft` Python API

`autodraft` 는 위 MobiTree 런타임을 감싸는 얇은 Python wrapper 입니다.
기존 shell script (`run_target.sh`, `run_mt_bench_sd.sh` 등) 와 CLI 동작은 그대로
유지되고, 추가로 다음 두 가지 사용법이 지원됩니다.

### 7-1. Repo clone 후 바로 import

```bash
git clone https://github.com/PJChoi1/MobiTree.git
cd MobiTree
# 터미널 A: target server
python examples/target.py
# 터미널 B: draft side (실제 generated text)
python examples/draft.py
```

또는 Python 코드에서:

```python
# 터미널 A — target server
from autodraft import serve_target

serve_target(
    host="0.0.0.0",
    port=26001,
    target_model="meta-llama/Llama-3.2-1B-Instruct",
    quantization="8bit",
    server_name="autodraft",
    hf_token=None,  # 또는 "hf_xxx"
)  # 이 호출은 서버 loop라 영원히 block 됩니다.

# 터미널 B — draft side
from autodraft import Autodraft

engine = Autodraft(
    draft_model="meta-llama/Llama-3.2-1B-Instruct",
    target_model="meta-llama/Llama-3.2-1B-Instruct",
    target_host="127.0.0.1",
    target_port=26001,
)
result = engine.run("input text", server_name="autodraft")
print(result["generated_text"])
```

MobiTree 는 항상 split-process 입니다 — draft 와 target 은 별도 프로세스로
실행되고 socket 으로 통신합니다. 같은 머신에서는 `target_host="127.0.0.1"`
(기본값), 다른 머신에서는 target 머신의 IP 를 적으면 됩니다.

`server_name` 은 양쪽이 일치해야 profile/reference cache 가 재사용됩니다 (default: ``"autodraft"``).

주의:
- `pip install` 없이 import 하려면 Python 의 `sys.path` 에 repo root 가 포함되어야 합니다.
  `examples/draft.py` 와 `examples/target.py` 는 자기 자신의 부모 디렉토리를
  `sys.path` 에 끼워넣는 4줄짜리 bootstrap 이 들어 있어 `python examples/draft.py`
  로 그대로 실행됩니다.
- 본인 코드에서 `pip install` 없이 쓰려면 다음 중 하나를 사용하세요:
  ```bash
  # (a) repo root 에서 -m 으로 실행
  python -m examples.draft

  # (b) PYTHONPATH 에 repo path 를 추가
  export PYTHONPATH=/path/to/MobiTree:$PYTHONPATH
  python my_script.py

  # (c) 본인 스크립트 첫 줄에 sys.path bootstrap 을 추가 (examples 와 동일한 패턴)
  ```
- target server 가 떠 있어야 draft 가 동작합니다. 먼저 `python examples/target.py`
  (또는 5-3 의 `./run_target.sh`) 로 target 을 실행해 두세요.
- gated 모델 (`meta-llama/*` 등) 을 사용하려면 HF access token 이 필요합니다. 두 가지 방법이 있습니다:
  ```python
  # (a) 생성자에 직접 전달
  engine = Autodraft(draft_model=..., target_model=..., hf_token="hf_xxx")

  # (b) 환경변수로 설정 (HF_TOKEN 또는 HUGGING_FACE_HUB_TOKEN)
  # bash:
  #   export HF_TOKEN=hf_xxx
  ```
  `hf_token` 인자는 `repr` 출력에서 `'***'` 로 마스킹됩니다. 이 token 은 ``run_local`` 안에서
  ``HF_TOKEN`` / ``HUGGING_FACE_HUB_TOKEN`` env var 로 export 되어 ``run_target.sh`` 로 띄운
  target server 도 같은 token 을 사용할 수 있게 됩니다 (target 쪽은 server 프로세스의 env 에
  자체적으로 token 을 export 하면 됩니다).

### 7-2. Editable install

```bash
pip install -e .
# 4/8-bit quantization 도 쓰려면:
pip install -e ".[quant]"
# 테스트/lint/build 까지:
pip install -e ".[dev]"
```

이렇게 하면 repo 밖 어디서든 `import autodraft` 가 가능합니다.

### 7-3. PyPI 설치 (배포 후)

```bash
pip install autodraft-sd
```

> Distribution 이름은 `autodraft-sd` 입니다. PyPI 에 `autodraft` 이름이 이미
> 사용 중이라 충돌을 피하기 위해 `-sd` (speculative decoding) 접미사를
> 붙였습니다. **import 이름은 그대로 `autodraft`** 이므로 코드에서는
> `from autodraft import Autodraft` 를 쓰면 됩니다 (Pillow 가 `pip install
> pillow` 인데 `from PIL` 로 import 하는 것과 같은 패턴).

PyPI wheel 에는 wrapper API (`autodraft/`) 와 함께 speculative decoding
runtime (`evaluation/`, `opt_classic/`) 도 포함됩니다. 따라서 source checkout
없이 `pip install autodraft-sd` 한 번으로 `engine.run(...)` 과
`serve_target(...)` 모두 동작합니다 (단, GPU 와 적절한 PyTorch wheel 이 필요
— 7-5 참조).

profile/reference cache 디렉토리는 기본적으로 **스크립트를 실행한 디렉토리의
`./data/`** 에 생성됩니다 (즉 `python my_script.py` 했다면 그 명령이 실행된
디렉토리의 `data/profile/`, `data/reference/`). source-checkout 사용자가 repo
root 에서 실행하던 동작과 자연스럽게 일치합니다.

처음 실행하면 (cache 가 없으면) target latency profile 과 draft latency
profile 을 자동으로 굽습니다 (수십 초~수 분). 두 번째 실행부터는 cache hit
하여 빠르게 시작됩니다. cache 파일명은 GPU 이름을 포함하므로 다른 GPU 로 옮기면
새로 굽습니다. GPU 이름을 직접 지정하려면:

```bash
export MOBITREE_DEVICE_NAME="rtx-5080"   # 자유 라벨
```

지정하지 않으면 `torch.cuda.get_device_name(0)` 으로 자동 감지하고, GPU 가
없으면 `"unknown-gpu"` 라벨을 사용합니다.

다른 cache 위치 (예: 여러 프로젝트가 cache 를 공유하게 하고 싶을 때) 를
쓰려면:

```bash
export MOBITREE_DATA_DIR=/path/to/shared/cache
```

`chat_ui/`, `data/` (벤치마크 데이터셋), `result/` 등 연구용 자산은 wheel 에
포함되지 않습니다. 그것들은 GitHub source checkout 으로만 접근하세요.

4/8-bit quantization (bitsandbytes 의존) 은 별도 extra 입니다:

```bash
pip install "autodraft-sd[quant]"
```

### 7-4. 다른 머신에 떠 있는 target 사용

target 측은 `python examples/target.py` 로 띄우거나 (server 머신에서) 기존
shell 방식 그대로:

```bash
./run_target.sh --host 0.0.0.0 --port 26001 \
                --base-model-path meta-llama/Llama-2-70b-chat-hf \
                --draft-model-path meta-llama/Llama-2-7b-chat-hf \
                --load-in-8bit --lazy-load --enable-auto-target-profile \
                --server-name my-server
```

클라이언트 측 (다른 머신):

```python
from autodraft import Autodraft

engine = Autodraft(
    draft_model="...",
    target_model="...",
    target_host="TARGET_SERVER_IP",
    target_port=26001,
)

result = engine.run("input text", proactive=True, cs=None, server_name="my-server")
```

### 7-5. PyTorch 설치에 대한 안내

`autodraft` 의 base wheel 은 PyTorch 를 의존성으로 가지지 않습니다 (`autodraft/`
자체는 stdlib 만 사용). in-process 실행에 필요한 torch / transformers /
accelerate / fastchat 등은 `[local]` extra 또는 기존 `requirements.txt` 로
설치하세요.

특히 PyTorch 는 CUDA 버전마다 wheel 이 다르므로, 본 repo 의 `requirements.txt`
에 명시된 `torch==2.7.0+cu128` 과 같이 환경에 맞는 wheel 을 직접 설치하는 것을
권장합니다.

---

## 8) PyPI 배포 (Trusted Publishing)

PyPI 배포는 **모든 push 에 대해 자동으로 수행되지 않습니다**. 다음 두 가지가
모두 만족할 때만 배포됩니다.

1. `pyproject.toml` 의 `version` 이 올라가 있다.
2. 그 version 에 매칭되는 git tag (`vX.Y.Z`) 가 push 된다.

### 8-1. Trusted Publishing 일회성 설정

1. PyPI 계정 생성 (없다면)
2. <https://pypi.org/manage/account/publishing/> 에서 **Add a new pending publisher** 선택
3. 다음 값을 입력:
   - **PyPI Project Name**: `autodraft-sd` (또는 `pyproject.toml` 의 `name` 값)
   - **Owner**: GitHub repository owner (예: `PJChoi1`)
   - **Repository name**: `MobiTree`
   - **Workflow name**: `publish.yml`
   - **Environment name**: `pypi`
4. GitHub repository 의 **Settings → Environments** 에서 `pypi` environment 를 생성

이렇게 하면 첫 배포부터 API token 없이 OIDC 기반으로 publish 가 됩니다.
`PYPI_API_TOKEN` secret 은 사용하지 않습니다.

### 8-2. 새 버전 배포

```bash
# pyproject.toml 의 version 을 0.1.0 → 0.1.1 등으로 올린 뒤
git add pyproject.toml
git commit -m "release: 0.1.1"
git tag v0.1.1
git push origin v0.1.1
```

`publish.yml` workflow 가 build → twine check → PyPI publish 를 자동으로 수행합니다.
workflow 는 시작 시 git tag 와 `pyproject.toml` 의 version 일치 여부를 검증하므로,
둘이 어긋나면 배포가 거부됩니다.

> PyPI 는 같은 version 을 다시 업로드할 수 없습니다. 코드를 수정할 때마다 자동
> 배포하지 마시고, version 을 올리고 tag 를 push 할 때만 배포하세요.
