handle ctrl_c signal

CUDA11.4.Dockerfile

@@ -1,5 +1,5 @@
-FROM nvidia/cuda:11.4.3-devel-ubuntu20.04 as chef
-RUN apt-get update && apt-get install -y curl
+FROM nvidia/cuda:11.4.3-cudnn8-devel-ubuntu20.04 as chef
+RUN apt-get update && apt-get install -y cmake g++ libclang-dev libssl-dev pkg-config python3-dev curl
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --default-toolchain stable -y
 ENV PATH=/root/.cargo/bin:$PATH
 RUN cargo install cargo-chef
@@ -10,14 +10,12 @@ COPY . .
 RUN cargo chef prepare  --recipe-path recipe.json
 
 FROM chef as builder
-ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update && apt-get install -y cmake g++ libclang-dev libssl-dev pkg-config python3-dev
 COPY --from=planner /app/recipe.json recipe.json
 RUN cargo chef cook --release --recipe-path recipe.json
 COPY . .
 RUN cargo build --release
 
-FROM nvidia/cuda:11.4.3-runtime-ubuntu20.04 as runtime
+FROM nvidia/cuda:11.4.3-cudnn8-runtime-ubuntu20.04 as runtime
 
 RUN apt-get update && apt-get install -y ca-certificates && rm -rf /var/lib/apt/lists/*
 WORKDIR /app

Cargo.toml

@@ -16,7 +16,7 @@ once_cell = "1.18"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_yaml = "0.9"
-tokio = { version = "1.33", features = ["macros", "rt-multi-thread", "sync"] }
+tokio = { version = "1.33", features = ["macros", "rt-multi-thread", "sync", "signal"] }
 tokio-stream = "0.1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

src/main.rs

@@ -70,7 +70,13 @@ async fn main() -> Result<(), std::io::Error> {
     let listener = TcpListener::bind(addr);
     let server = Server::new(listener);
 
-    server.run(app).await
+    select! {
+        res = server.run(app) => res,
+        _ = tokio::signal::ctrl_c() => {
+            tracing::info!("Shutting down");
+            Ok(())
+        },
+    }
 }
 
 #[derive(Deserialize, Debug)]