feat: add port to config

atoma-network · Apr 10, 2024 · 3cf2ff5 · 3cf2ff5
1 parent 1fa6475
commit 3cf2ff5
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/atoma-inference/src/jrpc_server/mod.rs b/atoma-inference/src/jrpc_server/mod.rs
@@ -6,15 +6,15 @@ use tokio::sync::{mpsc, oneshot};
 
 pub type RequestSender = mpsc::Sender<(Value, oneshot::Sender<Value>)>;
 
-pub async fn run(sender: RequestSender) {
+pub async fn run(sender: RequestSender, port: u64) {
     let (shutdown_signal_sender, mut shutdown_signal_receiver) = mpsc::channel::<()>(1);
     let app = Router::new()
         .route("/", post(jrpc_call))
         .route("/healthz", post(healthz))
         .layer(Extension(Arc::new(sender)))
         .layer(Extension(Arc::new(shutdown_signal_sender)));
 
-    let listener = tokio::net::TcpListener::bind("0.0.0.0:21212")
+    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{port}"))
         .await
         .unwrap();
     axum::serve(listener, app)

diff --git a/atoma-inference/src/main.rs b/atoma-inference/src/main.rs
@@ -19,6 +19,7 @@ async fn main() -> Result<(), ModelServiceError> {
         .expect("Incorrect private key bytes length");
 
     let private_key = PrivateKey::from(private_key_bytes);
+    let jrpc_port = model_config.jrpc_port();
     let mut service = ModelService::start(model_config, private_key, req_receiver)
         .expect("Failed to start inference service");
 
@@ -27,7 +28,7 @@ async fn main() -> Result<(), ModelServiceError> {
         Ok::<(), ModelServiceError>(())
     });
 
-    jrpc_server::run(req_sender).await;
+    jrpc_server::run(req_sender, jrpc_port).await;
 
     Ok(())
 }
diff --git a/atoma-inference/src/models/config.rs b/atoma-inference/src/models/config.rs
@@ -62,6 +62,7 @@ pub struct ModelsConfig {
     flush_storage: bool,
     models: Vec<ModelConfig>,
     tracing: bool,
+    jrpc_port: u64,
 }
 
 impl ModelsConfig {
@@ -71,13 +72,15 @@ impl ModelsConfig {
         flush_storage: bool,
         models: Vec<ModelConfig>,
         tracing: bool,
+        jrpc_port: u64,
     ) -> Self {
         Self {
             api_key,
             cache_dir,
             flush_storage,
             models,
             tracing,
+            jrpc_port,
         }
     }
 
@@ -101,6 +104,10 @@ impl ModelsConfig {
         self.tracing
     }
 
+    pub fn jrpc_port(&self) -> u64 {
+        self.jrpc_port
+    }
+
     pub fn from_file_path(config_file_path: PathBuf) -> Self {
         let builder = Config::builder().add_source(config::File::with_name(
             config_file_path.to_str().as_ref().unwrap(),
@@ -137,12 +144,18 @@ impl ModelsConfig {
             .parse()
             .unwrap();
 
+        let jrpc_port = std::env::var("JRPC_PORT")
+            .expect("Failed to retrieve jrpc port from .env file")
+            .parse()
+            .unwrap();
+
         Self {
             api_key,
             cache_dir,
             flush_storage,
             models,
             tracing,
+            jrpc_port,
         }
     }
 }