gpustack 0.1.0rc1__tar.gz

gpustack-0.1.0rc1/PKG-INFO

@@ -0,0 +1,33 @@
+Metadata-Version: 2.1
+Name: gpustack
+Version: 0.1.0rc1
+Summary: GPUStack
+Author: GPUStack Authors
+Author-email: contact@gpustack.ai
+Requires-Python: >=3.11,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: aiosqlite (>=0.20.0,<0.21.0)
+Requires-Dist: alembic (>=1.13.2,<2.0.0)
+Requires-Dist: apscheduler (>=3.10.4,<4.0.0)
+Requires-Dist: argon2-cffi (>=23.1.0,<24.0.0)
+Requires-Dist: asyncpg (>=0.29.0,<0.30.0)
+Requires-Dist: attrs (>=23.2.0,<24.0.0)
+Requires-Dist: colorama (>=0.4.6,<0.5.0)
+Requires-Dist: dataclasses-json (>=0.6.7,<0.7.0)
+Requires-Dist: fastapi (>=0.111.0,<0.112.0)
+Requires-Dist: httpx[socks] (>=0.27.0,<0.28.0)
+Requires-Dist: huggingface-hub (>=0.23.3,<0.24.0)
+Requires-Dist: inflection (>=0.5.1,<0.6.0)
+Requires-Dist: netifaces (>=0.11.0,<0.12.0)
+Requires-Dist: openai (>=1.31.1,<2.0.0)
+Requires-Dist: prometheus-client (>=0.20.0,<0.21.0)
+Requires-Dist: psutil (>=5.9.8,<6.0.0)
+Requires-Dist: pydantic-settings (>=2.2.1,<3.0.0)
+Requires-Dist: pyjwt (>=2.8.0,<3.0.0)
+Requires-Dist: python-multipart (>=0.0.9,<0.0.10)
+Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: setproctitle (>=1.3.3,<2.0.0)
+Requires-Dist: sqlalchemy[asyncio] (>=2.0.30,<3.0.0)
+Requires-Dist: sqlmodel (>=0.0.18,<0.0.19)

gpustack-0.1.0rc1/gpustack/__init__.py

@@ -0,0 +1,2 @@
+__version__ = '0.1.0rc1'
+__git_commit__ = 'cdab9a9'

gpustack-0.1.0rc1/gpustack/api/exceptions.py

