[backend] add environ to overwrite passed backend

hpcaitech · Oct 1, 2024 · e28a80a · e28a80a
1 parent 51ed242
commit e28a80a
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 4 deletions.
diff --git a/csrc/offload.cpp b/csrc/offload.cpp
@@ -26,9 +26,29 @@ iovec *tensors_to_iovec(const std::vector<at::Tensor> &tensors)
     return iovs;
 }
 
+std::string Offloader::get_default_backend() {
+    const char* env = getenv("TENSORNVME_BACKEND");
+    if (env == nullptr) {
+        return std::string("");
+    }
+    return std::string(env);
+}
+
 Offloader::Offloader(const std::string &filename, unsigned int n_entries, const std::string &backend) : filename(filename), space_mgr(SpaceManager(0))
-{
-    this->aio = create_asyncio(n_entries, backend);
+{   
+    std::string default_backend = get_default_backend();
+    if (default_backend.size() > 0) {
+        if (get_backends().count(default_backend) == 0) {
+            throw std::runtime_error("Cannot find backend: " + default_backend + ", please check if TENSORNVME_BACKEND is set correctly");
+        }
+        this->aio = create_asyncio(n_entries, default_backend);
+    } else {
+        if (get_backends().count(backend) == 0) {
+            throw std::runtime_error("Cannot find backend: " + backend + ", please check the passed backend is set correctly");
+        }
+        this->aio = create_asyncio(n_entries, backend);
+    }
+
     this->fd = open(filename.c_str(), O_RDWR | O_CREAT, S_IRUSR | S_IWUSR);
     this->aio->register_file(fd);
 }

diff --git a/include/offload.h b/include/offload.h
@@ -11,6 +11,7 @@
 #include "aio.h"
 #endif
 
+#include <cstdlib>
 class Offloader
 {
 public:
@@ -31,6 +32,7 @@ class Offloader
     void async_readv(const std::vector<at::Tensor> &tensors, const std::string &key, callback_t callback = nullptr);
     void sync_writev(const std::vector<at::Tensor> &tensors, const std::string &key);
     void sync_readv(const std::vector<at::Tensor> &tensors, const std::string &key);
+    static std::string get_default_backend();
 private:
     const std::string filename;
     int fd;

diff --git a/tensornvme/offload.py b/tensornvme/offload.py
@@ -7,8 +7,6 @@
 
 class DiskOffloader(Offloader):
     def __init__(self, dir_name: str, n_entries: int = 16, backend: str = 'uring') -> None:
-        assert backend in get_backends(
-        ), f'Unsupported backend: {backend}, please install tensornvme with this backend'
         if not os.path.exists(dir_name):
             os.mkdir(dir_name)
         assert os.path.isdir(dir_name)