← back to Exo
Fix race condition in mlx_distributed_init with concurrent instances (#1012)
4b65d5f89691e529a4a85bda7e57c0f6a81baa8a · 2025-12-27 11:13:26 -0500 · Heath Dutton🕴️
## Motivation
Fixes #1005
When multiple instances initialize concurrently with the same rank, they
overwrite each other's coordination files (hosts_{rank}.json), causing
"[jaccl] Malformed device file" errors and initialization failures.
## Changes
- Changed coordination filename from `./hosts_{rank}.json` to
`./hosts_{instance_id}_{rank}.json` to make it unique per instance
- Added cleanup in a finally block to remove coordination files after
initialization completes
- Applied fix to both MlxRingInstance and MlxJacclInstance cases
## Why It Works
Each instance now gets a unique coordination file based on its
instance_id, preventing concurrent instances from overwriting each
other's files. The cleanup logic ensures files are removed after use,
preventing accumulation and handling both success and failure cases.
## Test Plan
### Manual Testing
Code review and logic verification. The fix prevents the race condition
by ensuring filename uniqueness per instance.
### Automated Testing
No new tests added. Existing tests continue to pass.
---------
Co-authored-by: Ryuichi Leo Takashige <rl.takashige@gmail.com>
Files touched
M src/exo/worker/engines/mlx/utils_mlx.py
Diff
commit 4b65d5f89691e529a4a85bda7e57c0f6a81baa8a
Author: Heath Dutton🕴️ <heathdutton@gmail.com>
Date: Sat Dec 27 11:13:26 2025 -0500
Fix race condition in mlx_distributed_init with concurrent instances (#1012)
## Motivation
Fixes #1005
When multiple instances initialize concurrently with the same rank, they
overwrite each other's coordination files (hosts_{rank}.json), causing
"[jaccl] Malformed device file" errors and initialization failures.
## Changes
- Changed coordination filename from `./hosts_{rank}.json` to
`./hosts_{instance_id}_{rank}.json` to make it unique per instance
- Added cleanup in a finally block to remove coordination files after
initialization completes
- Applied fix to both MlxRingInstance and MlxJacclInstance cases
## Why It Works
Each instance now gets a unique coordination file based on its
instance_id, preventing concurrent instances from overwriting each
other's files. The cleanup logic ensures files are removed after use,
preventing accumulation and handling both success and failure cases.
## Test Plan
### Manual Testing
Code review and logic verification. The fix prevents the race condition
by ensuring filename uniqueness per instance.
### Automated Testing
No new tests added. Existing tests continue to pass.
---------
Co-authored-by: Ryuichi Leo Takashige <rl.takashige@gmail.com>
---
src/exo/worker/engines/mlx/utils_mlx.py | 90 +++++++++++++++++++--------------
1 file changed, 52 insertions(+), 38 deletions(-)
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 71114282..d7f410ce 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -21,6 +21,8 @@ try:
from mlx_lm.tokenizer_utils import load_tokenizer
except ImportError:
from mlx_lm.tokenizer_utils import load as load_tokenizer # type: ignore
+import contextlib
+
import mlx.core as mx
import mlx.nn as nn
from mlx_lm.utils import load_model
@@ -112,44 +114,56 @@ def mlx_distributed_init(
rank = bound_instance.bound_shard.device_rank
logger.info(f"Starting initialization for rank {rank}")
- # TODO: singleton instances
- match bound_instance.instance:
- case MlxRingInstance(hosts=hosts):
- hostfile = f"./hosts_{rank}.json"
- hosts_json = HostList.from_hosts(hosts).model_dump_json()
-
- with open(hostfile, "w") as f:
- _ = f.write(hosts_json)
-
- logger.info(f"rank {rank} hostfile: {hostfile} hosts: {hosts_json}")
-
- os.environ["MLX_HOSTFILE"] = hostfile
- os.environ["MLX_RANK"] = str(rank)
- os.environ["MLX_RING_VERBOSE"] = "1"
- group = mx.distributed.init(backend="ring", strict=True)
-
- case MlxJacclInstance(
- ibv_devices=ibv_devices, jaccl_coordinators=jaccl_coordinators
- ):
- # Use RDMA connectivity matrix
- devices_file = f"./hosts_{rank}.json"
- ibv_devices_json = json.dumps(ibv_devices)
-
- with open(devices_file, "w") as f:
- _ = f.write(ibv_devices_json)
-
- jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
-
- logger.info(f"rank {rank} MLX_IBV_DEVICES: {ibv_devices_json}")
- logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
- os.environ["MLX_IBV_DEVICES"] = devices_file
- os.environ["MLX_RANK"] = str(rank)
- os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
- group = mx.distributed.init(backend="jaccl", strict=True)
-
- logger.info(f"Rank {rank} mlx distributed initialization complete")
-
- return group
+ coordination_file = None
+ try:
+ # TODO: singleton instances
+ match bound_instance.instance:
+ case MlxRingInstance(hosts=hosts):
+ coordination_file = (
+ f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+ )
+ hosts_json = HostList.from_hosts(hosts).model_dump_json()
+
+ with open(coordination_file, "w") as f:
+ _ = f.write(hosts_json)
+
+ logger.info(
+ f"rank {rank} hostfile: {coordination_file} hosts: {hosts_json}"
+ )
+
+ os.environ["MLX_HOSTFILE"] = coordination_file
+ os.environ["MLX_RANK"] = str(rank)
+ os.environ["MLX_RING_VERBOSE"] = "1"
+ group = mx.distributed.init(backend="ring", strict=True)
+
+ case MlxJacclInstance(
+ ibv_devices=ibv_devices, jaccl_coordinators=jaccl_coordinators
+ ):
+ # Use RDMA connectivity matrix
+ coordination_file = (
+ f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
+ )
+ ibv_devices_json = json.dumps(ibv_devices)
+
+ with open(coordination_file, "w") as f:
+ _ = f.write(ibv_devices_json)
+
+ jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
+
+ logger.info(f"rank {rank} MLX_IBV_DEVICES: {ibv_devices_json}")
+ logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
+ os.environ["MLX_IBV_DEVICES"] = coordination_file
+ os.environ["MLX_RANK"] = str(rank)
+ os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
+ group = mx.distributed.init(backend="jaccl", strict=True)
+
+ logger.info(f"Rank {rank} mlx distributed initialization complete")
+
+ return group
+ finally:
+ with contextlib.suppress(FileNotFoundError):
+ if coordination_file:
+ os.remove(coordination_file)
def initialize_mlx(
← 1c1792f5 mlx: update to 0.30.1 and align coordinator naming with MLX
·
back to Exo
·
Separate out the Runner's behaviour into a "connect" phase a 8e9332d6 →