Merge branch 'main' into ymq

Author: gxuu

pyproject.toml

@@ -18,7 +18,7 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "bidict", 
+    "bidict",
     "cloudpickle",
     "psutil==7.0.0",
     "pycapnp==2.0.0",
@@ -45,8 +45,8 @@ scaler_object_storage_server = "scaler.entry_points.object_storage_server:main"
 [project.optional-dependencies]
 uvloop = ["uvloop"]
 graphblas = ["python-graphblas", "numpy==2.3.1"]
-gui = ["nicegui[plotly]==2.20.0"]
-all = ["python-graphblas", "numpy==2.3.1", "uvloop", "nicegui[plotly]==2.20.0"]
+gui = ["nicegui[plotly]==2.21.1"]
+all = ["python-graphblas", "numpy==2.3.1", "uvloop", "nicegui[plotly]==2.21.1"]
 
 [tool.scikit-build]
 wheel.packages = ["scaler"]
@@ -62,7 +62,7 @@ exclude = [
     "^venv.*$",
     "migrations",
     "__pycache__",
-    "^build.*$",
+    "^build_*_*.*$",
     "^capnproto.*$",
     "^boost.*$",
     "^scripts.*$",
@@ -77,7 +77,7 @@ extend-ignore = ["E203"]
 exclude = [
     "migrations",
     "__pycache__",
-    "./build",
+    "./build_*_*",
     "./capnproto*",
     "./boost*",
     "./scripts",

scaler/cluster/object_storage_server.py

@@ -1,7 +1,6 @@
 import multiprocessing
-import os
 
-from scaler.object_storage.object_storage_server import run_object_storage_server
+from scaler.object_storage.object_storage_server import ObjectStorageServer
 from scaler.utility.object_storage_config import ObjectStorageConfig
 
 
@@ -10,11 +9,12 @@ def __init__(self, storage_address: ObjectStorageConfig):
         multiprocessing.Process.__init__(self, name="ObjectStorageServer")
 
         self._storage_address = storage_address
-        self._on_server_ready_fd = os.eventfd(0, os.EFD_SEMAPHORE)
+
+        self._server = ObjectStorageServer()
 
     def wait_until_ready(self) -> None:
         """Blocks until the object storage server is available to server requests."""
-        os.eventfd_read(self._on_server_ready_fd)
+        self._server.wait_until_ready()
 
     def run(self) -> None:
-        run_object_storage_server(self._storage_address.host, self._storage_address.port, self._on_server_ready_fd)
+        self._server.run(self._storage_address.host, self._storage_address.port)

scaler/entry_points/object_storage_server.py

@@ -1,6 +1,6 @@
 import argparse
 
-from scaler.object_storage.object_storage_server import run_object_storage_server
+from scaler.object_storage.object_storage_server import ObjectStorageServer
 from scaler.utility.object_storage_config import ObjectStorageConfig
 
 
@@ -19,4 +19,4 @@ def get_args():
 
 def main():
     args = get_args()
-    run_object_storage_server(args.address.host, args.address.port)
+    ObjectStorageServer().run(args.address.host, args.address.port)

scaler/io/async_object_storage_connector.py

@@ -85,7 +85,7 @@ async def routine(self):
         if header.response_type != ObjectResponseHeader.ObjectResponseType.GetOK:
             return
 
-        pending_get_future = self._pending_get_requests.get(header.object_id)
+        pending_get_future = self._pending_get_requests.pop(header.object_id, None)
 
         if pending_get_future is None:
             logging.warning(f"unknown get-ok response for unrequested object_id={repr(header.object_id)}.")

scaler/object_storage/CMakeLists.txt

@@ -19,21 +19,28 @@ target_link_libraries(cc_object_storage_server PUBLIC CapnProto::capnp)
 target_include_directories(cc_object_storage_server PUBLIC ${CMAKE_BINARY_DIR})
 
 install(TARGETS cc_object_storage_server
-        LIBRARY DESTINATION scaler/object_storage/)
+        LIBRARY DESTINATION scaler/object_storage/
+        COMPONENT cc_object_storage_server)
 
 find_package(Python3 COMPONENTS Development.Module REQUIRED)
 
 add_library(object_storage_server SHARED)
 
 set_target_properties(object_storage_server PROPERTIES  PREFIX "")
 
+# Forces the .so suffix, as Python only recognizes these (macOS uses .dylib by default)
+set_target_properties(object_storage_server PROPERTIES SUFFIX ".so")
+
 target_sources(object_storage_server PRIVATE pymod_object_storage_server.cpp)
-target_include_directories(object_storage_server PRIVATE ${Python3_INCLUDE_DIRS})
 target_link_libraries(object_storage_server PRIVATE cc_object_storage_server
-                                            PRIVATE ${Python3_LIBRARIES}
-)
+                                            PRIVATE Python3::Module)
 
