From df62a7d8e9c2b5cfecbb6a312eda46020f57b478 Mon Sep 17 00:00:00 2001 From: Cromefire_ Date: Fri, 22 Dec 2023 14:15:20 +0100 Subject: [PATCH] Fixed build and slimmed down container --- Cargo.lock | 1 + crates/tabby/src/worker.rs | 9 +++++++- ee/tabby-webserver/src/hub/api.rs | 2 ++ ee/tabby-webserver/src/hub/mod.rs | 2 ++ ee/tabby-webserver/src/lib.rs | 10 +------- rocm.Dockerfile | 38 ++++++++++++++++++++++--------- 6 files changed, 41 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ac879cdf0ec1..4a2dc9600cce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3712,6 +3712,7 @@ dependencies = [ "filenamify", "glob", "home", + "juniper", "lazy_static", "reqwest", "serde", diff --git a/crates/tabby/src/worker.rs b/crates/tabby/src/worker.rs index f859550026da..3e88b444ad30 100644 --- a/crates/tabby/src/worker.rs +++ b/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, }, ) diff --git a/ee/tabby-webserver/src/hub/api.rs b/ee/tabby-webserver/src/hub/api.rs index 61de89b2eecc..62ce8dc46636 100644 --- a/ee/tabby-webserver/src/hub/api.rs +++ b/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, + pub accelerators: Vec, } pub static REGISTER_WORKER_HEADER: HeaderName = HeaderName::from_static("x-tabby-register-worker"); diff --git a/ee/tabby-webserver/src/hub/mod.rs b/ee/tabby-webserver/src/hub/mod.rs index de459527c856..8dcc3b55bc3a 100644 --- a/ee/tabby-webserver/src/hub/mod.rs +++ b/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)) diff --git a/ee/tabby-webserver/src/lib.rs b/ee/tabby-webserver/src/lib.rs index af79bbf12383..6c38e1217755 100644 --- a/ee/tabby-webserver/src/lib.rs +++ b/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, diff --git a/rocm.Dockerfile b/rocm.Dockerfile index 87ccd51acf11..dddbd373417a 100644 --- a/rocm.Dockerfile +++ b/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 \