@@ -0,0 +1,133 @@
+from fastapi import FastAPI, Request, status
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+import httpx
+from pydantic import BaseModel
+
+
+class HTTPException(Exception):
+    def __init__(self, status_code: int, reason: str, message: str):
+        self.status_code = status_code
+        self.reason = reason
+        self.message = message
+
+
+def http_exception_factory(status_code: int, reason: str, default_message: str):
+    class_name = reason + "Exception"
+    return type(
+        class_name,
+        (HTTPException,),
+        {
+            "__init__": lambda self, message=default_message: super(
+                self.__class__, self
+            ).__init__(status_code, reason, message)
+        },
+    )
+
+
+AlreadyExistsException = http_exception_factory(
+    status.HTTP_409_CONFLICT, "AlreadyExists", "Already exists"
+)
+NotFoundException = http_exception_factory(
+    status.HTTP_404_NOT_FOUND, "NotFound", "Not found"
+)
+UnauthorizedException = http_exception_factory(
+    status.HTTP_401_UNAUTHORIZED, "Unauthorized", "Unauthorized"
+)
+ForbiddenException = http_exception_factory(
+    status.HTTP_403_FORBIDDEN, "Forbidden", "Forbidden"
+)
+InvalidException = http_exception_factory(
+    status.HTTP_422_UNPROCESSABLE_ENTITY, "Invalid", "Invalid input"
+)
+BadRequestException = http_exception_factory(
+    status.HTTP_400_BAD_REQUEST, "BadRequest", "Bad request"
+)
+InternalServerErrorException = http_exception_factory(
+    status.HTTP_500_INTERNAL_SERVER_ERROR,
+    "InternalServerError",
+    "Internal server error",
+)
+ServiceUnavailableException = http_exception_factory(
+    status.HTTP_503_SERVICE_UNAVAILABLE, "ServiceUnavailable", "Service unavailable"
+)
+
+
+def raise_if_response_error(response: httpx.Response):
+    if response.status_code < status.HTTP_400_BAD_REQUEST:
+        return
+
+    error = ErrorResponse.model_validate(response.json())
+
+    if response.status_code == status.HTTP_404_NOT_FOUND:
+        raise NotFoundException(error.message)
+
+    if (
+        response.status_code == status.HTTP_409_CONFLICT
+        and error.reason == "AlreadyExists"
+    ):
+        raise AlreadyExistsException(error.message)
+
+    if response.status_code == status.HTTP_401_UNAUTHORIZED:
+        raise UnauthorizedException(error.message)
+
+    if response.status_code == status.HTTP_403_FORBIDDEN:
+        raise ForbiddenException(error.message)
+
+    if response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY:
+        raise InvalidException(error.message)
+
+    if response.status_code == status.HTTP_400_BAD_REQUEST:
+        raise BadRequestException(error.message)
+
+    if response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR:
+        raise InternalServerErrorException(error.message)
+
+    if response.status_code == status.HTTP_503_SERVICE_UNAVAILABLE:
+        raise ServiceUnavailableException(error.message)
+
+    raise HTTPException(error.code, error.reason, error.message)
+
+
+class ErrorResponse(BaseModel):
+    code: int
+    reason: str
+    message: str
+
+
+error_responses = {
+    404: {"model": ErrorResponse},
+    409: {"model": ErrorResponse},
+    401: {"model": ErrorResponse},
+    403: {"model": ErrorResponse},
+    422: {"model": ErrorResponse},
+    400: {"model": ErrorResponse},
+    500: {"model": ErrorResponse},
+    503: {"model": ErrorResponse},
+}
+
+
+def register_handlers(app: FastAPI):
+    @app.exception_handler(HTTPException)
+    async def http_exception_handler(request: Request, exc: HTTPException):
+        return JSONResponse(
+            status_code=exc.status_code,
+            content=ErrorResponse(
+                code=exc.status_code, reason=exc.reason, message=exc.message
+            ).model_dump(),
+        )
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_exception_handler(request, exc: RequestValidationError):
+        message = f"{len(exc.errors())} validation errors:\n"
+        for err in exc.errors():
+            message += f"  {err}\n"
+
+        return JSONResponse(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            content=ErrorResponse(
+                code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                reason="Invalid",
+                message=message,
+            ).model_dump(),
+        )

gpustack-0.1.0rc1/gpustack/api/middlewares.py

@@ -0,0 +1,122 @@
+from datetime import date
+import json
+import logging
+import time
+from fastapi import Request, Response
+from fastapi.responses import StreamingResponse
+from jwt import DecodeError, ExpiredSignatureError
+from starlette.middleware.base import BaseHTTPMiddleware
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from gpustack.schemas.model_usage import ModelUsage
+from gpustack.schemas.models import Model
+from gpustack.schemas.users import User
+from gpustack.security import JWT_TOKEN_EXPIRE_MINUTES, JWTManager
+from gpustack.server.auth import SESSION_COOKIE_NAME
+from gpustack.server.db import get_engine
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+logger = logging.getLogger(__name__)
+
+
+class ModelUsageMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        response = await call_next(request)
+        if (
+            not request.url.path == "/v1-openai/chat/completions"
+            or response.status_code != 200
+        ):
+            return response
+
+        stream: bool = getattr(request.state, "stream", False)
+        if stream:
+            response = await self.handle_streaming_response(request, response)
+        else:
+            response_body = b"".join([chunk async for chunk in response.body_iterator])
+            try:
+                completion_dict = json.loads(response_body)
+                chat_completion = ChatCompletion(**completion_dict)
+                await self.process_model_usage(request, chat_completion)
+            except Exception as e:
+                logger.error(f"Error processing model usage: {e}")
+            response = Response(content=response_body, headers=dict(response.headers))
+
+        return response
+
+    async def handle_streaming_response(
+        self, request: Request, response: StreamingResponse
+    ):
+        async def streaming_generator():
+            try:
+                async for chunk in response.body_iterator:
+                    data = chunk.decode("utf-8")
+                    yield chunk
+                    if '"completion_tokens":' in data:
+                        completion_dict = json.loads(data.split('data: ')[-1])
+                        completion_chunk = ChatCompletionChunk(**completion_dict)
+                        await self.process_model_usage(request, completion_chunk)
+                        break
+
+                async for chunk in response.body_iterator:
+                    yield chunk
+            except Exception as e:
+                logger.error(f"Error processing streaming response: {e}")
+
+        return StreamingResponse(streaming_generator(), headers=dict(response.headers))
+
+    async def process_model_usage(
+        self, request: Request, chat_completion: ChatCompletion | ChatCompletionChunk
+    ):
+        completion_tokens = chat_completion.usage.completion_tokens
+        prompt_tokens = chat_completion.usage.prompt_tokens
+        user: User = request.state.user
+        model: Model = request.state.model
+        fields = {
+            "user_id": user.id,
+            "model_id": model.id,
+            "date": date.today(),
+        }
+        model_usage = ModelUsage(
+            **fields,
+            completion_token_count=completion_tokens,
+            prompt_token_count=prompt_tokens,
+            request_count=1,
+            operation="chat_completion",
+        )
+        async with AsyncSession(get_engine()) as session:
+            current_model_usage = await ModelUsage.one_by_fields(session, fields)
+            if current_model_usage:
+                current_model_usage.completion_token_count += completion_tokens
+                current_model_usage.prompt_token_count += prompt_tokens
+                current_model_usage.request_count += 1
+                await current_model_usage.update(session)
+            else:
+                await ModelUsage.create(session, model_usage)
+
+
+class RefreshTokenMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        response = await call_next(request)
+
+        jwt_manager: JWTManager = request.app.state.jwt_manager
+        token = request.cookies.get(SESSION_COOKIE_NAME)
+
+        if token:
+            try:
+                payload = jwt_manager.decode_jwt_token(token)
+                if payload:
+                    # Check if the token is about to expire (less than 5 minutes left)
+                    if payload['exp'] - time.time() < 5 * 60:
+                        new_token = jwt_manager.create_jwt_token(
+                            username=payload['sub']
+                        )
+                        response.set_cookie(
+                            key=SESSION_COOKIE_NAME,
+                            value=new_token,
+                            httponly=True,
+                            max_age=JWT_TOKEN_EXPIRE_MINUTES * 60,
+                            expires=JWT_TOKEN_EXPIRE_MINUTES * 60,
+                        )
+            except (ExpiredSignatureError, DecodeError):
+                pass
+
+        return response

