Improve network isolation for serve up

This PR improves the network isolation for the docker compose flow of spinning up the grpc and http-gateway servers. We no longer use the `host` network mode, and instead use the `bridge` network mode. We have also added all the default gRPC and HTTP addresses/ports to constants.
autonomi-ai · Jan 13, 2024 · a35d9ff · a35d9ff
1 parent c7597c0
commit a35d9ff
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 22 deletions.
diff --git a/nos/cli/serve.py b/nos/cli/serve.py
@@ -18,7 +18,7 @@
 from rich.console import Console
 from rich.tree import Tree
 
-from nos.constants import NOS_TMP_DIR
+from nos.constants import DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT, NOS_TMP_DIR
 
 
 NOS_SERVE_TMP_DIR = NOS_TMP_DIR / "serve"
@@ -44,7 +44,10 @@ class ServeOptions:
     http: bool = field(default=False)
     """Whether to use HTTP for the server."""
 
-    http_port: int = field(default=8000)
+    http_host: str = field(default=DEFAULT_HTTP_HOST)
+    """HTTP host to use for the server."""
+
+    http_port: int = field(default=DEFAULT_HTTP_PORT)
     """HTTP port to use for the server."""
 
     http_workers: int = field(default=1)
@@ -75,6 +78,23 @@ class ServeOptions:
     """Environment file to use for the server."""
 
 
