Fixed build and slimmed down container

crates/tabby/src/worker.rs

@@ -91,7 +91,13 @@ struct WorkerContext {
 impl WorkerContext {
     async fn new(kind: WorkerKind, args: &WorkerArgs) -> Self {
         let (cpu_info, cpu_count) = read_cpu_info();
-        let cuda_devices = read_cuda_devices().unwrap_or_default();
+        let accelerators = read_accelerators();
+
+        // For compatibility
+        let mut cuda_devices = vec![];
+        for accelerator in &accelerators {
+            cuda_devices.push(accelerator.display_name.clone());
+        }
 
         Self {
             client: tabby_webserver::public::create_client(
@@ -105,6 +111,7 @@ impl WorkerContext {
                     arch: ARCH.to_string(),
                     cpu_info,
                     cpu_count: cpu_count as i32,
+                    cuda_devices,
                     accelerators,
                 },
             )

ee/tabby-webserver/src/hub/api.rs

@@ -3,6 +3,7 @@ use axum::{headers::Header, http::HeaderName};
 use hyper::Request;
 use serde::{Deserialize, Serialize};
 use tabby_common::api::{
+    accelerator::Accelerator,
     code::{CodeSearch, CodeSearchError, SearchResponse},
     event::RawEventLogger,
 };
@@ -109,6 +110,7 @@ pub struct RegisterWorkerRequest {
     pub cpu_info: String,
     pub cpu_count: i32,
     pub cuda_devices: Vec<String>,
+    pub accelerators: Vec<Accelerator>,
 }
 
 pub static REGISTER_WORKER_HEADER: HeaderName = HeaderName::from_static("x-tabby-register-worker");

ee/tabby-webserver/src/hub/mod.rs

@@ -45,6 +45,7 @@ pub(crate) async fn ws_handler(
 
     let addr = format!("http://{}:{}", addr.ip(), request.port);
 
+    #[allow(deprecated)]
     let worker = Worker {
         name: request.name,
         kind: request.kind,
@@ -54,6 +55,7 @@ pub(crate) async fn ws_handler(
         cpu_info: request.cpu_info,
         cpu_count: request.cpu_count,
         cuda_devices: request.cuda_devices,
+        accelerators: request.accelerators,
     };
 
     ws.on_upgrade(move |socket| handle_socket(state, socket, worker))

ee/tabby-webserver/src/lib.rs

@@ -17,7 +17,6 @@ pub mod public {
 
 use std::sync::Arc;
 
-use api::Hub;
 use axum::{
     extract::State,
     http::Request,
@@ -29,14 +28,7 @@ use juniper_axum::{graphiql, graphql, playground};
 pub use schema::create_schema;
 use schema::{Schema, ServiceLocator};
 use service::create_service_locator;
-use tabby_common::api::{
-    accelerator::{Accelerator, DeviceType},
-    code::{CodeSearch, SearchResponse},
-    event::RawEventLogger,
-};
-use tokio::sync::Mutex;
-use tracing::{error, warn};
-use websocket::WebSocketTransport;
+use tabby_common::api::{code::CodeSearch, event::RawEventLogger};
 
 pub async fn attach_webserver(
     api: Router,

rocm.Dockerfile

@@ -1,12 +1,26 @@
-ARG UBUNTU_VERSION=22.04
-# This needs to generally match the container host's environment.
-ARG ROCM_VERSION=5.7
-# Target the CUDA build image
-ARG BASE_ROCM_DEV_CONTAINER="rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}-complete"
-# Target the CUDA runtime image
-ARG BASE_ROCM_RUN_CONTAINER="rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}-complete"
+# Ubuntu version to be used as base
+ARG UBUNTU_VERSION=jammy
+# URL to the amdgpu-install debian package
+ARG AMDGPU_INSTALL_URL=https://repo.radeon.com/amdgpu-install/6.0/ubuntu/${UBUNTU_VERSION}/amdgpu-install_6.0.60000-1_all.deb
 
-FROM ${BASE_ROCM_DEV_CONTAINER} as build
+FROM ubuntu:${UBUNTU_VERSION} as hipblas_base
+
+ARG AMDGPU_INSTALL_URL
+
+# Install ROCm
+RUN apt-get update &&  \
+    apt-get install -y curl ca-certificates &&  \
+    curl -Lo /tmp/amdgpu-install.deb "${AMDGPU_INSTALL_URL}" && \
+    apt-get install -y /tmp/amdgpu-install.deb && \
+    rm /tmp/amdgpu-install.deb && \
+    apt-get update && \
+    apt-get install -y "hipblas" && \
+    apt-get purge -y curl ca-certificates && \
+    apt-get autoremove -y && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+FROM hipblas_base as build
 
 # Rust toolchain version
 ARG RUST_TOOLCHAIN=stable
@@ -15,23 +29,25 @@ ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         curl \
+        ca-certificates \
         pkg-config \
         libssl-dev \
         protobuf-compiler \
         git \
         cmake \
+        hipblas-dev \
+        build-essential \
         && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
 # setup rust.
-RUN curl https://sh.rustup.rs -sSf | bash -s -- --default-toolchain ${RUST_TOOLCHAIN} -y
+RUN curl https://sh.rustup.rs -sSf | sh -s -- --profile minimal --no-modify-path --default-toolchain ${RUST_TOOLCHAIN} -y
 ENV PATH="/root/.cargo/bin:${PATH}"
 
 WORKDIR /root/workspace
 
 RUN mkdir -p /opt/tabby/bin
-RUN mkdir -p /opt/tabby/lib
 RUN mkdir -p target
 
 COPY . .
@@ -41,7 +57,7 @@ RUN --mount=type=cache,target=/usr/local/cargo/registry \
     cargo build --features rocm --release --package tabby && \
     cp target/release/tabby /opt/tabby/bin/
 
-FROM ${BASE_ROCM_RUN_CONTAINER} as runtime
+FROM hipblas_base as runtime
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \