-
Notifications
You must be signed in to change notification settings - Fork 35
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
clean up shortfin llm integration tests
- Loading branch information
Showing
6 changed files
with
522 additions
and
348 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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",), | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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") |
Oops, something went wrong.