@@ -71,7 +71,11 @@ impl ManagerServer {
71
71
connect_timeout : Duration ,
72
72
) -> PyResult < Self > {
73
73
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 ( ) ?;
75
79
let manager = runtime
76
80
. block_on ( manager:: Manager :: new (
77
81
replica_id,
@@ -127,7 +131,11 @@ impl ManagerClient {
127
131
#[ new]
128
132
fn new ( py : Python < ' _ > , addr : String , connect_timeout : Duration ) -> PyResult < Self > {
129
133
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 ( ) ?;
131
139
let client = runtime
132
140
. block_on ( manager:: manager_client_new ( addr, connect_timeout) )
133
141
. map_err ( |e| PyRuntimeError :: new_err ( e. to_string ( ) ) ) ?;
@@ -294,7 +302,10 @@ fn lighthouse_main(py: Python<'_>) -> PyResult<()> {
294
302
let mut args = env:: args ( ) ;
295
303
args. next ( ) ; // discard binary arg
296
304
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 ( ) ?;
298
309
rt. block_on ( lighthouse_main_async ( opt) )
299
310
. map_err ( |e| PyRuntimeError :: new_err ( e. to_string ( ) ) ) ?;
300
311
Ok ( ( ) )
@@ -345,7 +356,11 @@ impl LighthouseServer {
345
356
let heartbeat_timeout_ms = heartbeat_timeout_ms. unwrap_or ( 5000 ) ;
346
357
347
358
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 ( ) ?;
349
364
350
365
let lighthouse = rt
351
366
. block_on ( lighthouse:: Lighthouse :: new ( lighthouse:: LighthouseOpt {
0 commit comments