gpustack-0.1.0rc1/gpustack/chat/manager.py

@@ -0,0 +1,184 @@
+import os
+import sys
+from typing import List, Optional
+
+from colorama import Fore, Style
+from openai import OpenAI
+from pydantic import model_validator
+from pydantic_settings import BaseSettings
+from tqdm import tqdm
+
+from gpustack.client.generated_clientset import ClientSet
+from gpustack.schemas.models import (
+    ModelCreate,
+    ModelInstance,
+    ModelInstanceStateEnum,
+    SourceEnum,
+)
+from gpustack.server.bus import Event
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionUserMessageParam,
+    ChatCompletionAssistantMessageParam,
+)
+
+
+class ChatConfig(BaseSettings):
+    debug: bool = False
+    model: str
+    prompt: Optional[str] = None
+    base_url: str = os.getenv("GPUSTACK_SERVER_URL", "http://127.0.0.1")
+    api_key: Optional[str] = os.getenv("GPUSTACK_API_KEY")
+
+    @model_validator(mode="after")
+    def check_api_key(self):
+        if self.base_url != "http://127.0.0.1" and not self.api_key:
+            raise ValueError(
+                "API key is required. Please set GPUSTACK_API_KEY env var."
+            )
+        elif self.base_url == "http://127.0.0.1" and not self.api_key:
+            self.api_key = "local"
+
+
+def parse_arguments(args) -> ChatConfig:
+    return ChatConfig(debug=args.debug, model=args.model, prompt=args.prompt)
+
+
+def print_completion_result(message):
+    # move cursor to the end of previous line
+    sys.stdout.write("\033[F\033[1000C")
+    print(message)
+
+
+def print_error(message):
+    print(f"{Fore.RED}{message}{Style.RESET_ALL}")
+
+
+class ChatManager:
+    def __init__(self, cfg: ChatConfig) -> None:
+        self._model_name = cfg.model
+        self._prompt = cfg.prompt
+        self._clientset = ClientSet(base_url=cfg.base_url, api_key=cfg.api_key)
+        self._openai_client = OpenAI(
+            base_url=f"{cfg.base_url}/v1-openai", api_key=cfg.api_key
+        )
+        self._history: List[ChatCompletionMessageParam] = []
+
+    def start(self):
+        self._ensure_model()
+
+        if self._prompt:
+            self.chat_completion(self._prompt)
+            return
+
+        user_input = None
+        while True:
+            user_input = input(">")
+            if user_input == "\\q" or user_input == "\\quit":
+                break
+            elif user_input == "\\?" or user_input == "\\h" or user_input == "\\help":
+                self._print_help()
+                continue
+            elif user_input == "\\c" or user_input == "\\clear":
+                self._clear_context()
+                continue
+            elif not user_input.strip():
+                continue
+
+            try:
+                self.chat_completion(user_input)
+            except Exception as e:
+                print_error(e)
+
+    @staticmethod
+    def _print_help():
+        print("Commands:")
+        print("  \\q or \\quit - Quit the chat")
+        print("  \\c or \\clear - Clear chat context in prompt")
+        print("  \\h or \\? or \\help - Print this help message")
+
+    def _clear_context(self):
+        self._history = []
+        print("Chat context cleared.")
+
+    def _ensure_model(self):
+        models = self._clientset.models.list()
+        for model in models.items:
+            if model.name == self._model_name:
+                self._model = model
+                break
+
+        if not hasattr(self, "_model"):
+            self._create_model()
+
+        self._wait_for_model_ready()
+
+    def _create_model(self):
+        model_create = ModelCreate(
+            name=self._model_name,
+            source=SourceEnum.OLLAMA_LIBRARY,
+            ollama_library_model_name=self._model_name,
+        )
+        created = self._clientset.models.create(model_create=model_create)
+        self._model = created
+
+    def _wait_for_model_ready(self):
+        def stop_when_running(event: Event) -> bool:
+            if event.data["id"] == self._model.id and event.data["state"] == "Running":
+                return True
+            elif event.data["state"] == ModelInstanceStateEnum.error:
+                raise Exception(f"Error running model: {event.data['state_message']}")
+            return False
+
+        with tqdm(
+            total=0,
+            desc=f"Preparing {self._model_name} model...",
+            bar_format="{desc}",
+            leave=False,
+        ) as pbar:
+            current_progress = 0
+
+            def print_progress(event: Event):
+                nonlocal current_progress
+                mi = ModelInstance.model_validate(event.data)
+                if mi.download_progress is not None:
+                    increment = mi.download_progress - current_progress
+                    if increment <= 0:
+                        return
+
+                    if pbar.total == 0:
+                        pbar.total = 100
+                        pbar.bar_format = "{l_bar}{bar}{r_bar}"
+                        pbar.set_description(f"Downloading {self._model_name} model")
+                        pbar.reset()
+
+                    pbar.update(increment)
+                    current_progress = mi.download_progress
+
+            self._clientset.model_instances.watch(
+                stop_condition=stop_when_running,
+                callback=print_progress,
+                params={"model_id": self._model.id},
+            )
+
+    def chat_completion(self, prompt: str):
+        self._history.append(
+            ChatCompletionUserMessageParam(role="user", content=prompt)
+        )
+
+        completion = self._openai_client.chat.completions.create(
+            model=self._model_name,
+            messages=self._history,
+            stream=True,
+        )
+
+        result = ""
+        for chunk in completion:
+            if chunk.choices[0].delta.content:
+                result += chunk.choices[0].delta.content
+                print(chunk.choices[0].delta.content, end="", flush=True)
+
+        self._history.append(
+            ChatCompletionAssistantMessageParam(role="assistant", content=result)
+        )
+        print()

gpustack-0.1.0rc1/gpustack/client/__init__.py

@@ -0,0 +1,4 @@
+from .generated_clientset import ClientSet
+
+
+__all__ = ["ClientSet"]

gpustack-0.1.0rc1/gpustack/client/generated_clientset.py

@@ -0,0 +1,39 @@
+import base64
+from .generated_http_client import HTTPClient
+from typing import Optional
+
+from .generated_worker_client import WorkerClient
+from .generated_model_client import ModelClient
+from .generated_model_instance_client import ModelInstanceClient
+from .generated_user_client import UserClient
+
+
+class ClientSet:
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        headers: Optional[dict] = None,
+    ):
+        if headers is None:
+            headers = {}
+
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        elif username and password:
+            base64_credentials = base64.b64encode(
+                f"{username}:{password}".encode()
+            ).decode()
+            headers["Authorization"] = f"Basic {base64_credentials}"
+
+        self.base_url = base_url
+        self.headers = headers
+
+        http_client = HTTPClient(base_url=base_url).with_headers(headers)
+
+        self.workers = WorkerClient(http_client)
+        self.models = ModelClient(http_client)
+        self.model_instances = ModelInstanceClient(http_client)
+        self.users = UserClient(http_client)