← back to Exo
Try harder to clean up processes nicely (#1889)
8d81811b894d5d3387c17b1924fec9e31af62bbc · 2026-04-14 16:37:49 +0100 · rltakashige
## Motivation
Model loading is actually quite reliable now. No need to kill if you
have a slow SSD or it's a massive model; the user can shut the instance
down if necessary.
This was a major cause of signal=9 issues although not the only one (can
happen during inference too?).
The reason signal=9 is so bad is that RDMA will no longer work until
restart if this ever happens.
## Changes
- no more model load timeout
- no more crazy sigkills
- try harder to clean up processes on model shutdown
## Test Plan
### Manual Testing
Tested with some RDMA instances
Files touched
M src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/llm_inference/runner.pyM src/exo/worker/runner/runner_supervisor.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
Diff
commit 8d81811b894d5d3387c17b1924fec9e31af62bbc
Author: rltakashige <rl.takashige@gmail.com>
Date: Tue Apr 14 16:37:49 2026 +0100
Try harder to clean up processes nicely (#1889)
## Motivation
Model loading is actually quite reliable now. No need to kill if you
have a slow SSD or it's a massive model; the user can shut the instance
down if necessary.
This was a major cause of signal=9 issues although not the only one (can
happen during inference too?).
The reason signal=9 is so bad is that RDMA will no longer work until
restart if this ever happens.
## Changes
- no more model load timeout
- no more crazy sigkills
- try harder to clean up processes on model shutdown
## Test Plan
### Manual Testing
Tested with some RDMA instances
---
src/exo/worker/engines/mlx/auto_parallel.py | 85 +++-------------------
src/exo/worker/engines/mlx/utils_mlx.py | 24 +-----
src/exo/worker/runner/llm_inference/runner.py | 8 --
src/exo/worker/runner/runner_supervisor.py | 30 ++++++--
.../worker/tests/unittests/test_mlx/conftest.py | 8 +-
.../test_mlx/test_pipeline_prefill_callbacks.py | 4 +-
6 files changed, 40 insertions(+), 119 deletions(-)
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 3b6e8421..29c5320e 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -1,10 +1,8 @@
-import os
-import threading
from abc import ABC, abstractmethod
from collections.abc import Callable
from functools import partial
from inspect import signature
-from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
+from typing import TYPE_CHECKING, Literal, Protocol, cast
import mlx.core as mx
import mlx.nn as nn
@@ -63,7 +61,6 @@ from exo.worker.runner.bootstrap import logger
if TYPE_CHECKING:
from mlx_lm.models.cache import Cache
-TimeoutCallback = Callable[[], None]
LayerLoadedCallback = Callable[[int, int], None] # (layers_loaded, total_layers)
@@ -82,38 +79,6 @@ def clear_prefill_sends() -> None:
_pending_prefill_sends.clear()
-def eval_with_timeout(
- mlx_item: Any, # pyright: ignore[reportAny]
- timeout_seconds: float = 60.0,
- on_timeout: TimeoutCallback | None = None,
-) -> None:
- """Evaluate MLX item with a hard timeout.
-
- If on_timeout callback is provided, it will be called before terminating
- the process. This allows the runner to send a failure event before exit.
- """
- completed = threading.Event()
-
- def watchdog() -> None:
- if not completed.wait(timeout=timeout_seconds):
- logger.error(
- f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. "
- "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. "
- "Terminating process."
- )
- if on_timeout is not None:
- on_timeout()
- os._exit(1)
-
- watchdog_thread = threading.Thread(target=watchdog, daemon=True)
- watchdog_thread.start()
-
- try:
- mx.eval(mlx_item) # pyright: ignore[reportAny]
- finally:
- completed.set()
-
-
class _LayerCallable(Protocol):
"""Structural type that any compatible layer must satisfy.
@@ -491,8 +456,6 @@ def patch_tensor_model[T](model: T) -> T:
def tensor_auto_parallel(
model: nn.Module,
group: mx.distributed.Group,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
all_to_sharded_linear = partial(
@@ -612,9 +575,7 @@ def tensor_auto_parallel(
else:
raise ValueError(f"Unsupported model type: {type(model)}")
- model = tensor_parallel_sharding_strategy.shard_model(
- model, timeout_seconds, on_timeout, on_layer_loaded
- )
+ model = tensor_parallel_sharding_strategy.shard_model(model, on_layer_loaded)
return patch_tensor_model(model)
@@ -638,8 +599,6 @@ class TensorParallelShardingStrategy(ABC):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module: ...
@@ -648,15 +607,13 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(LlamaModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
# Force load weights before sharding to avoid FAST_SYNCH deadlock
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -704,15 +661,13 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(DeepseekV3Model, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
# Shard the self attention
if layer.self_attn.q_lora_rank is None:
@@ -789,19 +744,13 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(GLM4MoeLiteModel, model)
total = len(model.layers) # type: ignore
for i, layer in enumerate(model.layers): # type: ignore
layer = cast(Glm4MoeLiteDecoderLayer, layer)
- eval_with_timeout(
- layer.parameters(),
- timeout_seconds / total,
- on_timeout,
- )
+ mx.eval(layer.parameters())
if layer.self_attn.q_lora_rank is None: # type: ignore
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
@@ -935,14 +884,12 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(MiniMaxModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
# Shard the self attention
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -976,8 +923,6 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(
@@ -985,7 +930,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
# Shard the self attention
if isinstance(layer, Qwen3MoeDecoderLayer):
layer.self_attn.q_proj = self.all_to_sharded_linear(
@@ -1137,14 +1082,12 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(Glm4MoeModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -1185,15 +1128,13 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(GptOssMoeModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -1228,15 +1169,13 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(Step35Model, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -1273,15 +1212,13 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
model: nn.Module,
- timeout_seconds: float,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(NemotronHModel, model)
rank = self.group.rank()
total = len(model.layers)
for i, layer in enumerate(model.layers):
- eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
+ mx.eval(layer.parameters())
mixer = layer.mixer
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 63f1b6ed..1bcc6135 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -59,8 +59,6 @@ from exo.shared.types.worker.shards import (
)
from exo.worker.engines.mlx.auto_parallel import (
LayerLoadedCallback,
- TimeoutCallback,
- eval_with_timeout,
get_inner_model,
get_layers,
pipeline_auto_parallel,
@@ -84,10 +82,6 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
)
-class ModelLoadingTimeoutError(Exception):
- pass
-
-
class HostList(RootModel[list[str]]):
@classmethod
def from_hosts(cls, hosts: list[Host]) -> "HostList":
@@ -169,7 +163,6 @@ def initialize_mlx(
def load_mlx_items(
bound_instance: BoundInstance,
group: Group | None,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> "tuple[Model, TokenizerWrapper, VisionProcessor | None]":
if group is None:
@@ -201,7 +194,6 @@ def load_mlx_items(
model, tokenizer = shard_and_load(
bound_instance.bound_shard,
group=group,
- on_timeout=on_timeout,
on_layer_loaded=on_layer_loaded,
)
end_time = time.perf_counter()
@@ -230,7 +222,6 @@ def load_mlx_items(
def shard_and_load(
shard_metadata: ShardMetadata,
group: Group,
- on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> tuple[nn.Module, TokenizerWrapper]:
model_path = build_model_path(shard_metadata.model_card.model_id)
@@ -260,27 +251,16 @@ def shard_and_load(
logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
- # Estimate timeout based on model size (5x default for large queued workloads)
- base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300"))
- model_size = get_weights_size(shard_metadata)
- timeout_seconds = base_timeout + model_size.in_gb
- logger.info(
- f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
- f"(model size: {model_size.in_gb:.1f}GB)"
- )
-
match shard_metadata:
case TensorShardMetadata():
logger.info(f"loading model from {model_path} with tensor parallelism")
- model = tensor_auto_parallel(
- model, group, timeout_seconds, on_timeout, on_layer_loaded
- )
+ model = tensor_auto_parallel(model, group, on_layer_loaded)
case PipelineShardMetadata():
logger.info(f"loading model from {model_path} with pipeline parallelism")
model = pipeline_auto_parallel(
model, group, shard_metadata, on_layer_loaded=on_layer_loaded
)
- eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+ mx.eval(model.parameters())
case CfgShardMetadata():
raise ValueError(
"CfgShardMetadata is not supported for text model loading - "
diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py
index 53fc2f19..55650113 100644
--- a/src/exo/worker/runner/llm_inference/runner.py
+++ b/src/exo/worker/runner/llm_inference/runner.py
@@ -40,7 +40,6 @@ from exo.shared.types.worker.runner_response import (
from exo.shared.types.worker.runners import (
RunnerConnected,
RunnerConnecting,
- RunnerFailed,
RunnerIdle,
RunnerLoaded,
RunnerLoading,
@@ -182,12 +181,6 @@ class Runner:
)
self.acknowledge_task(task)
- def on_model_load_timeout() -> None:
- self.update_status(
- RunnerFailed(error_message="Model loading timed out")
- )
- time.sleep(0.5)
-
def on_layer_loaded(layers_loaded: int, total: int) -> None:
self.update_status(
RunnerLoading(layers_loaded=layers_loaded, total_layers=total)
@@ -203,7 +196,6 @@ class Runner:
) = load_mlx_items(
self.bound_instance,
self.generator.group,
- on_timeout=on_model_load_timeout,
on_layer_loaded=on_layer_loaded,
)
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index b55c0bd4..8a48c6bc 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -136,12 +136,30 @@ class RunnerSupervisor:
"Runner process didn't shutdown succesfully, terminating"
)
self.runner_process.terminate()
- self.runner_process.join(timeout=5)
- # This is overkill but it's not technically bad, just unnecessary.
- if self.runner_process.is_alive():
- logger.critical("Runner process didn't respond to SIGTERM, killing")
- self.runner_process.kill()
- self.runner_process.join(timeout=5)
+ self.runner_process.join(timeout=10)
+
+ if not self.runner_process.is_alive():
+ logger.warning("Terminated nicely in the first attempt!")
+
+ else:
+ # Try really hard to terminate
+ for i in range(2, 11):
+ self.runner_process.terminate()
+ self.runner_process.join(timeout=2)
+ if not self.runner_process.is_alive():
+ logger.warning(f"That took {i} attempts :)")
+ break
+ # Try even harder to kill
+ else:
+ logger.critical(
+ "Runner process didn't respond to SIGTERM, killing"
+ )
+ j = 0
+ while self.runner_process.is_alive():
+ j += 1
+ self.runner_process.kill()
+ self.runner_process.join(timeout=5)
+ logger.warning(f"That took {j} attempts :(")
else:
logger.info("Runner process succesfully terminated")
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index c09bfe1e..98bd4694 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -96,9 +96,7 @@ def run_gpt_oss_pipeline_device(
n_layers=24,
)
- model, tokenizer = shard_and_load(
- shard_meta, group, on_timeout=None, on_layer_loaded=None
- )
+ model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
model = cast(Model, model)
# Generate a prompt of exact token length
@@ -174,9 +172,7 @@ def run_gpt_oss_tensor_parallel_device(
n_layers=24,
)
- model, tokenizer = shard_and_load(
- shard_meta, group, on_timeout=None, on_layer_loaded=None
- )
+ model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
model = cast(Model, model)
base_text = "The quick brown fox jumps over the lazy dog. "
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
index be723380..161ad7d7 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
@@ -174,9 +174,7 @@ def _run_pipeline_device(
n_layers=TOTAL_LAYERS,
)
- model, tokenizer = shard_and_load(
- shard_meta, group, on_timeout=None, on_layer_loaded=None
- )
+ model, tokenizer = shard_and_load(shard_meta, group, on_layer_loaded=None)
model = cast(Any, model)
prompt, task = _build_prompt(tokenizer, prompt_tokens)
← f2709dcd Add prefix cache flag to exo bench (#1888)
·
back to Exo
·
Add gemma 4 tensor parallelism (#1891) b8eaf707 →