connect to host and xpu.

Signed-off-by: Klaus Ma <klausm@nvidia.com>
openbce · Feb 23, 2024 · cb54cbd · cb54cbd
1 parent afb2bf4
commit cb54cbd
Show file tree

Hide file tree

Showing 7 changed files with 143 additions and 56 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/shim/Cargo.toml b/shim/Cargo.toml
@@ -19,6 +19,7 @@ prost = "0.12"
 prost-types = "0.12"
 prost-build = "0.12"
 async-trait = "0.1"
+tower = "0.4"
 
 [build-dependencies]
 tonic-build = { workspace = true }
diff --git a/shim/src/cfg/mod.rs b/shim/src/cfg/mod.rs
@@ -17,10 +17,11 @@ use clap::Parser;
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 pub struct Options {
-    // The address of Unix socket for Chariot shim.
-    // #[arg(short, long, default_value=cri::DEFAULT_UNIX_SOCKET)]
-    // pub address: String,
     /// The address of CRI server in XPU.
-    #[arg(short, long)]
-    pub xpu_address: String,
+    #[arg(long)]
+    pub xpu_cri: String,
+
+    /// The address of CRI server in host.
+    #[arg(long)]
+    pub host_cri: String,
 }
diff --git a/shim/src/cri/image.rs b/shim/src/cri/image.rs
@@ -11,7 +11,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use tokio::net::UnixStream;
 use tonic::transport::Channel;
+use tonic::transport::{Endpoint, Uri};
+use tower::service_fn;
 use tracing::info;
 
 use self::crirpc::image_service_client::ImageServiceClient;
@@ -26,30 +29,59 @@ use crate::common::ChariotError;
 
 pub struct ImageShim {
     pub xpu_client: ImageServiceClient<Channel>,
+    pub host_client: ImageServiceClient<Channel>,
 }
 
 impl ImageShim {
-    pub async fn connect(cri_addr: String) -> Result<Self, ChariotError> {
-        let mut image_client = ImageServiceClient::connect(cri_addr)
+    pub async fn connect(host_cri: String, xpu_cri: String) -> Result<Self, ChariotError> {
+        let mut xpu_client = ImageServiceClient::connect(xpu_cri)
             .await
             .map_err(|e| ChariotError::NetworkError(e.to_string()))?;
 
-        let resp = image_client
+        // Log XPU image FS info
+        let resp = xpu_client
             .image_fs_info(ImageFsInfoRequest {})
             .await
             .map_err(|e| ChariotError::CriError(e.to_string()))?;
         let fs_info = resp.into_inner();
 
         for fs in fs_info.container_filesystems {
-            info!("Container FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
+            info!("XPU container FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
         }
 
         for fs in fs_info.image_filesystems {
-            info!("Image FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
+            info!("XPU image FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
+        }
+
+        let channel = Endpoint::try_from("http://[::]:50051")
+            .map_err(|e| ChariotError::NetworkError(e.to_string()))?
+            .connect_with_connector(service_fn(move |_: Uri| {
+                let host_path = host_cri.clone();
+                UnixStream::connect(host_path)
+            }))
+            .await
+            .map_err(|e| ChariotError::NetworkError(e.to_string()))?;
+
+        let mut host_client = ImageServiceClient::new(channel);
+
+        // Log host image FS info
+        let resp = host_client
+            .image_fs_info(ImageFsInfoRequest {})
+            .await
+            .map_err(|e| ChariotError::CriError(e.to_string()))?;
+        let fs_info = resp.into_inner();
+
+        for fs in fs_info.container_filesystems {
+            info!("Host container FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
+        }
+
+        for fs in fs_info.image_filesystems {
+            info!("Host image FS: {:?}", fs.fs_id.map(|i| i.mountpoint),);
         }
 
         Ok(ImageShim {
-            xpu_client: image_client,
+            xpu_client,
+            host_client,
         })
     }
 }
@@ -60,7 +92,7 @@ impl crirpc::image_service_server::ImageService for ImageShim {
         &self,
         request: tonic::Request<ListImagesRequest>,
     ) -> Result<tonic::Response<ListImagesResponse>, tonic::Status> {
-        let mut client = self.xpu_client.clone();
+        let mut client = self.host_client.clone();
 
         client.list_images(request).await
     }
@@ -69,7 +101,7 @@ impl crirpc::image_service_server::ImageService for ImageShim {
         &self,
         request: tonic::Request<ImageStatusRequest>,
     ) -> Result<tonic::Response<ImageStatusResponse>, tonic::Status> {
-        let mut client = self.xpu_client.clone();
+        let mut client = self.host_client.clone();
 
         client.image_status(request).await
     }
@@ -78,7 +110,7 @@ impl crirpc::image_service_server::ImageService for ImageShim {
         &self,
         request: tonic::Request<PullImageRequest>,
     ) -> Result<tonic::Response<PullImageResponse>, tonic::Status> {
-        let mut client = self.xpu_client.clone();
+        let mut client = self.host_client.clone();
 
         client.pull_image(request).await
     }
@@ -87,7 +119,7 @@ impl crirpc::image_service_server::ImageService for ImageShim {
         &self,
         request: tonic::Request<RemoveImageRequest>,
     ) -> Result<tonic::Response<RemoveImageResponse>, tonic::Status> {
-        let mut client = self.xpu_client.clone();
+        let mut client = self.host_client.clone();
 
         client.remove_image(request).await
     }
@@ -96,7 +128,7 @@ impl crirpc::image_service_server::ImageService for ImageShim {
         &self,
         request: tonic::Request<ImageFsInfoRequest>,
     ) -> Result<tonic::Response<ImageFsInfoResponse>, tonic::Status> {
-        let mut client = self.xpu_client.clone();
+        let mut client = self.host_client.clone();
 
         client.image_fs_info(request).await
     }