Skip to content

Commit e602c05

Browse files
committed
tokio: limit number of threads and set names
1 parent f0a4061 commit e602c05

File tree

1 file changed

+19
-4
lines changed

1 file changed

+19
-4
lines changed

src/lib.rs

+19-4
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,11 @@ impl ManagerServer {
7171
connect_timeout: Duration,
7272
) -> PyResult<Self> {
7373
py.allow_threads(move || {
74-
let runtime = Runtime::new()?;
74+
let runtime = tokio::runtime::Builder::new_multi_thread()
75+
.worker_threads(4)
76+
.thread_name("torchft-manager")
77+
.enable_all()
78+
.build()?;
7579
let manager = runtime
7680
.block_on(manager::Manager::new(
7781
replica_id,
@@ -127,7 +131,11 @@ impl ManagerClient {
127131
#[new]
128132
fn new(py: Python<'_>, addr: String, connect_timeout: Duration) -> PyResult<Self> {
129133
py.allow_threads(move || {
130-
let runtime = Runtime::new()?;
134+
let runtime = tokio::runtime::Builder::new_multi_thread()
135+
.worker_threads(4)
136+
.thread_name("torchft-mgrclnt")
137+
.enable_all()
138+
.build()?;
131139
let client = runtime
132140
.block_on(manager::manager_client_new(addr, connect_timeout))
133141
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
@@ -294,7 +302,10 @@ fn lighthouse_main(py: Python<'_>) -> PyResult<()> {
294302
let mut args = env::args();
295303
args.next(); // discard binary arg
296304
let opt = lighthouse::LighthouseOpt::from_iter(args);
297-
let rt = Runtime::new()?;
305+
let rt = tokio::runtime::Builder::new_multi_thread()
306+
.thread_name("torchft-lighths")
307+
.enable_all()
308+
.build()?;
298309
rt.block_on(lighthouse_main_async(opt))
299310
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
300311
Ok(())
@@ -345,7 +356,11 @@ impl LighthouseServer {
345356
let heartbeat_timeout_ms = heartbeat_timeout_ms.unwrap_or(5000);
346357

347358
py.allow_threads(move || {
348-
let rt = Runtime::new()?;
359+
let rt = tokio::runtime::Builder::new_multi_thread()
360+
.worker_threads(4)
361+
.thread_name("torchft-lighths")
362+
.enable_all()
363+
.build()?;
349364

350365
let lighthouse = rt
351366
.block_on(lighthouse::Lighthouse::new(lighthouse::LighthouseOpt {

0 commit comments

Comments
 (0)