Skip to content

Commit

Permalink
clean up shortfin llm integration tests
Browse files Browse the repository at this point in the history
  • Loading branch information
renxida committed Dec 26, 2024
1 parent f5e9cb4 commit 2a33f86
Show file tree
Hide file tree
Showing 6 changed files with 522 additions and 348 deletions.
17 changes: 17 additions & 0 deletions app_tests/integration_tests/llm/device_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from typing import Tuple
from dataclasses import dataclass


@dataclass
class DeviceSettings:
compile_flags: Tuple[str]
server_flags: Tuple[str]


CPU = DeviceSettings(
compile_flags=(
"-iree-hal-target-backends=llvm-cpu",
"--iree-llvmcpu-target-cpu=host",
),
server_flags=("--device=local-task",),
)
227 changes: 227 additions & 0 deletions app_tests/integration_tests/llm/model_management.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
"""Module for managing model artifacts through various processing stages."""
import logging
from pathlib import Path
import subprocess
from dataclasses import dataclass
from typing import Optional, Tuple
from enum import Enum, auto

logger = logging.getLogger(__name__)


class ModelSource(Enum):
HUGGINGFACE = auto()
LOCAL = auto()
AZURE = auto()


@dataclass
class AzureConfig:
"""Configuration for Azure blob storage downloads."""

account_name: str
container_name: str
blob_path: str
auth_mode: str = "key"


@dataclass
class ModelConfig:
"""Configuration for model source and settings."""

model_file: str
tokenizer_id: str
batch_sizes: Tuple[int, ...]
device_settings: "DeviceSettings"
source: ModelSource
repo_id: Optional[str] = None
local_path: Optional[Path] = None
azure_config: Optional[AzureConfig] = None

def __post_init__(self):
if self.source == ModelSource.HUGGINGFACE and not self.repo_id:
raise ValueError("repo_id required for HuggingFace models")
elif self.source == ModelSource.LOCAL and not self.local_path:
raise ValueError("local_path required for local models")
elif self.source == ModelSource.AZURE and not self.azure_config:
raise ValueError("azure_config required for Azure models")


@dataclass
class ModelArtifacts:
"""Container for all paths related to model artifacts."""

weights_path: Path
tokenizer_path: Path
mlir_path: Path
vmfb_path: Path
config_path: Path


class ModelStageManager:
"""Manages different stages of model processing with caching behavior."""

def __init__(self, base_dir: Path, config: ModelConfig):
self.base_dir = base_dir
self.config = config
self.model_dir = self._get_model_dir()
self.model_dir.mkdir(parents=True, exist_ok=True)

def _get_model_dir(self) -> Path:
"""Creates and returns appropriate model directory based on source."""
if self.config.source == ModelSource.HUGGINGFACE:
return self.base_dir / self.config.repo_id.replace("/", "_")
elif self.config.source == ModelSource.LOCAL:
return self.base_dir / "local" / self.config.local_path.stem
elif self.config.source == ModelSource.AZURE:
return (
self.base_dir
/ "azure"
/ self.config.azure_config.blob_path.replace("/", "_")
)
raise ValueError(f"Unsupported model source: {self.config.source}")

def _download_from_huggingface(self) -> Path:
"""Downloads model from HuggingFace."""
model_path = self.model_dir / self.config.model_file
if not model_path.exists():
logger.info(f"Downloading model {self.config.repo_id} from HuggingFace")
subprocess.run(
f"huggingface-cli download --local-dir {self.model_dir} {self.config.repo_id} {self.config.model_file}",
shell=True,
check=True,
)
return model_path

def _copy_from_local(self) -> Path:
"""Copies model from local filesystem."""
import shutil

model_path = self.model_dir / self.config.model_file
if not model_path.exists():
logger.info(f"Copying local model from {self.config.local_path}")
shutil.copy2(self.config.local_path, model_path)
return model_path

def _download_from_azure(self) -> Path:
"""Downloads model from Azure blob storage."""
model_path = self.model_dir / self.config.model_file
if not model_path.exists():
logger.info(
f"Downloading model from Azure blob storage: {self.config.azure_config.blob_path}"
)
subprocess.run(
[
"az",
"storage",
"blob",
"download",
"--account-name",
self.config.azure_config.account_name,
"--container-name",
self.config.azure_config.container_name,
"--name",
self.config.azure_config.blob_path,
"--file",
str(model_path),
"--auth-mode",
self.config.azure_config.auth_mode,
],
check=True,
)
return model_path

def prepare_tokenizer(self) -> Path:
"""Downloads and prepares tokenizer."""
tokenizer_path = self.model_dir / "tokenizer.json"
if not tokenizer_path.exists():
logger.info(f"Downloading tokenizer {self.config.tokenizer_id}")
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(self.config.tokenizer_id)
tokenizer.save_pretrained(self.model_dir)
return tokenizer_path

