Thanks Code Rabbit

Author: grahamking

components/frontend/src/dynamo/frontend/main.py

@@ -166,6 +166,8 @@ def parse_args():
 
     if flags.static_endpoint and (not flags.model_name or not flags.model_path):
         parser.error("--static-endpoint requires both --model-name and --model-path")
+    if bool(flags.tls_cert_path) ^ bool(flags.tls_key_path):  # ^ is XOR
+        parser.error("--tls-cert-path and --tls-key-path must be provided together")
 
     return flags
 

launch/dynamo-run/src/flags.rs

@@ -50,11 +50,11 @@ pub struct Flags {
     pub http_port: u16,
 
     /// TLS certificate file
-    #[arg(long)]
+    #[arg(long, requires = "tls_key_path")]
     pub tls_cert_path: Option<PathBuf>,
 
     /// TLS certificate key file
-    #[arg(long)]
+    #[arg(long, requires = "tls_cert_path")]
     pub tls_key_path: Option<PathBuf>,
 
     /// The name of the model we are serving

lib/bindings/python/Cargo.lock

@@ -10,6 +10,7 @@ use dynamo_llm::entrypoint::input::Input;
 use dynamo_llm::entrypoint::EngineConfig as RsEngineConfig;
 use dynamo_llm::entrypoint::RouterConfig as RsRouterConfig;
 use dynamo_llm::kv_router::KvRouterConfig as RsKvRouterConfig;
+use dynamo_llm::local_model::DEFAULT_HTTP_PORT;
 use dynamo_llm::local_model::{LocalModel, LocalModelBuilder};
 use dynamo_llm::mocker::protocols::MockEngineArgs;
 use dynamo_runtime::protocols::Endpoint as EndpointId;
@@ -128,6 +129,13 @@ impl EntrypointArgs {
             })?),
             None => None,
         };
+        if (tls_cert_path.is_some() && tls_key_path.is_none())
+            || (tls_cert_path.is_none() && tls_key_path.is_some())
+        {
+            return Err(pyo3::exceptions::PyValueError::new_err(
+                "tls_cert_path and tls_key_path must be provided together",
+            ));
+        }
         Ok(EntrypointArgs {
             engine_type,
             model_path,
@@ -138,7 +146,7 @@ impl EntrypointArgs {
             template_file,
             router_config,
             kv_cache_block_size,
-            http_port: http_port.unwrap_or(0),
+            http_port: http_port.unwrap_or(DEFAULT_HTTP_PORT),
             tls_cert_path,
             tls_key_path,
             extra_engine_args,

lib/bindings/python/rust/llm/entrypoint.rs

@@ -37,7 +37,13 @@ pub async fn run(runtime: Runtime, engine_config: EngineConfig) -> anyhow::Resul
                 .tls_key_path(Some(tls_key_path.to_path_buf()))
                 .port(local_model.http_port())
         }
-        (_, _) => service_v2::HttpService::builder().port(local_model.http_port()),
+        (None, None) => service_v2::HttpService::builder().port(local_model.http_port()),
+        (_, _) => {
+            // CLI should prevent us ever getting here
+            anyhow::bail!(
+                "Both --tls-cert-path and --tls-key-path must be provided together to enable TLS"
+            );
+        }
     };
     http_service_builder =
         http_service_builder.with_request_template(engine_config.local_model().request_template());

lib/llm/src/entrypoint/input/http.rs

@@ -199,7 +199,7 @@ impl HttpService {
     pub async fn run(&self, cancel_token: CancellationToken) -> Result<()> {
         let address = format!("{}:{}", self.host, self.port);
         let protocol = if self.enable_tls { "HTTPS" } else { "HTTP" };
-        tracing::info!(address, "Starting {protocol} service on: {address}");
+        tracing::info!(protocol, address, "Starting HTTP(S) service");
 
         let router = self.router.clone();
         let observer = cancel_token.child_token();
@@ -220,26 +220,27 @@ impl HttpService {
 
             // aws_lc_rs is the default but other crates pull in `ring` also,
             // so rustls doesn't know which one to use. Tell it.
-            let rustls_out = rustls::crypto::aws_lc_rs::default_provider().install_default();
-            // The Err type is `Arc<Self>`. Very strange.
-            if rustls_out.is_err() {
-                anyhow::bail!(
-                    "Failed installing AWS-LC-RS as the TLS crypto provider. Cannot start web server."
-                );
+            if let Err(e) = rustls::crypto::aws_lc_rs::default_provider().install_default() {
+                tracing::debug!("TLS crypto provider already installed: {e:?}");
             }
 
             let config = RustlsConfig::from_pem_file(cert_path, key_path)
                 .await
                 .map_err(|e| anyhow::anyhow!("Failed to create TLS config: {}", e))?;
 
-            let server = axum_server::bind_rustls(addr, config).serve(router.into_make_service());
+            let handle = axum_server::Handle::new();
+            let server = axum_server::bind_rustls(addr, config)
+                .handle(handle.clone())
+                .serve(router.into_make_service());
 
             tokio::select! {
                 result = server => {
                     result.map_err(|e| anyhow::anyhow!("HTTPS server error: {}", e))?;
                 }
                 _ = observer.cancelled() => {
                     tracing::info!("HTTPS server shutdown requested");
+                    handle.graceful_shutdown(Some(Duration::from_secs(5)));
+                    // TODO: Do we need to wait?
                 }
             }
         } else {

lib/llm/src/http/service/service_v2.rs

@@ -38,7 +38,8 @@ const DEFAULT_NAME: &str = "dynamo";
 const DEFAULT_KV_CACHE_BLOCK_SIZE: u32 = 16;
 
 /// We can't have it default to 0, so pick something
-const DEFAULT_HTTP_PORT: u16 = 8080;
+/// 'pub' because the bindings use it for consistency.
+pub const DEFAULT_HTTP_PORT: u16 = 8080;
 
 pub struct LocalModelBuilder {
     model_path: Option<PathBuf>,

lib/llm/src/local_model.rs

@@ -276,14 +276,11 @@ impl Component {
         let component_clone = self.clone();
         let mut hierarchies = self.parent_hierarchy();
         hierarchies.push(self.hierarchy());
-        debug_assert_eq!(
-            hierarchies
-                .last()
-                .cloned()
-                .unwrap_or_default()
-                .to_lowercase(),
-            self.service_name().to_lowercase()
-        ); // it happens that in component, hierarchy and service name are the same
+        debug_assert!(hierarchies
+            .last()
+            .map(|x| x.as_str())
+            .unwrap_or_default()
+            .eq_ignore_ascii_case(&self.service_name())); // it happens that in component, hierarchy and service name are the same
 
         // Register a metrics callback that scrapes component statistics
         let metrics_callback = Arc::new(move || {