From 5fba80f27f1a920d7873e71075e5a66899300822 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Thu, 24 Sep 2026 09:27:43 -0700 Subject: [PATCH 1/2] fix register_group caching --- mlx/distributed/distributed.cpp | 7 +++---- python/tests/ring_test_distributed.py | 6 ++++++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/mlx/distributed/distributed.cpp b/mlx/distributed/distributed.cpp index 34be2ec600..d8fca8b160 100644 --- a/mlx/distributed/distributed.cpp +++ b/mlx/distributed/distributed.cpp @@ -148,12 +148,11 @@ get_backends() { } Group register_group(std::shared_ptr group, std::string bk) { - auto& backends = get_backends(); if (group == nullptr) { - group = std::make_shared(); - } else { - backends.insert({"any", group}); + return Group(std::make_shared()); } + auto& backends = get_backends(); + backends.insert({"any", group}); backends.insert({std::move(bk), group}); return Group(group); } diff --git a/python/tests/ring_test_distributed.py b/python/tests/ring_test_distributed.py index 31db2ed5f7..3053e32ac1 100644 --- a/python/tests/ring_test_distributed.py +++ b/python/tests/ring_test_distributed.py @@ -23,6 +23,12 @@ def test_groups(self): with self.assertRaises(RuntimeError): sub = world.split(world.rank() % 2) + def test_strict_init_after_fallback(self): + # A failed non-strict init must not satisfy a later strict init + self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) + with self.assertRaises(RuntimeError): + mx.distributed.init(backend="jaccl", strict=True) + def test_all_reduce_extra(self): world = mx.distributed.init() dtypes = [ From 1cf598419c9c6b060208672c5a73823680ce9f03 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Thu, 24 Sep 2026 10:27:24 -0700 Subject: [PATCH 2/2] clear JACCL environment --- python/tests/ring_test_distributed.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/tests/ring_test_distributed.py b/python/tests/ring_test_distributed.py index 3053e32ac1..8b4de9ccae 100644 --- a/python/tests/ring_test_distributed.py +++ b/python/tests/ring_test_distributed.py @@ -24,10 +24,14 @@ def test_groups(self): sub = world.split(world.rank() % 2) def test_strict_init_after_fallback(self): - # A failed non-strict init must not satisfy a later strict init - self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) - with self.assertRaises(RuntimeError): - mx.distributed.init(backend="jaccl", strict=True) + with mlx_tests.scoped_env( + JACCL_IBV_DEVICES=None, + MLX_IBV_DEVICES=None, + ): + self.assertEqual(mx.distributed.init(backend="jaccl").size(), 1) + + with self.assertRaises(RuntimeError): + mx.distributed.init(backend="jaccl", strict=True) def test_all_reduce_extra(self): world = mx.distributed.init()