diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 03921269..a4b5c68f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -127,9 +127,8 @@ jobs: - name: Start metadata server run: | - cd mooncake-transfer-engine/example/http-metadata-server-python - pip install aiohttp - python ./bootstrap_server.py & + source test_env/bin/activate + mooncake_http_metadata_server --port 8080 & shell: bash - name: Run tests diff --git a/mooncake-wheel/mooncake/http_metadata_server.py b/mooncake-wheel/mooncake/http_metadata_server.py new file mode 100644 index 00000000..2e30cce7 --- /dev/null +++ b/mooncake-wheel/mooncake/http_metadata_server.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +""" +HTTP Metadata Server for Mooncake. + +This module provides a simple HTTP server for storing and retrieving metadata +used by Mooncake. It can be used as an alternative to etcd for metadata storage. +""" + +import argparse +import asyncio +import logging +import os +import signal +import sys +import threading +from enum import Enum +from time import sleep + +from aiohttp import web + + +class KVPoll(Enum): + """Status of the KV server.""" + Failed = 0 + Bootstrapping = 1 + WaitingForInput = 2 + Transferring = 3 + Success = 4 + + +class KVBootstrapServer: + """HTTP server for storing and retrieving metadata.""" + + def __init__(self, port: int, host: str = "0.0.0.0"): + """Initialize the server. + + Args: + port: Port to listen on + host: Host to bind to (default: 0.0.0.0) + """ + self.port = port + self.host = host + self.app = web.Application() + self.store = dict() + self.lock = asyncio.Lock() + self._loop = None + self._runner = None + self.thread = None + self._setup_routes() + + def run(self): + """Start the server in a background thread.""" + self.thread = threading.Thread(target=self._run_server, daemon=True) + self.thread.start() + logging.info(f"HTTP Metadata Server started on {self.host}:{self.port}") + return self.thread + + def _setup_routes(self): + """Set up the HTTP routes.""" + self.app.router.add_route('*', '/metadata', self._handle_metadata) + + async def _handle_metadata(self, request: web.Request): + """Handle metadata requests.""" + key = request.query.get('key', '') + + if request.method == 'GET': + return await self._handle_get(key) + elif request.method == 'PUT': + return await self._handle_put(key, request) + elif request.method == 'DELETE': + return await self._handle_delete(key) + return web.Response(text='Method not allowed', status=405, + content_type='application/json') + + async def _handle_get(self, key): + """Handle GET requests.""" + async with self.lock: + value = self.store.get(key) + if value is None: + return web.Response(text='metadata not found', status=404, + content_type='application/json') + return web.Response(body=value, status=200, + content_type='application/json') + + async def _handle_put(self, key, request): + """Handle PUT requests.""" + data = await request.read() + async with self.lock: + self.store[key] = data + return web.Response(text='metadata updated', status=200, + content_type='application/json') + + async def _handle_delete(self, key): + """Handle DELETE requests.""" + async with self.lock: + if key not in self.store: + return web.Response(text='metadata not found', status=404, + content_type='application/json') + del self.store[key] + return web.Response(text='metadata deleted', status=200, + content_type='application/json') + + def _run_server(self): + """Run the server in the current thread.""" + try: + # Event Loop + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + + self._runner = web.AppRunner(self.app) + self._loop.run_until_complete(self._runner.setup()) + + site = web.TCPSite(self._runner, host=self.host, port=self.port) + self._loop.run_until_complete(site.start()) + self._loop.run_forever() + except Exception as e: + logging.error(f"Server error: {str(e)}") + finally: + # Cleanup + if self._runner is not None: + self._loop.run_until_complete(self._runner.cleanup()) + if self._loop is not None: + self._loop.close() + + def close(self): + """Shutdown the server.""" + if self._loop is not None and self._loop.is_running(): + self._loop.call_soon_threadsafe(self._loop.stop) + logging.info("Stopping server loop...") + + if self.thread is not None and self.thread.is_alive(): + self.thread.join(timeout=2) + logging.info("Server thread stopped") + + def poll(self) -> KVPoll: + """Poll the server status.""" + if self.thread is None: + return KVPoll.Failed + if not self.thread.is_alive(): + return KVPoll.Failed + return KVPoll.Success + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="HTTP Metadata Server for Mooncake", + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--port", + type=int, + default=8080, + help="Port to listen on" + ) + parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Host to bind to" + ) + parser.add_argument( + "--log-level", + type=str, + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + default="INFO", + help="Logging level" + ) + return parser.parse_args() + + +def main(): + """Main entry point for the mooncake_http_metadata_server command.""" + args = parse_args() + + # Configure logging + logging.basicConfig( + level=getattr(logging, args.log_level), + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + ) + + # Create and start the server + server = KVBootstrapServer(port=args.port, host=args.host) + server.run() + + # Handle graceful shutdown + def signal_handler(sig, frame): + logging.info("Shutting down...") + server.close() + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + # Keep the main thread alive + try: + while True: + sleep(1) + except KeyboardInterrupt: + server.close() + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/mooncake-wheel/pyproject.toml b/mooncake-wheel/pyproject.toml index d95e0f05..5384f03a 100644 --- a/mooncake-wheel/pyproject.toml +++ b/mooncake-wheel/pyproject.toml @@ -10,6 +10,7 @@ authors = [ { name = "Mooncake Authors" } ] requires-python = ">=3.8" +dependencies = ["aiohttp", "requests"] keywords = ["mooncake", "data transfer", "kv cache", "llm inference"] classifiers = [ "Programming Language :: Python :: 3", @@ -26,6 +27,7 @@ Issues = "https://github.com/kvcache-ai/Mooncake/issues" [project.scripts] mooncake_master = "mooncake.cli:main" +mooncake_http_metadata_server = "mooncake.http_metadata_server:main" mc_store_rest_server = "mooncake.mooncake_store_service:main" [tool.setuptools]