Open browser to fastapi docs page when running model_serve (#54)

* open browser

* add serving test basic

* add option to automatically open browser

tests/test_serve_model.py

@@ -0,0 +1,22 @@
+from ray import serve
+
+from xinfer.serve import serve_model
+
+
+def test_serve_model():
+    serve_model("vikhyatk/moondream2", blocking=False, open_api_docs=False)
+
+    serve.shutdown()
+
+
+def test_serve_model_custom_deployment():
+    """Test model serving with custom deployment options"""
+    deployment_kwargs = {"num_replicas": 2, "ray_actor_options": {"num_cpus": 2}}
+    handle = serve_model(
+        "vikhyatk/moondream2",
+        deployment_kwargs=deployment_kwargs,
+        blocking=False,
+        open_api_docs=False,
+    )
+    assert handle.deployment_id.name == "XInferModel"
+    serve.shutdown()

xinfer/serve.py

@@ -68,6 +68,7 @@ def serve_model(
     host: str = "127.0.0.1",
     port: int = 8000,
     blocking: bool = True,
+    open_api_docs: bool = True,
     **model_kwargs,
 ):
     deployment_kwargs = deployment_kwargs or {}
@@ -84,12 +85,26 @@ def serve_model(
     app = deployment.bind(model_id, **model_kwargs)
 
     try:
-        handle = serve.run(app, blocking=blocking)
-        if not blocking:
+        handle = serve.run(app)
+        logger.info(f"Open FastAPI docs at http://{host}:{port}/docs")
+        if open_api_docs:
+            import webbrowser
+
+            webbrowser.open(f"http://{host}:{port}/docs")
+
+        if blocking:
+            try:
+                while True:
+                    time.sleep(1)
+            except (KeyboardInterrupt, SystemExit):
+                logger.info("Receiving shutdown signal. Cleaning up...")
+                serve.shutdown()
+        else:
             logger.info(
                 "Running server in non-blocking mode, remember to call serve.shutdown() to stop the server"
             )
-            return handle  # Return handle without shutting down
-    except (KeyboardInterrupt, SystemExit):
-        logger.info("Receiving shutdown signal. Cleaning up...")
+            return handle
+    except Exception as e:
+        logger.error(f"Error starting server: {str(e)}")
         serve.shutdown()
+        raise