[Feature] Proposal for design of new memory manager in C++

[PranshuPandya]

Checked other resources

• I added a very descriptive title to this issue.
• I am sure the issue hasn't been already addressed by searching through https://github.com/agiresearch/AIOS/issues.

The feature, motivation and pitch

Please find the proposal doc here: https://docs.google.com/document/d/1ZX3daMbhELwOUsCx46_EWUeGbdsj5RLcaRi3P8dZtKc/edit?usp=sharing

cc: @dongyuanjushi

Alternatives

No response

Additional context

No response

[dongyuanjushi]

The proposal looks good to me. A question would be how will you plan to integrate the c++ code into the python codebase? Could you give some examples of how to integrate?

[do-min-o]

Hi @dongyuanjushi, Thanks for reviewing the proposal. We are planning to use a C++ library called pybind11/pybind11.h. By this library we can do something like this:
Code in C++

// my_class.cpp
#include <pybind11/pybind11.h>
#include <string>

namespace py = pybind11;

class MyClass {
public:
    MyClass(const std::string &name) : name(name) {}

    void set_name(const std::string &new_name) {
        name = new_name;
    }

    std::string get_name() const {
        return name;
    }

    std::string greet() const {
        return "Hello, " + name + "!";
    }

private:
    std::string name;
};

// Expose the class to Python
PYBIND11_MODULE(my_class, m) {
    py::class_<MyClass>(m, "MyClass")
        .def(py::init<const std::string &>())  // Constructor
        .def("set_name", &MyClass::set_name)  // Method
        .def("get_name", &MyClass::get_name)  // Method
        .def("greet", &MyClass::greet);       // Method
}

Usage in Python

from my_class import MyClass

# Create an instance of the C++ class
obj = MyClass("Alice")

# Call methods
print(obj.greet())  # Outputs: Hello, Alice!
obj.set_name("Bob")
print(obj.get_name())  # Outputs: Bob
print(obj.greet())  # Outputs: Hello, Bob!

[do-min-o]

We can also create a wrapper over the C++ class in Python if we need to expose a Python module. Something like this

// core_module.cpp
#include <pybind11/pybind11.h>
#include <string>

namespace py = pybind11;

class CoreClass {
public:
    CoreClass(const std::string& name) : name(name) {}

    void set_name(const std::string& new_name) { name = new_name; }
    std::string get_name() const { return name; }
    std::string greet() const { return "Hello, " + name + "!"; }

private:
    std::string name;
};

PYBIND11_MODULE(core_module, m) {
    py::class_<CoreClass>(m, "CoreClass")
        .def(py::init<const std::string&>())
        .def("set_name", &CoreClass::set_name)
        .def("get_name", &CoreClass::get_name)
        .def("greet", &CoreClass::greet);
}

# wrapper.py

from core_module import CoreClass

class PythonWrapper:
    def __init__(self, name):
        """
        Initialize the Python wrapper.
        """
        self._core_instance = CoreClass(name)  # Use the C++ CoreClass internally

    def set_name(self, name):
        """
        Set the name in the C++ instance.
        """
        if not isinstance(name, str):
            raise ValueError("Name must be a string")
        self._core_instance.set_name(name)

    def get_name(self):
        """
        Get the name from the C++ instance.
        """
        return self._core_instance.get_name()

    def greet(self):
        """
        Get the greeting from the C++ instance.
        """
        return self._core_instance.greet()

    def custom_behavior(self):
        """
        Add Python-only functionality here.
        """
        name = self.get_name()
        return f"Custom Greeting for: {name}"

In any case, our core C++ implementation will not change. So we can move forward by first writing a C++ implementation.

[dongyuanjushi]

This looks sound. You can go ahead for the first iteration of this implementation. During development, it would be necessary to update the installation and launch script as it introduces c++ compilation.

[dongyuanjushi]

Btw it could be necessary to give a the time/space efficiency comparison between the c++ version after integration and the pure python version.

[do-min-o]

Can you elaborate on time/space efficiency comparison? Are we looking for benchmark results between Python and C++ implementation?

[do-min-o]

One quick question on the single memory manager class. I can see a function named address_request, but it only uses a single operation_type. How are we using other functions? It seems we are not, so how do we currently manage memory in the runtime?

[do-min-o]

How is the memory module currently being utilized? The paper indicates that it is intended for tasks like context switching, but at the moment, it doesn't appear to be functioning in that capacity. It would be helpful to have a brief sync to discuss the project's current status.

[dongyuanjushi]

For now, the memory manager is not fully integrated yet. After integration, the agent will call the memory api from cerebrum (https://github.com/agiresearch/Cerebrum) to invoke memory system calls in the memory manager.

[dongyuanjushi]

One quick question on the single memory manager class. I can see a function named address_request, but it only uses a single operation_type. How are we using other functions? It seems we are not, so how do we currently manage memory in the runtime?

The memory system calls are expected to have five different calls: alloc, read, write, delete, retrieve.