def export_model(self, weights_path: Path) -> Tuple[Path, Path]:
"""Exports model to MLIR format."""
bs_string = ",".join(map(str, self.config.batch_sizes))
mlir_path = self.model_dir / "model.mlir"
config_path = self.model_dir / "config.json"

logger.info(
"Exporting model with following settings:\n"
f" MLIR Path: {mlir_path}\n"
f" Config Path: {config_path}\n"
f" Batch Sizes: {bs_string}"
)

subprocess.run(
[
"python",
"-m",
"sharktank.examples.export_paged_llm_v1",
"--block-seq-stride=16",
f"--{weights_path.suffix.strip('.')}-file={weights_path}",
f"--output-mlir={mlir_path}",
f"--output-config={config_path}",
f"--bs={bs_string}",
],
check=True,
)

logger.info(f"Model successfully exported to {mlir_path}")
return mlir_path, config_path

def compile_model(self, mlir_path: Path) -> Path:
"""Compiles model to VMFB format."""
vmfb_path = self.model_dir / "model.vmfb"
logger.info(f"Compiling model to {vmfb_path}")

compile_command = [
"iree-compile",
str(mlir_path),
"-o",
str(vmfb_path),
]
compile_command.extend(self.config.device_settings.compile_flags)

subprocess.run(compile_command, check=True)
logger.info(f"Model successfully compiled to {vmfb_path}")
return vmfb_path


class ModelProcessor:
"""Main interface for processing models through all stages."""

def __init__(self, base_dir: Path):
self.base_dir = Path(base_dir)

def process_model(self, config: ModelConfig) -> ModelArtifacts:
"""Process model through all stages and return paths to all artifacts."""
manager = ModelStageManager(self.base_dir, config)

# Stage 1: Download weights and tokenizer (cached)
if config.source == ModelSource.HUGGINGFACE:
weights_path = manager._download_from_huggingface()
elif config.source == ModelSource.LOCAL:
weights_path = manager._copy_from_local()
elif config.source == ModelSource.AZURE:
weights_path = manager._download_from_azure()
else:
raise ValueError(f"Unsupported model source: {config.source}")

tokenizer_path = manager.prepare_tokenizer()

# Stage 2: Export model (fresh every time)
mlir_path, config_path = manager.export_model(weights_path)

# Stage 3: Compile model (fresh every time)
vmfb_path = manager.compile_model(mlir_path)

return ModelArtifacts(
weights_path=weights_path,
tokenizer_path=tokenizer_path,
mlir_path=mlir_path,
vmfb_path=vmfb_path,
config_path=config_path,
)
97 changes: 97 additions & 0 deletions app_tests/integration_tests/llm/server_management.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
"""Handles server lifecycle and configuration."""
import json
import socket
from contextlib import closing
from dataclasses import dataclass, field
import subprocess
import time
import requests
from pathlib import Path
import sys
from typing import Optional

from .device_settings import DeviceSettings
from .model_management import ModelArtifacts


@dataclass
class ServerConfig:
"""Configuration for server instance."""

port: int
artifacts: ModelArtifacts
device_settings: DeviceSettings

# things we need to write to config
prefix_sharing_algorithm: str = "none"


class ServerManager:
"""Manages server lifecycle and configuration."""

@staticmethod
def find_available_port() -> int:
"""Finds an available port for the server."""
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1]

def __init__(self, config: ServerConfig):
self.config = config

def write_config(self) -> Path:
"""Creates server config by extending the exported model config."""
source_config_path = self.config.artifacts.config_path
server_config_path = (
source_config_path.parent
/ f"server_config_{self.config.prefix_sharing_algorithm}.json"
)

# Read the exported config as base
with open(source_config_path) as f:
config = json.load(f)

# Update with server-specific settings
config.update(
{
"paged_kv_cache": {
"prefix_sharing_algorithm": self.config.prefix_sharing_algorithm
}
}
)

# Write the extended config
with open(server_config_path, "w") as f:
json.dump(config, f)
6
return server_config_path

def start(self) -> subprocess.Popen:
"""Starts the server process."""
config_path = self.write_config()
cmd = [
sys.executable,
"-m",
"shortfin_apps.llm.server",
f"--tokenizer_json={self.config.artifacts.tokenizer_path}",
f"--model_config={config_path}",
f"--vmfb={self.config.artifacts.vmfb_path}",
f"--parameters={self.config.artifacts.weights_path}",
f"--port={self.config.port}",
]
cmd.extend(self.config.device_settings.server_flags)
process = subprocess.Popen(cmd)
self._wait_for_server(timeout=10)
return process

def _wait_for_server(self, timeout: int = 10):
"""Waits for server to be ready."""
start = time.time()
while time.time() - start < timeout:
try:
requests.get(f"http://localhost:{self.config.port}/health")
return
except requests.exceptions.ConnectionError:
time.sleep(1)
raise TimeoutError(f"Server failed to start within {timeout} seconds")
Loading

0 comments on commit 2a33f86

Please sign in to comment.