feat(node-agent): add cni

add cni and format code

Signed-off-by: Martin Moreira de Jesus <[email protected]>
Signed-off-by: Noé Tarbouriech <[email protected]>
Co-Authored-By: Noé Tarbouriech <[email protected]>

Author: Martin Moreira de Jesus

node-agent/src/args.rs

@@ -9,6 +9,9 @@ pub struct CliArguments {
     #[arg(long, default_value_t = 3, env)]
     pub lifecycle_retries: i32,
 
+    /// Enable the orka CNI in containerd.
+    #[arg(long, default_value_t = false, env)]
+    pub cni: bool,
 
     /// The address of the scheduler to connect the node agent to.
     #[arg(long, default_value = "[::]", env)]

node-agent/src/main.rs

@@ -3,27 +3,28 @@ mod workload_manager;
 
 use std::process::exit;
 
-use clap::Parser;
-use orka_proto::scheduler_agent::{lifecycle_service_client::LifecycleServiceClient, ConnectionRequest, status_update_service_client::StatusUpdateServiceClient, DisconnectionNotice};
-use tokio::sync::mpsc::{Sender, Receiver, self};
-use tracing::{info, error, warn};
-use uuid::Uuid;
 use anyhow::Result;
+use clap::Parser;
+use orka_proto::scheduler_agent::{
+    lifecycle_service_client::LifecycleServiceClient,
+    status_update_service_client::StatusUpdateServiceClient, ConnectionRequest,
+    DisconnectionNotice,
+};
+use tokio::sync::mpsc::{self, Receiver, Sender};
+use tracing::{error, info, warn};
 use tracing_log::AsTrace;
+use uuid::Uuid;
 
-use crate::workload_manager::grpc::server::GrpcServer;
 use crate::args::CliArguments;
+use crate::workload_manager::grpc::server::GrpcServer;
 use crate::workload_manager::node::metrics::stream_node_status;
 
 async fn execute_node_lifecycle(
     node_id: Uuid,
     node_agent_port: u16,
     scheduler_connection_string: String,
 ) -> Result<()> {
-    info!(
-        "Connecting to scheduler on {}",
-        scheduler_connection_string
-    );
+    info!("Connecting to scheduler on {}", scheduler_connection_string);
 
     loop {
         tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
@@ -100,7 +101,7 @@ async fn main() {
     info!("Node ID: {}", node_id);
 
     let (tx, mut rx): (Sender<i32>, Receiver<i32>) = mpsc::channel(1);
-    
+
     let lifecycle_tx = tx.clone();
     let lifecycle_connection_string = scheduler_connection_string.clone();
 
@@ -125,7 +126,8 @@ async fn main() {
             }
         }
 
-        error!("
+        error!(
+            "
             Node lifecycle failed, initiating graceful shutdown"
         );
 
@@ -134,9 +136,9 @@ async fn main() {
 
     // start grpc server
     tokio::spawn(async move {
-        error!("Starting gRPC server on {}:{}",
-            args.node_agent_address,
-            args.node_agent_port
+        error!(
+            "Starting gRPC server on {}:{}",
+            args.node_agent_address, args.node_agent_port
         );
 
         let grpc = GrpcServer::new(args.node_agent_address, args.node_agent_port);

node-agent/src/workload_manager/container/client.rs

@@ -1,8 +1,10 @@
+use clap::Parser;
 use containerd_client::{
     connect,
     services::v1::{
-        containers_client::ContainersClient, tasks_client::TasksClient,
-        GetContainerRequest, GetContainerResponse, KillRequest, WaitRequest, WaitResponse, DeleteTaskRequest, DeleteContainerRequest, GetResponse, GetRequest, MetricsRequest, MetricsResponse,
+        containers_client::ContainersClient, tasks_client::TasksClient, DeleteContainerRequest,
+        DeleteTaskRequest, GetContainerRequest, GetContainerResponse, GetRequest, GetResponse,
+        KillRequest, MetricsRequest, MetricsResponse, WaitRequest, WaitResponse,
     },
     with_namespace,
 };
@@ -12,6 +14,7 @@ use tonic::{Request, Status};
 
 use super::error::ContainerClientError;
 
+use crate::args::CliArguments;
 use orka_proto::node_agent::Workload;
 
 const NAMESPACE: &str = "default";
@@ -172,9 +175,16 @@ impl ContainerClient {
             .iter()
             .map(|value| format!("--env={}", value));
 
+        let cni = if CliArguments::parse().cni {
+            vec!["--cni"]
+        } else {
+            vec![]
+        };
+
         // TODO - use containerd library to create container instead of ctr
         let command = Command::new("ctr")
             .arg("run")
+            .args(cni.as_slice())
             .arg("--detach")
             .args(env)
             .arg(&workload.image)