-target_link_options(object_storage_server PRIVATE "-Wl,-rpath,$ORIGIN")
+if(APPLE)
+    target_link_options(object_storage_server PRIVATE "-Wl,-rpath,@loader_path")
+else()
+    target_link_options(object_storage_server PRIVATE "-Wl,-rpath,$ORIGIN")
+endif()
 
 install(TARGETS object_storage_server
-        LIBRARY DESTINATION scaler/object_storage/)
+        LIBRARY DESTINATION scaler/object_storage/
+        COMPONENT object_storage_server)

scaler/object_storage/io_helper.cpp

@@ -1,5 +1,9 @@
 #include "io_helper.h"
 
+#include <cstdint>
+#include <exception>
+#include <iostream>
+
 #include <capnp/message.h>
 #include <capnp/serialize.h>
 
@@ -15,8 +19,6 @@
 #include <boost/asio/use_awaitable.hpp>
 #include <boost/asio/write.hpp>
 #include <boost/system/system_error.hpp>
-#include <exception>
-#include <iostream>
 
 #include "protocol/object_storage.capnp.h"
 #include "scaler/object_storage/constants.h"
@@ -65,12 +67,12 @@ awaitable<void> read_request_header(tcp::socket& socket, ObjectRequestHeader& he
         if (e.code() == boost::asio::error::eof) {
             std::cerr << "Remote end closed, nothing to read.\n";
         } else {
-            std::cerr << "exception throwned, read error e.what() = " << e.what() << '\n';
+            std::cerr << "exception thrown, read error e.what() = " << e.what() << '\n';
         }
         throw e;
     } catch (std::exception& e) {
         // TODO: make this a log, capnp header corruption is an err.
-        std::cerr << "exception throwned, header not a capnp e.what() = " << e.what() << '\n';
+        std::cerr << "exception thrown, header not a capnp e.what() = " << e.what() << '\n';
 
         throw e;
     }
@@ -93,6 +95,11 @@ awaitable<void> read_request_payload(tcp::socket& socket, ObjectRequestHeader& h
         co_return;
     }
 
+    if (header.payloadLength > SIZE_MAX) {
+        std::cerr << "payload length is larger than SIZE_MAX = " << SIZE_MAX << '\n';
+        std::terminate();
+    }
+
     payload.resize(header.payloadLength);
     try {
         std::size_t n = co_await boost::asio::async_read(socket, boost::asio::buffer(payload), use_awaitable);

scaler/object_storage/io_helper.h

@@ -22,7 +22,5 @@ boost::asio::awaitable<void> read_request_payload(
 boost::asio::awaitable<void> write_response(
     boost::asio::ip::tcp::socket& socket, ObjectResponseHeader& header, std::span<const unsigned char> payload_view);
 
-
-
 };  // namespace object_storage
 };  // namespace scaler

scaler/object_storage/object_storage_server.h

@@ -1,14 +1,18 @@
 #pragma once
 
+#include <algorithm>
+#include <iostream>
+#include <map>
+#include <unistd.h>
+#include <utility>
+
 #include <boost/asio/awaitable.hpp>
 #include <boost/asio/co_spawn.hpp>
 #include <boost/asio/detached.hpp>
 #include <boost/asio/io_context.hpp>
 #include <boost/asio/ip/tcp.hpp>
 #include <boost/asio/signal_set.hpp>
 #include <boost/system/system_error.hpp>
-#include <iostream>
-#include <map>
 
 #include "protocol/object_storage.capnp.h"
 #include "scaler/object_storage/defs.h"
