Merge pull request #2 from zihaoccc/device_info

fetch device info and including in network registry
This commit is contained in:
Zihao Chen 2024-07-18 00:46:52 -05:00 committed by GitHub
commit c4b159958b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 21 additions and 3 deletions

View file

@ -183,18 +183,23 @@ void printAntigmaLogo() {
)" << '\n';
}
static ggml_backend_t create_backend() {
static ggml_backend_t create_backend(Registry& registry) {
ggml_backend_t backend = NULL;
printAntigmaLogo();
#ifdef GGML_USE_CUDA
std::string device_info = "";
fprintf(stderr, "%s: using CUDA backend\n", __func__);
backend = ggml_backend_cuda_init(0); // init device 0
get_cuda_device_info(0, *device_info);
registry.add_field("device_info", device_info);
if (!backend) {
fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);
}
#elif GGML_USE_METAL
fprintf(stderr, "%s: using Metal backend\n", __func__);
backend = ggml_backend_metal_init();
std::string device_info = get_metal_device_info();
registry.add_field("device_info", device_info);
if (!backend) {
fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);
}
@ -235,7 +240,7 @@ int main(int argc, char * argv[]) {
fprintf(stderr, "Invalid parameters\n");
return 1;
}
ggml_backend_t backend = create_backend();
ggml_backend_t backend = create_backend(registry);
if (!backend) {
fprintf(stderr, "Failed to create backend\n");
return 1;
@ -253,7 +258,7 @@ int main(int argc, char * argv[]) {
if (registry.register_with_central(endpoint)) {
print_colored("Registered successfully", "32");
} else {
print_colored("Registered successfully", "31");
print_colored("Registered failed", "31");
}
start_rpc_server(backend, endpoint.c_str(), free_mem, total_mem);
ggml_backend_free(backend);

View file

@ -44,6 +44,8 @@ GGML_API void ggml_backend_metal_log_set_callback(ggml_log_callback log_callback
GGML_API ggml_backend_t ggml_backend_metal_init(void);
GGML_API char* get_metal_device_info(void);
GGML_API bool ggml_backend_is_metal(ggml_backend_t backend);
GGML_API GGML_CALL ggml_backend_buffer_t ggml_backend_metal_buffer_from_ptr(void * data, size_t size, size_t max_size);

View file

@ -2960,6 +2960,10 @@ static ggml_guid_t ggml_backend_cuda_guid() {
return &guid;
}
static void get_cuda_device_info(int device, std::string& info) {
info = "test";
}
GGML_CALL ggml_backend_t ggml_backend_cuda_init(int device) {
if (device < 0 || device >= ggml_backend_cuda_get_device_count()) {
GGML_CUDA_LOG_ERROR("%s: invalid device %d\n", __func__, device);

View file

@ -289,6 +289,13 @@ static void * ggml_metal_host_malloc(size_t n) {
return data;
}
char* get_metal_device_info(void) {
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
char* info = [[device name] UTF8String];
return info;
}
static struct ggml_metal_context * metal_init(int n_cb) {
GGML_METAL_LOG_INFO("%s: allocating\n", __func__);