Skip to main content

Custom Backend

We have found that a major problem in business is that the preset backends (computational backend/scheduling backend/RPC backend/cross-process backend, etc.) cannot cover all requirements. Typically, extending the backend is not the task of the user, but rather the task of the library developer. Torchpipe holds a different perspective, believing that the backend itself is also an API oriented towards the user. Therefore, the backend must be designed to be simple enough. Drawing on the design of frameworks such as GStreamer and FFmpeg, and targeting modern C++ and Python, Torchpipe hopes that the backend:

  • Is granular enough
  • Does not require any special input/output data types, with the backend author responsible for type validation and management
  • As an API, the user only needs to pay attention to the usage rules of basic types (dict) and the optional logging component, both of which can be re-implemented and globally overridden. No other C++ elements are required. Although this approach can easily lead to code confusion, it greatly gives the user freedom to extend.
  • Provides AOT compilation, especially in Python environments.

Basic Types​

any​

Similar to std::any in C++17, we have defined a type-erased container, ipipe::any, with an almost identical interface. It can accept any type of data except char* and unsigned char*. Data of the any type can be retrieved using any_cast to obtain the actual data. 

T data = any_cast<T>(any_data);
// or
auto* pdata = any_cast<T>(&any_data);

dict​

As a data carrier, similar to Python's dict, we have also defined the following dict in C++:

#ifndef CUSTOM_DICT
using dict = std::shared_ptr<std::unordered_map<std::string, ipipe::any>>;
#else
// todo :  custom dict with bidirectional binding between C++ and Python
#endif

Backend​

Torchpipe limits the basic elements of the backend to:

  • Initialization: parameter configuration
  • Forward: input/output interface
  • Forward: batch range of data

The backend interface is defined as follows:

class Backend {
 public:
  // The configuration will be passed through to the first parameter of `init`
  virtual bool init(const std::unordered_map<std::string, std::string>& config, dict dict_config) {
    return true;
  };

  virtual void forward(const std::vector<dict>& input_dicts) = 0;

  virtual uint32_t max() const { return 1; };
  virtual uint32_t min() const { return 1; };
  ... //Irrelevant Functions
};

You can customize the backend, for example, the internal Identity can be defined equivalently in the following CustomIdentity.cpp file:

#include "torchpipe/extension.h"

namespace ipipe {
class CustomIdentity : public Backend {
 public:
  void forward(const std::vector<dict>& input_dicts) override {
    (*input_dicts[0])[TASK_RESULT_KEY] = input_dicts[0]->at(TASK_DATA_KEY);
  }
  uint32_t max() const override final { return 1; }; // or directly inherit from SingleBackend
};

// Register with the system
IPIPE_REGISTER(Backend, CustomIdentity, "CustomIdentity");
}  // namespace ipipe

Compared to the triton inference server, the backend here has fewer requirements. Users generally only need to extend init/forward. If batching functionality is required, then max/min needs to be modified. In our practice, providing these elements in the backend is sufficient as long as the basic assumptions are met. Torchpipe implements all the functions through the composite of backends. In fact, torchpipe itself is a composite of backends.

Compilation​

For convenient compilation, we use the torch.utils.cpp_extension AOT compilation facility to complete the rapid custom extension of the backend in the Python environment.

from torchpipe.utils.cpp_extension import load
from torchpipe import pipe
## Load C++
load(sources=["CustomIdentity.cpp"])

## Initialization
model = pipe({"backend":"CustomIdentity", "instance_num":"2"})

## Forward
input = {"data":"123"}
model(input)

## Check result
assert(input["result"] == b"123")

Or you can:

tp.utils.cpp_extension.load_filter(
    name = 'Skip', 
    sources='status forward(dict data){return status::Skip;}',
    sources_header="")



tp.utils.cpp_extension.load_backend(
    name = 'identity', 
    sources='void forward(dict data){(*data)["result"] = (*data)["data"];}',
    sources_header="")
model = tp.pipe({"backend":'identity'})
input = {"data":2}
model(input)
assert input["result"] == 2

Binding with Python​

When using Python as the front-end language, the back-end is called from Python and the results are returned to Python, requiring type conversion.

From Python Types to Any​

Assuming the input data is

input = {"data":py_object}

The framework converts the Python object py_object to a C++ object according to certain rules;

For the Python dictionary input, the corresponding C++ code is:

dict input_cpp = std::make_shared<std::unordered_map<std::string, ipipe::any>>();
(*input_cpp)["data"] = PY2CPP(py_object);

After the conversion, the data will be sent to the back-end for processing as a dict. Users can specify the data type T in the back-end and obtain the real data through the following API:

T data = any_cast<T>((*input_cpp)["data"]);
// or
auto* pdata = any_cast<T>(&(*input_cpp)["data"]);

If T is inconsistent with the actual C++ type, an exception will be thrown or a null pointer will be returned.

From any to Python Types​

When the back-end writes the results and other key values to input_cpp, if the Python interface is needed, the framework will automatically transmit the data corresponding to all key values except data to the Python end according to certain rules.