diff --git a/python/src/cluster.rs b/python/src/cluster.rs index f5da9298a..aa4260ce2 100644 --- a/python/src/cluster.rs +++ b/python/src/cluster.rs @@ -143,7 +143,7 @@ pub struct PyExecutor { #[pymethods] impl PyExecutor { - #[pyo3(signature = (bind_port=None, bind_host =None, scheduler_host = None, scheduler_port = None))] + #[pyo3(signature = (bind_port=None, bind_host =None, scheduler_host = None, scheduler_port = None, concurrent_tasks = None))] #[new] pub fn new( py: Python, @@ -151,6 +151,7 @@ impl PyExecutor { bind_host: Option, scheduler_host: Option, scheduler_port: Option, + concurrent_tasks: Option, ) -> PyResult { let mut config = ExecutorProcessConfig::default(); if let Some(port) = bind_port { @@ -169,6 +170,10 @@ impl PyExecutor { config.scheduler_host = host; } + if let Some(concurrent_tasks) = concurrent_tasks { + config.concurrent_tasks = concurrent_tasks as usize + } + config.override_logical_codec = Some(Arc::new(PyLogicalCodec::try_new(py)?)); config.override_physical_codec = Some(Arc::new(PyPhysicalCodec::try_new(py)?));