bug fixes (#17)

Author: jeffreyqdd

examples/add/main.cpp

@@ -5,7 +5,7 @@ using namespace std;
 int main() {
 	pyscheduler::PyManager manager;
 	pyscheduler::PyManager::InvokeHandler add =
-		manager.getPythonModule("examples.add.python_modules.add", "invoke");
+		manager.loadPythonModule("examples.add.python_modules.add", "invoke");
 	std::cout << add.invoke<int64_t>(3000, -1234) << std::endl;
 	return 0;
 }

examples/add/python_modules/tensor_mul.py

@@ -0,0 +1,9 @@
+import torch
+import torch.utils.dlpack
+
+
+def invoke(a, b):
+    a_tensor = torch.utils.dlpack.from_dlpack(a)
+    b_tensor = torch.utils.dlpack.from_dlpack(b)
+    c_tensor = a_tensor * b_tensor
+    return torch.utils.dlpack.to_dlpack(c_tensor)

include/pyscheduler/details/move_only_impl.hpp

@@ -0,0 +1,26 @@
+#ifdef __INTELLISENSE__
+#	include "pyscheduler/move_only.hpp"
+#endif
+
+#include <memory>
+
+namespace pyscheduler {
+
+template <typename R, typename... Args>
+template <typename F>
+MoveOnlyFunction<R(Args...)>::MoveOnlyFunction(F&& f)
+	: _obj(new F(std::move(f)), [](void* p) { delete static_cast<F*>(p); })
+	, _invoke([](void* p, Args... args) -> R {
+		return (*static_cast<F*>(p))(std::forward<Args>(args)...);
+	}) { }
+
+template <typename R, typename... Args>
+R MoveOnlyFunction<R(Args...)>::operator()(Args... args) {
+	return _invoke(_obj.get(), std::forward<Args>(args)...);
+}
+
+template <typename R, typename... Args>
+MoveOnlyFunction<R(Args...)>::operator bool() const {
+	return (bool)_obj;
+}
+} // namespace pyscheduler

include/pyscheduler/details/pyscheduler_impl.hpp

@@ -1,4 +1,9 @@
-#include "pyscheduler/pyscheduler.hpp"
+#ifdef __INTELLISENSE__
+#	include "pyscheduler/pyscheduler.hpp"
+#endif
+
+#include "pyscheduler/move_only.hpp"
+#include <cassert>
 #include <chrono>
 
 namespace pyscheduler {
@@ -10,27 +15,90 @@ PyManager::InvokeHandler::InvokeHandler(size_t id, std::unique_ptr<PyManager> ma
 	: _id(id)
 	, _manager(std::move(manager)) { }
 
-const std::shared_ptr<std::pair<pybind11::module_, pybind11::object>>
+const std::shared_ptr<std::pair<pybind11::module_, pybind11::object>>&
 PyManager::InvokeHandler::getModuleAndFunc() {
 	// need to lock py_mutex because we don't want a vector resize
 	// to happen during lookup
 
+	// no need to acquire gil because not incrementing python reference count
+
 	// should allow multiple reads concurrently which don't mutate state
 	PyManager::SharedState& state = _manager->shared();
 	std::shared_lock lock(state.py_mutex);
 	return _manager->shared().py_modules.at(_id);
 }
 
+template <typename ReturnType, typename... Args>
+ReturnType PyManager::InvokeHandler::invoke(Args&&... args) {
+	auto mod_and_func = getModuleAndFunc();
+	pybind11::gil_scoped_acquire gil;
+	pybind11::object result = mod_and_func->second(std::forward<Args>(args)...);
+	return result.cast<ReturnType>();
+}
+
+template <typename Callback, typename... Args>
+auto PyManager::InvokeHandler::invoke(Callback&& callback, Args&&... args)
+	-> std::invoke_result_t<Callback, pybind11::object> {
+	auto mod_and_func = getModuleAndFunc();
+	pybind11::gil_scoped_acquire gil;
+	pybind11::object result = mod_and_func->second(std::forward<Args>(args)...);
+	return callback(result);
+}
+template <typename Callback, typename... Args>
+auto PyManager::InvokeHandler::queue_invoke(Callback&& callback, Args&&... args)
+	-> std::future<std::invoke_result_t<Callback, pybind11::object>> {
+	// Need to wrap a promise inside a shared_ptr because Promises are not
+	// copy constructable (requirement enforced by appending to task queue)
+	//
+	// solution was to wrap a promise inside a shared pointer, which is
+	// copy constructable
+	using ReturnType = std::invoke_result_t<Callback, pybind11::object>;
+	using PromisePtr = std::shared_ptr<std::promise<ReturnType>>;
+
+	auto args_tuple = std::make_tuple(std::forward<Args>(args)...);
+	PromisePtr promise_ptr = std::make_shared<std::promise<ReturnType>>();
+	std::future<ReturnType> future = promise_ptr->get_future();
+
+	// Dear reader, I'm sorry
+	// this section creates a closure that executes a python method with
+	// the provided arguments
+	//
+	// the return result from the python function is processed using the
+	// callback function, and the value from that is stored into the
+	// promise.
+
+	auto mod_and_func = getModuleAndFunc();
+	auto method = [this,
+				   callback = std::move(callback),
+				   mod_and_func = std::move(mod_and_func),
+				   args_tuple = std::move(args_tuple),
+				   promise_ptr]() mutable {
+		pybind11::gil_scoped_acquire gil;
+		pybind11::object result = std::apply(
+			[&mod_and_func](auto&&... unpackedArgs) {
+				return mod_and_func->second(std::forward<decltype(unpackedArgs)>(unpackedArgs)...);
+			},
+			args_tuple);
+
+		promise_ptr->set_value(callback(result));
+	};
+	PyManager::shared().task_queue.enqueue(std::move(method));
+	return future;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // Impl PyManager
 ///////////////////////////////////////////////////////////////////////////////
 
+PyManager::SharedState PyManager::_instance;
+
 PyManager::PyManager() {
+	// this lock should be dropped at return so that postcondition (Python Interpreter Initialized)
+	// is guaranteed
 	std::unique_lock lock(shared().py_mutex);
 	if(shared().arc.fetch_add(1) == 0) {
 		shared().main_worker = std::thread(&PyManager::mainLoop, this);
 	}
-
 	// small cost paid to block until interpreter is initalized
 	while(!shared().interpreter_initialized)
 		continue;
@@ -49,8 +117,8 @@ PyManager::~PyManager() {
 	}
 }
 
-PyManager::InvokeHandler PyManager::getPythonModule(const std::string& module_name,
-													const std::string& entry_point) {
+PyManager::InvokeHandler PyManager::loadPythonModule(const std::string& module_name,
+													 const std::string& entry_point) {
 
 	SharedState& state = shared();
 
@@ -67,17 +135,17 @@ PyManager::InvokeHandler PyManager::getPythonModule(const std::string& module_na
 
 		pybind11::gil_scoped_acquire gil;
 
-		pybind11::module_ mod = pybind11::module_::import(module_name.c_str());
-
-		if(!mod) {
-			PyErr_Print();
+		pybind11::module_ mod;
+		try {
+			mod = pybind11::module_::import(module_name.c_str());
+		} catch(pybind11::error_already_set& e) {
 			throw std::invalid_argument("Could not import module: " + module_name);
 		}
 
-		pybind11::object func = mod.attr(entry_point.c_str());
-
-		if(!func) {
-			PyErr_Print();
+		pybind11::object func;
+		try {
+			func = mod.attr(entry_point.c_str());
+		} catch(pybind11::error_already_set& e) {
 			throw std::invalid_argument("Could not find the '" + entry_point +
 										"' method in module " + module_name);
 		}
@@ -89,7 +157,6 @@ PyManager::InvokeHandler PyManager::getPythonModule(const std::string& module_na
 				std::make_pair(mod, func)));
 		state.py_invoke_handler_map[module_name] = id;
 		lock.unlock();
-
 		return PyManager::InvokeHandler(id, std::make_unique<PyManager>());
 	}
 
@@ -130,8 +197,9 @@ void PyManager::mainLoop() {
 			std::vector<std::thread> sub_workers;
 			for(size_t i = 0; i < NUM_WORKERS; i++) {
 				sub_workers.emplace_back(std::thread([i]() {
-					while(shared().threads_active) {
-						std::function<void()> task;
+					// worker should only end if stop signal is set and queue is empty
+					while(shared().threads_active || shared().task_queue.size_approx() > 0) {
+						MoveOnlyFunction<void()> task;
 
 						// have a small timeout so threads can wake up and check if they
 						// need to exit.
@@ -145,7 +213,6 @@ void PyManager::mainLoop() {
 							continue;
 						}
 
-						// std::cout << shared().task_queue.size_approx() << std::endl;
 						task();
 					}
 				}));
@@ -167,10 +234,9 @@ void PyManager::mainLoop() {
 		// we clear and drop all items in queue to safely free
 		// memory.
 		while(shared().task_queue.size_approx() > 0) {
-			std::function<void()> black_box;
+			MoveOnlyFunction<void()> black_box;
 			shared().task_queue.try_dequeue(black_box);
 		}
-
 	} // end python interpreter
 }
 } // namespace pyscheduler

include/pyscheduler/move_only.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <memory>
+
+namespace pyscheduler {
+
+template <typename Signature>
+class MoveOnlyFunction;
+
+template <typename R, typename... Args>
+class MoveOnlyFunction<R(Args...)> {
+public:
+	MoveOnlyFunction() = default;
+
+	template <typename F>
+	MoveOnlyFunction(F&& f);
+
+	MoveOnlyFunction(MoveOnlyFunction&&) noexcept = default;
+	MoveOnlyFunction& operator=(MoveOnlyFunction&&) noexcept = default;
+
+	MoveOnlyFunction(const MoveOnlyFunction&) = delete;
+	MoveOnlyFunction& operator=(const MoveOnlyFunction&) = delete;
+
+	R operator()(Args... args);
+	explicit operator bool() const;
+
+private:
+	using InvokeFn = R (*)(void*, Args&&...);
+	using DestroyFn = void (*)(void*);
+
+	std::unique_ptr<void, DestroyFn> _obj{ nullptr, nullptr };
+	InvokeFn _invoke{ nullptr };
+};
+
+} // namespace pyscheduler
+
+#include "pyscheduler/details/move_only_impl.hpp"