Merge pull request #134 from Dstack-TEE/teepod-ui-opt

teepod: Optimize UI

Author: kvinwang

teepod/rpc/proto/teepod_rpc.proto

@@ -98,14 +98,33 @@ message UpgradeAppRequest {
   bytes encrypted_env = 3;
   // Optional update user config. Leave empty to not update.
   string user_config = 4;
+  // Optional update port mapping.
+  bool update_ports = 5;
+  // Port mapping
+  repeated PortMapping ports = 7;
+}
+
+// Message for Status request
+message StatusRequest {
+  // List of VM IDs
+  repeated string ids = 1;
+  // Brief (Don't include VM configuration)
+  bool brief = 2;
+  // Filter by keyword
+  string keyword = 3;
+  // Page number
+  uint32 page = 4;
+  // Page size
+  uint32 page_size = 5;
 }
 
-// Message for VM list response
 message StatusResponse {
   // List of VMs
   repeated VmInfo vms = 1;
   // Port mapping enabled
   bool port_mapping_enabled = 2;
+  // Total number of VMs
+  uint32 total = 3;
 }
 
 message ImageListResponse {
@@ -206,7 +225,7 @@ service Teepod {
   rpc ResizeVm(ResizeVmRequest) returns (google.protobuf.Empty);
 
   // RPC to list all VMs
-  rpc Status(google.protobuf.Empty) returns (StatusResponse);
+  rpc Status(StatusRequest) returns (StatusResponse);
   // RPC to list all available images
   rpc ListImages(google.protobuf.Empty) returns (ImageListResponse);
 

teepod/src/app.rs

@@ -16,7 +16,7 @@ use std::net::IpAddr;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, MutexGuard};
 use supervisor_client::SupervisorClient;
-use teepod_rpc::{self as pb, GpuInfo, VmConfiguration};
+use teepod_rpc::{self as pb, GpuInfo, StatusRequest, StatusResponse, VmConfiguration};
 use tracing::{error, info};
 
 pub use image::{Image, ImageInfo};
@@ -122,6 +122,7 @@ impl App {
         &self,
         work_dir: impl AsRef<Path>,
         cids_assigned: &HashMap<String, u32>,
+        auto_start: bool,
     ) -> Result<()> {
         let vm_work_dir = VmWorkDir::new(work_dir.as_ref());
         let manifest = vm_work_dir.manifest().context("Failed to read manifest")?;
@@ -160,6 +161,9 @@ impl App {
             }
             teapot.add(VmState::new(vm_config));
         };
+        if auto_start && vm_work_dir.started().unwrap_or_default() {
+            self.start_vm(&vm_id).await?;
+        }
         Ok(())
     }
 
@@ -285,7 +289,7 @@ impl App {
                 let entry = entry.context("Failed to read directory entry")?;
                 let vm_path = entry.path();
                 if vm_path.is_dir() {
-                    if let Err(err) = self.load_vm(vm_path, &occupied_cids).await {
+                    if let Err(err) = self.load_vm(vm_path, &occupied_cids, true).await {
                         error!("Failed to load VM: {err:?}");
                     }
                 }
@@ -302,7 +306,7 @@ impl App {
         Ok(())
     }
 
-    pub async fn list_vms(&self) -> Result<Vec<pb::VmInfo>> {
+    pub async fn list_vms(&self, request: StatusRequest) -> Result<StatusResponse> {
         let vms = self
             .supervisor
             .list()
@@ -315,19 +319,43 @@ impl App {
         let mut infos = self
             .lock()
             .iter_vms()
+            .filter(|vm| {
+                if !request.ids.is_empty() && !request.ids.contains(&vm.config.manifest.id) {
+                    return false;
+                }
+                if request.keyword.is_empty() {
+                    true
+                } else {
+                    vm.config.manifest.name.contains(&request.keyword)
+                        || vm.config.manifest.id.contains(&request.keyword)
+                        || vm.config.manifest.app_id.contains(&request.keyword)
+                        || vm.config.manifest.image.contains(&request.keyword)
+                }
+            })
+            .cloned()
+            .collect::<Vec<_>>();
+        infos.sort_by(|a, b| {
+            a.config
+                .manifest
+                .created_at_ms
+                .cmp(&b.config.manifest.created_at_ms)
+        });
+
+        let total = infos.len() as u32;
+        let vms = paginate(infos, request.page, request.page_size)
             .map(|vm| {
                 vm.merged_info(
                     vms.get(&vm.config.manifest.id),
                     &self.work_dir(&vm.config.manifest.id),
                 )
             })
+            .map(|info| info.to_pb(&self.config.gateway))
             .collect::<Vec<_>>();
-
-        infos.sort_by(|a, b| a.manifest.created_at_ms.cmp(&b.manifest.created_at_ms));
-        let gw = &self.config.gateway;
-
-        let lst = infos.into_iter().map(|info| info.to_pb(gw)).collect();
-        Ok(lst)
+        Ok(StatusResponse {
+            vms,
+            port_mapping_enabled: self.config.cvm.port_mapping.enabled,
+            total,
+        })
     }
 
     pub fn list_images(&self) -> Result<Vec<(String, ImageInfo)>> {
@@ -579,6 +607,21 @@ impl App {
     }
 }
 
+fn paginate<T>(items: Vec<T>, page: u32, page_size: u32) -> impl Iterator<Item = T> {
+    let skip;
+    let take;
+    if page == 0 || page_size == 0 {
+        skip = 0;
+        take = items.len();
+    } else {
+        let page = page - 1;
+        let start = page * page_size;
+        skip = start as usize;
+        take = page_size as usize;
+    }
+    items.into_iter().skip(skip).take(take)
+}
+
 #[derive(Clone)]
 pub struct VmState {
     pub(crate) config: Arc<VmConfig>,