+@serve_cli.command("generate", help="Generate the NOS server Dockerfiles, without building it.")
+def _serve_generate(
+    config_filename: str = typer.Option(None, "-c", "--config", help="Serve configuration filename."),
+    target: str = typer.Option(None, "--target", help="Serve a specific target.", show_default=False),
+    tag: str = typer.Option("autonomi/nos:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
+    prod: bool = typer.Option(
+        False,
+        "-p",
+        "--prod",
+        help="Run with production flags (slimmer images, no dev. dependencies).",
+        show_default=False,
+    ),
+) -> None:
+    """Main entrypoint for custom NOS runtime Dockerfile generation."""
+    _serve(config_filename=config_filename, runtime="auto", target=target, tag=tag, generate=True, prod=prod)
+
+
 @serve_cli.command("build", help="Build the NOS server locally.")
 def _serve_build(
     config_filename: str = typer.Option(None, "-c", "--config", help="Serve configuration filename."),
@@ -99,7 +119,8 @@ def _serve_up(
     target: str = typer.Option(None, "--target", help="Serve a specific target.", show_default=True),
     tag: str = typer.Option("autonomi/nos:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
     http: bool = typer.Option(False, "--http", help="Serve with HTTP gateway.", show_default=True),
-    http_port: int = typer.Option(8000, "--http-port", help="HTTP port to use.", show_default=True),
+    http_host: str = typer.Option("0.0.0.0", "--http-host", help="HTTP host to use.", show_default=True),
+    http_port: int = typer.Option(DEFAULT_HTTP_PORT, "--http-port", help="HTTP port to use.", show_default=True),
     http_workers: int = typer.Option(1, "--http-workers", help="HTTP max workers.", show_default=True),
     logging_level: str = typer.Option("INFO", "--logging-level", help="Logging level.", show_default=True),
     home_directory: str = typer.Option(
@@ -132,6 +153,7 @@ def _serve_up(
         target=target,
         tag=tag,
         http=http,
+        http_host=http_host,
         http_port=http_port,
         http_workers=http_workers,
         logging_level=logging_level,
@@ -152,12 +174,14 @@ def _serve(
     target: str = None,
     tag: str = "autonomi/nos:{target}",
     http: bool = False,
-    http_port: int = 8000,
+    http_host: str = DEFAULT_HTTP_HOST,
+    http_port: int = DEFAULT_HTTP_PORT,
     http_workers: int = 1,
     logging_level: str = "INFO",
     home_directory: str = "~/.nosd",
     daemon: bool = False,
     reload: bool = False,
+    generate: bool = False,
     build: bool = False,
     prod: bool = False,
     env_file: str = None,
@@ -303,15 +327,22 @@ def _serve(
                 f"[bold green]✓[/bold green] Successfully generated Dockerfile (target=[bold white]{docker_target}[/bold white], filename=[bold white]{filename}[/bold white])."
             ).add(f"[green]`{cmd}`[/green]")
             print(tree)
-            with redirect_stdout_to_logger(level="DEBUG"):
-                builder.build(filename=filename, target=docker_target, tags=[image_name])
+
+            # If the `--generate` flag is specified, then we do not build the image
+            if not generate:
+                with redirect_stdout_to_logger(level="DEBUG"):
+                    builder.build(filename=filename, target=docker_target, tags=[image_name])
             print(f"[green]✓[/green] Successfully built Docker image (image=[bold white]{image_name}[/bold white]).")
 
         # Copy the dockerfiles to the current working directory if debug is enabled.
         for _docker_target, filename in dockerfiles.items():
-            if debug:
+            if debug or generate:
                 shutil.copyfile(filename, Path.cwd() / Path(filename).name)
 
+        # If the `--generate` flag is specified, then we can stop here
+        if generate:
+            return
+
         # Check if the image was built
         if image_name is None:
             raise ValueError(f"Failed to build target={target}, cannot proceed.")
@@ -356,6 +387,7 @@ def _serve(
         image=image_name,
         gpu=runtime == "gpu",
         http=http,
+        http_host=http_host,
         http_port=http_port,
         http_workers=http_workers,
         http_env="prod" if prod else "dev",

diff --git a/nos/cli/templates/docker-compose.serve.yml.j2 b/nos/cli/templates/docker-compose.serve.yml.j2
@@ -4,10 +4,11 @@ services:
 {%- if http %}
   nos-http-gateway:
     image: {{ image }}
-    command: nos-http-server --port {{ http_port }} --workers {{ http_workers }} {% if reload %}--reload --reload-dir {{ reload_dir }} {% endif %}
+    command: nos-http-server --host {{ http_host }} --port {{ http_port }} --workers {{ http_workers }} {% if reload %}--reload --reload-dir {{ reload_dir }} {% endif %}
     environment:
       - NOS_HOME=/app/.nos
       - NOS_LOGGING_LEVEL={{ logging_level }}
+      - NOS_GRPC_HOST=nos-server
       {%- if http_env %}
       - NOS_HTTP_ENV={{ http_env }}
       {%- endif %}
@@ -25,7 +26,8 @@ services:
       - {{ volume }}
         {%- endfor %}
       {%- endif %}
-    network_mode: host
+    ports:
+      - {{ http_port }}:{{ http_port }}
     ipc: host
     depends_on:
       - nos-server
@@ -52,7 +54,8 @@ services:
       - {{ volume }}
         {%- endfor %}
       {%- endif %}
-    network_mode: host
+    ports:
+      - 50051:50051
     ipc: host
 {%- if gpu %}
     deploy:

diff --git a/nos/constants.py b/nos/constants.py
@@ -15,8 +15,14 @@
 NOS_LOG_DIR.mkdir(parents=True, exist_ok=True)
 NOS_TMP_DIR.mkdir(parents=True, exist_ok=True)
 
+DEFAULT_HTTP_HOST = os.getenv("NOS_HTTP_HOST", "127.0.0.1")
 DEFAULT_HTTP_PORT = int(os.getenv("NOS_HTTP_PORT", 8000))
+DEFAULT_HTTP_ADDRESS = f"{DEFAULT_HTTP_HOST}:{DEFAULT_HTTP_PORT}"
+
+DEFAULT_GRPC_HOST = os.getenv("NOS_GRPC_HOST", "[::]")
 DEFAULT_GRPC_PORT = int(os.getenv("NOS_GRPC_PORT", 50051))
+DEFAULT_GRPC_ADDRESS = f"{DEFAULT_GRPC_HOST}:{DEFAULT_GRPC_PORT}"
+
 GRPC_MAX_MESSAGE_LENGTH = 32 * 1024 * 1024  # 32 MB
 GRPC_MAX_WORKER_THREADS = int(os.getenv("NOS_GRPC_MAX_WORKER_THREADS", 4))
 

diff --git a/nos/server/_service.py b/nos/server/_service.py
@@ -15,7 +15,7 @@
 from nos.common import FunctionSignature, ModelSpec, ModelSpecMetadataCatalog, dumps, loads
 from nos.common.shm import NOS_SHM_ENABLED, SharedMemoryDataDict, SharedMemoryTransportManager
 from nos.constants import (  # noqa F401
-    DEFAULT_GRPC_PORT,  # noqa F401
+    DEFAULT_GRPC_ADDRESS,
     GRPC_MAX_MESSAGE_LENGTH,
     GRPC_MAX_WORKER_THREADS,
     NOS_PROFILING_ENABLED,
@@ -382,7 +382,7 @@ async def Stream(
 
 
 async def async_serve_impl(
-    address: str = f"[::]:{DEFAULT_GRPC_PORT}",
+    address: str = DEFAULT_GRPC_ADDRESS,
     wait_for_termination: bool = True,
     catalog: str = None,
 ):
@@ -420,7 +420,7 @@ async def async_serve_impl(
 
 
 def async_serve(
-    address: str = f"[::]:{DEFAULT_GRPC_PORT}",
+    address: str = DEFAULT_GRPC_ADDRESS,
     max_workers: int = GRPC_MAX_WORKER_THREADS,
     wait_for_termination: bool = True,
     catalog: str = None,
@@ -438,7 +438,7 @@ def main():
     import argparse
 
     parser = argparse.ArgumentParser(description="Inference service")
-    parser.add_argument("-a", "--address", type=str, default=f"[::]:{DEFAULT_GRPC_PORT}", help="gRPC server address")
+    parser.add_argument("-a", "--address", type=str, default=DEFAULT_GRPC_ADDRESS, help="gRPC server address")
     parser.add_argument("-c", "--catalog", type=str, default=None, help="Model catalog")
     args = parser.parse_args()
     logger.debug(f"args={args}")

diff --git a/nos/server/http/_service.py b/nos/server/http/_service.py
@@ -15,8 +15,9 @@
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
 
-from nos.client import DEFAULT_GRPC_PORT, Client
+from nos.client import Client
 from nos.common.tasks import TaskType
+from nos.constants import DEFAULT_GRPC_ADDRESS
 from nos.logging import logger
 from nos.protoc import import_module
 from nos.version import __version__
@@ -62,7 +63,7 @@ class InferenceService:
     version: str = field(default="v1")
     """NOS version."""
 
-    address: str = field(default=f"[::]:{DEFAULT_GRPC_PORT}")
+    address: str = field(default=DEFAULT_GRPC_ADDRESS)
     """gRPC address."""
 
     env: str = field(default=HTTP_ENV)
@@ -132,9 +133,7 @@ def as_path(file: SpooledTemporaryFile, suffix: str, chunk_size_mb: int = 4 * 10
 _model_table: Dict[str, ChatModel] = {}
 
 
-def app_factory(
-    version: str = HTTP_API_VERSION, address: str = f"[::]:{DEFAULT_GRPC_PORT}", env: str = HTTP_ENV
-) -> FastAPI:
+def app_factory(version: str = HTTP_API_VERSION, address: str = DEFAULT_GRPC_ADDRESS, env: str = HTTP_ENV) -> FastAPI:
     """Create a FastAPI factory application for the NOS REST API gateway.
 
     Args:
@@ -469,11 +468,11 @@ def main():
 
     import uvicorn
 
-    from nos.constants import DEFAULT_HTTP_PORT
+    from nos.constants import DEFAULT_HTTP_HOST, DEFAULT_HTTP_PORT
     from nos.logging import logger
 
     parser = argparse.ArgumentParser(description="NOS REST API Service")
-    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host address")
+    parser.add_argument("--host", type=str, default=DEFAULT_HTTP_HOST, help="Host address")
     parser.add_argument("--port", type=int, default=DEFAULT_HTTP_PORT, help="Port number")
     parser.add_argument("--workers", type=int, default=1, help="Number of workers")
     parser.add_argument(

diff --git a/nos/version.py b/nos/version.py
@@ -1 +1 @@
-__version__ = "0.1.4"
+__version__ = "0.1.5"