Kye
10 months ago
parent
57eb57e7bd
commit
0ddd37514c
@ -0,0 +1,45 @@
|
||||
# Import necessary libraries
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from pydantic import BaseModel
|
||||
|
||||
# from swarms import ToolAgent
|
||||
from swarms.utils.json_utils import base_model_schema_to_json
|
||||
|
||||
# Load the pre-trained model and tokenizer
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"databricks/dolly-v2-12b",
|
||||
load_in_4bit=True,
|
||||
device_map="auto",
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
|
||||
|
||||
|
||||
class Schema(BaseModel):
|
||||
name: str
|
||||
agent: int
|
||||
is_student: bool
|
||||
courses: list[str]
|
||||
|
||||
|
||||
json_schema = str(base_model_schema_to_json(Schema))
|
||||
print(json_schema)
|
||||
|
||||
# # Define the task to generate a person's information
|
||||
# task = (
|
||||
# "Generate a person's information based on the following schema:"
|
||||
# )
|
||||
|
||||
# # Create an instance of the ToolAgent class
|
||||
# agent = ToolAgent(
|
||||
# name="dolly-function-agent",
|
||||
# description="Ana gent to create a child data",
|
||||
# model=model,
|
||||
# tokenizer=tokenizer,
|
||||
# json_schema=json_schema,
|
||||
# )
|
||||
|
||||
# # Run the agent to generate the person's information
|
||||
# generated_data = agent.run(task)
|
||||
|
||||
# # Print the generated data
|
||||
# print(f"Generated data: {generated_data}")
|
||||
@ -1,93 +0,0 @@
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::wrap_pyfunction;
|
||||
use pyo3::types::IntoPyDict;
|
||||
use rayon::{ThreadPool, ThreadPoolBuilder};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use std::thread;
|
||||
|
||||
|
||||
#[pymodule]
|
||||
fn rust_module(py: Python, m: &PyModule) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(concurrent_exec, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// This function wraps Python code in Rust concurrency for ultra high performance.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `py_codes` - A vector of string slices that holds the Python codes to be executed.
|
||||
/// * `timeout` - An optional duration to specify a timeout for the Python code execution.
|
||||
/// * `num_threads` - The number of threads to use for executing the Python code.
|
||||
/// * `error_handler` - A function to handle errors during Python code execution.
|
||||
/// * `log_function` - A function to log the execution of the Python code.
|
||||
/// * `result_handler` - A function to handle the results of the Python code execution.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```
|
||||
/// let py_codes = vec!["print('Hello, World!')", "print('Hello, Rust!')"];
|
||||
/// let timeout = Some(Duration::from_secs(5));
|
||||
/// let num_threads = 4;
|
||||
/// let error_handler = |e| eprintln!("Error: {}", e);
|
||||
/// let log_function = |s| println!("Log: {}", s);
|
||||
/// let result_handler = |r| println!("Result: {:?}", r);
|
||||
/// execute_python_codes(py_codes, timeout, num_threads, error_handler, log_function, result_handler);
|
||||
/// ```
|
||||
|
||||
#[pyfunction]
|
||||
pub fn concurrent_exec<F, G, H>(
|
||||
py_codes: Vec<&str>,
|
||||
timeout: Option<Duration>,
|
||||
num_threads: usize,
|
||||
error_handler: F,
|
||||
log_function: G,
|
||||
result_handler: H,
|
||||
) -> PyResult<Vec<PyResult<()>>>
|
||||
where
|
||||
F: Fn(&str),
|
||||
G: Fn(&str),
|
||||
H: Fn(&PyResult<()>),
|
||||
{
|
||||
let gil = Python::acquire_gil();
|
||||
let py = gil.python();
|
||||
let py_codes = Arc::new(Mutex::new(py_codes));
|
||||
let results = Arc::new(Mutex::new(Vec::new()));
|
||||
let pool = ThreadPool::new(num_threads);
|
||||
|
||||
pool.install(|| {
|
||||
py_codes.par_iter().for_each(|code| {
|
||||
let locals = [("__name__", "__main__")].into_py_dict(py);
|
||||
let globals = [("__name__", "__main__")].into_py_dict(py);
|
||||
|
||||
log_function(&format!("Executing Python code: {}", code));
|
||||
let result = py.run(code, Some(globals), Some(locals));
|
||||
|
||||
match timeout {
|
||||
Some(t) => {
|
||||
let now = Instant::now();
|
||||
let timeout_thread = thread::spawn(move || {
|
||||
while now.elapsed() < t {
|
||||
if let Ok(_) = result {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if now.elapsed() >= t {
|
||||
error_handler(&format!("Python code execution timed out: {}", code));
|
||||
}
|
||||
});
|
||||
|
||||
timeout_thread.join().unwrap();
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
results.lock().unwrap().push(result.clone(result));
|
||||
result_handler(&result);
|
||||
});
|
||||
});
|
||||
|
||||
pool.join();
|
||||
Ok(results.lock().unwrap().clone())
|
||||
}
|
||||
@ -1,71 +0,0 @@
|
||||
use pyo3::prelude::*;
|
||||
use rustacuda::prelude::*;
|
||||
use rustacuda::memory::DeviceBox;
|
||||
use std::error::Error;
|
||||
use std::ffi::CString;
|
||||
|
||||
#[pymodule]
|
||||
fn rust_cuda(_py: Python, m: &PyModule) -> PyResult<()> {
|
||||
#[pyfn(m, "execute_on_device")]
|
||||
fn execute_on_device(py: Python, device_id: u32, a: f32, b: f32) -> PyResult<f32> {
|
||||
/// The result of executing the CUDA operation.
|
||||
let result = py.allow_threads(|| {
|
||||
execute_cuda(device_id, a, b)
|
||||
});
|
||||
match result {
|
||||
Ok(res) => Ok(res),
|
||||
Err(err) => Err(PyErr::new::<pyo3::exceptions::PyException, _>(format!("{}", err))),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn execute_cuda(device_id: u32, a: f32, b: f32) -> Result<f32, Box<dyn Error>> {
|
||||
rustacuda::init(CudaFlags::empty())?;
|
||||
let device = Device::get_device(device_id)?;
|
||||
/// Creates a new CUDA context and pushes it onto the current thread's stack.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `flags` - The flags to be used when creating the context.
|
||||
/// * `device` - The device on which the context will be created.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// The newly created CUDA context.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if the context creation fails.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use swarms::cuda_wrapper::Context;
|
||||
///
|
||||
/// let device = 0;
|
||||
/// let context = Context::create_and_push(ContextFlags::MAP_HOST | ContextFlags::SCHED_AUTO, device)?;
|
||||
/// ```
|
||||
pub fn create_and_push(flags: ContextFlags, device: i32) -> Result<Context, CudaError> {
|
||||
// implementation goes here
|
||||
}
|
||||
let context = Context::create_and_push(ContextFlags::MAP_HOST | ContextFlags::SCHED_AUTO, device)?;
|
||||
let module_data = CString::new(include_str!("../resources/add.ptx"))?;
|
||||
let module = Module::load_from_string(&module_data)?;
|
||||
let stream = Stream::new(StreamFlags::NON_BLOCKING, None)?;
|
||||
let mut x = DeviceBox::new(&a)?;
|
||||
let mut y = DeviceBox::new(&b)?;
|
||||
let mut result = DeviceBox::new(&0.0f32)?;
|
||||
unsafe {
|
||||
launch!(module.sum<<<1, 1, 0, stream>>>(
|
||||
x.as_device_ptr(),
|
||||
y.as_device_ptr(),
|
||||
result.as_device_ptr(),
|
||||
1
|
||||
))?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
let mut result_host = 0.0f32;
|
||||
result.copy_to(&mut result_host)?;
|
||||
Ok(result_host)
|
||||
}
|
||||
@ -1,37 +0,0 @@
|
||||
use std::fs::File;
|
||||
use std::io::prelude::*;
|
||||
use std::time::Instant;
|
||||
use std::io::{BufReader, io};
|
||||
use ranyon::prelude::{IntoParallelRefIterator, ParallelIterator};
|
||||
|
||||
fn read_file(path: &str) -> Vec<String> {
|
||||
/// Reads the contents of a file located at the specified path.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `path` - The path to the file.
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A `Result` containing a vector of strings representing the lines of the file if the file was successfully read,
|
||||
/// or an `io::Error` if there was an error reading the file.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```
|
||||
/// use std::io;
|
||||
/// use std::fs::File;
|
||||
/// use std::io::BufReader;
|
||||
///
|
||||
/// fn read_file(path: &str) -> io::Result<Vec<String>> {
|
||||
/// let contents: io::Result<Vec<String>> = BufReader::new(File::open(path).expect("Could not open file"))
|
||||
/// .lines()
|
||||
/// .collect();
|
||||
/// contents
|
||||
/// }
|
||||
/// ```
|
||||
let contents: io::Result<Vec<String>> = BufReader::new(File::open(path).expect("Could not open file"))
|
||||
.lines()
|
||||
.collect();
|
||||
return contents.expect("Could not read file");
|
||||
}
|
||||
@ -1,113 +0,0 @@
|
||||
/// This module provides a multi-threading processor for executing Python modules and functions in parallel.
|
||||
/// It utilizes the `rayon` crate for parallel processing and the `pyo3` crate for interacting with the Python interpreter.
|
||||
/// The `multithreading_processor` function takes a vector of `PythonModule` structs and the number of threads to use.
|
||||
/// Each `PythonModule` struct contains the name of the Python module, the name of the function to call, and any arguments to pass to the function.
|
||||
/// The function imports the Python module, calls the specified function, and sends any errors encountered back to the main thread.
|
||||
/// If an import error occurs, a `PythonError::ImportError` is returned.
|
||||
/// If a function call error occurs, a `PythonError::FunctionError` is returned.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::wrap_pyfunction;
|
||||
use rayon::prelude::*;
|
||||
use std::sync::mpsc::{channel, Sender};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use log::{info, error};
|
||||
|
||||
struct PythonModule<'a> {
|
||||
name: &'a str,
|
||||
function: &'a str,
|
||||
}
|
||||
|
||||
enum PythonError {
|
||||
ImportError(String),
|
||||
FunctionError(String),
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn my_module(py: Python, m: &PyModule) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(process_python_modules, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
|
||||
/// The function returns `Ok(())` if all modules are processed successfully.
|
||||
/// Note: This code assumes that the necessary dependencies (`pyo3`, `rayon`, `log`) are already imported and initialized.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `modules` - A vector of `PythonModule` structs representing the Python modules and functions to execute.
|
||||
/// * `num_threads` - The number of threads to use for parallel processing.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use pyo3::types::PyModule;
|
||||
/// use pyo3::types::PyResult;
|
||||
/// use pyo3::prelude::*;
|
||||
///
|
||||
/// struct PythonModule<'a> {
|
||||
/// name: &'a str,
|
||||
/// function: &'a str,
|
||||
/// args: Vec<&'a str>,
|
||||
/// }
|
||||
///
|
||||
/// #[pymodule]
|
||||
/// fn multithreading_processor(modules: Vec<PythonModule>, num_threads: usize) -> Result<(), PythonError> {
|
||||
/// // Function implementation
|
||||
/// Ok(())
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns a `PythonError` if an import error or a function call error occurs.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// This function does not panic.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// This function is safe to call, but it assumes that the necessary dependencies (`pyo3`, `rayon`, `log`) are already imported and initialized.
|
||||
// Initialize Python interpreter
|
||||
#[pyfunction]
|
||||
fn process_python_modules(modules: Vec<PythonModule>, num_threads: usize) -> Result<(), PythonError> {
|
||||
|
||||
let gil = Python::acquire_gil();
|
||||
let py = gil.python();
|
||||
|
||||
// Set the global thread pool's configuration
|
||||
rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(num_threads)
|
||||
.build_global()
|
||||
.unwrap();
|
||||
|
||||
// Create a channel to send errors from threads to the main thread
|
||||
let (tx, rx) = channel();
|
||||
let tx = Arc::new(Mutex::new(tx));
|
||||
|
||||
// Process each Python module in parallel
|
||||
modules.par_iter().for_each(|module| {
|
||||
let result = PyModule::import(py, module.name)
|
||||
.map_err(|_| PythonError::ImportError(module.name.to_string()))
|
||||
.and_then(|m| m.call0(module.function)
|
||||
.map_err(|_| PythonError::FunctionError(module.function.to_string())));
|
||||
|
||||
if let Err(e) = result {
|
||||
let tx = tx.lock().unwrap();
|
||||
tx.send(e).unwrap();
|
||||
}
|
||||
});
|
||||
|
||||
// Check for errors
|
||||
drop(tx); // Close the sender
|
||||
for error in rx {
|
||||
match error {
|
||||
PythonError::ImportError(module) => error!("Failed to import module {}", module),
|
||||
PythonError::FunctionError(function) => error!("Failed to call function {}", function),
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@ -1,59 +0,0 @@
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyList;
|
||||
use rayon::prelude::*;
|
||||
use std::fs;
|
||||
use std::time::Instant;
|
||||
|
||||
// Define the new execute function
|
||||
fn exec_concurrently(script_path: &str, threads: usize) -> PyResult<()> {
|
||||
(0..threads).into_par_iter().for_each(|_| {
|
||||
Python::with_gil(|py| {
|
||||
let sys = py.import("sys").unwrap();
|
||||
let path: &PyList = match sys.getattr("path") {
|
||||
Ok(path) => match path.downcast() {
|
||||
Ok(path) => path,
|
||||
Err(e) => {
|
||||
eprintln!("Failed to downcast path: {:?}", e);
|
||||
return;
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
eprintln!("Failed to get path attribute: {:?}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = path.append("lib/python3.11/site-packages") {
|
||||
eprintln!("Failed to append path: {:?}", e);
|
||||
}
|
||||
|
||||
let script = fs::read_to_string(script_path).unwrap();
|
||||
py.run(&script, None, None).unwrap();
|
||||
});
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn main() -> PyResult<()> {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
let threads = 20;
|
||||
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: {} <path_to_python_script>", args[0]);
|
||||
std::process::exit(1);
|
||||
}
|
||||
let script_path = &args[1];
|
||||
|
||||
let start = Instant::now();
|
||||
|
||||
// Call the execute function
|
||||
exec_concurrently(script_path, threads)?;
|
||||
|
||||
let duration = start.elapsed();
|
||||
match fs::write("/tmp/elapsed.time", format!("booting time: {:?}", duration)) {
|
||||
Ok(_) => println!("Successfully wrote elapsed time to /tmp/elapsed.time"),
|
||||
Err(e) => eprintln!("Failed to write elapsed time: {:?}", e),
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@ -0,0 +1,42 @@
|
||||
import concurrent.futures
|
||||
from typing import Any, Callable, Dict, List
|
||||
from inspect import iscoroutinefunction
|
||||
import asyncio
|
||||
|
||||
|
||||
# Helper function to run an asynchronous function in a synchronous way
|
||||
def run_async_function_in_sync(
|
||||
func: Callable, *args, **kwargs
|
||||
) -> Any:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
coroutine = func(*args, **kwargs)
|
||||
return loop.run_until_complete(coroutine)
|
||||
|
||||
|
||||
# Main omni function for parallel execution
|
||||
def omni_parallel_function_caller(
|
||||
function_calls: List[Dict[str, Any]]
|
||||
) -> List[Any]:
|
||||
results = []
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future_to_call = {}
|
||||
for call in function_calls:
|
||||
func = call["function"]
|
||||
args = call.get("args", ())
|
||||
kwargs = call.get("kwargs", {})
|
||||
|
||||
if iscoroutinefunction(func):
|
||||
# Wrap and execute asynchronous function in a separate process
|
||||
future = executor.submit(
|
||||
run_async_function_in_sync, func, *args, **kwargs
|
||||
)
|
||||
else:
|
||||
# Directly execute synchronous function in a thread
|
||||
future = executor.submit(func, *args, **kwargs)
|
||||
|
||||
future_to_call[future] = call
|
||||
|
||||
for future in concurrent.futures.as_completed(future_to_call):
|
||||
results.append(future.result())
|
||||
return results
|
||||
@ -0,0 +1,17 @@
|
||||
import json
|
||||
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def base_model_schema_to_json(model: BaseModel):
|
||||
"""
|
||||
Converts the JSON schema of a base model to a formatted JSON string.
|
||||
|
||||
Args:
|
||||
model (BaseModel): The base model for which to generate the JSON schema.
|
||||
|
||||
Returns:
|
||||
str: The JSON schema of the base model as a formatted JSON string.
|
||||
"""
|
||||
return json.dumps(model.model_json_schema(), indent=2)
|
||||
Loading…
Reference in new issue