@@ -81,10 +85,10 @@ class ObjectStorageServer {
 
             case reqType::GET_OBJECT: {
                 responseHeader.respType = respType::GET_O_K;
-                if (objectIDToMeta[requestHeader.objectID].object)
-                    responseHeader.payloadLength =
-                        std::min(objectIDToMeta[requestHeader.objectID].object->size(), requestHeader.payloadLength);
-                else
+                if (objectIDToMeta[requestHeader.objectID].object) {
+                    uint64_t objectSize = static_cast<uint64_t>(objectIDToMeta[requestHeader.objectID].object->size());
+                    responseHeader.payloadLength = std::min(objectSize, requestHeader.payloadLength);
+                } else
                     return false;
                 break;
             }
@@ -106,8 +110,8 @@ class ObjectStorageServer {
 private:
     awaitable<void> write_once(Meta meta) {
         if (meta.requestHeader.reqType == reqType::GET_OBJECT) {
-            meta.responseHeader.payloadLength =
-                std::min(objectIDToMeta[meta.responseHeader.objectID].object->size(), meta.requestHeader.payloadLength);
+            uint64_t objectSize = static_cast<uint64_t>(objectIDToMeta[meta.responseHeader.objectID].object->size());
+            meta.responseHeader.payloadLength = std::min(objectSize, meta.requestHeader.payloadLength);
         }
 
         auto payload_view = getMemoryViewForResponsePayload(meta.responseHeader);
@@ -131,7 +135,8 @@ class ObjectStorageServer {
 #ifndef NDEBUG
 public:
 #endif
-    int on_server_ready_fd;
+    int _onServerReadyReader;
+    int _onServerReadyWriter;
 
     std::map<scaler::object_storage::object_id_t, ObjectWithMeta> objectIDToMeta;
     std::map<std::size_t, shared_object_t> objectHashToObject;
@@ -165,31 +170,46 @@ class ObjectStorageServer {
         }
     }
 
-    static int createServerReadyEventfd() {
-        int on_server_ready_fd = eventfd(0, EFD_SEMAPHORE);
-        if (on_server_ready_fd == -1) {
-            std::cerr << "create on_server_ready_fd failed: errno=" << errno << std::endl;
+    void initServerReadyFds() {
+        int pipeFds[2];
+        int ret = pipe(pipeFds);
+
+        if (ret != 0) {
+            std::cerr << "create on server ready FDs failed: errno=" << errno << std::endl;
             std::terminate();
         }
 
-        return on_server_ready_fd;
+        this->_onServerReadyReader = pipeFds[0];
+        this->_onServerReadyWriter = pipeFds[1];
     }
 
-    void setServerReadyEventfd() {
+    void setServerReadyFd() {
         uint64_t value = 1;
-        ssize_t ret = write(this->on_server_ready_fd, &value, sizeof (uint64_t));
+        ssize_t ret = write(this->_onServerReadyWriter, &value, sizeof (uint64_t));
 
         if (ret != sizeof (uint64_t)) {
-            std::cerr << "write to on_server_ready_fd failed: errno=" << errno << std::endl;
+            std::cerr << "write to _onServerReadyWriter failed: errno=" << errno << std::endl;
             std::terminate();
         }
     }
 
+    void closeServerReadyFds() {
+        std::array<int, 2> fds { this->_onServerReadyReader, this->_onServerReadyWriter };
+
+        for (auto fd : fds) {
+            if (close(fd) != 0) {
+                std::cerr << "close failed: errno=" << errno << std::endl;
+                std::terminate();
+            }
+        }
+    }
+
     awaitable<void> listener(boost::asio::ip::tcp::endpoint endpoint) {
+
         auto executor = co_await boost::asio::this_coro::executor;
         tcp::acceptor acceptor(executor, endpoint);
 
-        setServerReadyEventfd();
+        setServerReadyFd();
 
         for (;;) {
             auto shared_socket = std::make_shared<tcp::socket>(executor);
@@ -201,13 +221,12 @@ class ObjectStorageServer {
     }
 
 public:
-    ObjectStorageServer(int on_server_ready_fd = -1)
-    {
-        if (on_server_ready_fd < 0) {
-            this->on_server_ready_fd = createServerReadyEventfd();
-        } else {
-            this->on_server_ready_fd = on_server_ready_fd;
-        }
+    ObjectStorageServer() {
+        this->initServerReadyFds();
+    }
+
+    ~ObjectStorageServer() {
+        this->closeServerReadyFds();
     }
 
     void run(std::string name, std::string port) {
@@ -220,7 +239,6 @@ class ObjectStorageServer {
             signals.async_wait([&](auto, auto) { io_context.stop(); });
 
             co_spawn(io_context, listener(res.begin()->endpoint()), detached);
-
             io_context.run();
         } catch (std::exception& e) {
             std::cerr << "Exception: " << e.what() << std::endl;
@@ -230,10 +248,10 @@ class ObjectStorageServer {
 
     void waitUntilReady() {
         uint64_t value;
-        ssize_t ret = read(this->on_server_ready_fd, &value, sizeof (uint64_t));
+        ssize_t ret = read(this->_onServerReadyReader, &value, sizeof (uint64_t));
 
         if (ret != sizeof (uint64_t)) {
-            std::cerr << "read from on_server_ready_fd failed: errno=" << errno << std::endl;
+            std::cerr << "read from _onServerReadyReader failed: errno=" << errno << std::endl;
             std::terminate();
         }
     }