| 139 | } |
| 140 | |
| 141 | Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { |
| 142 | static std::unordered_map<std::string, std::shared_ptr<detail::GroupImpl>> |
| 143 | backends; |
| 144 | |
| 145 | // Already initialized so return the group. |
| 146 | if (auto g = backends.find(bk); g != backends.end()) { |
| 147 | return Group(g->second); |
| 148 | } |
| 149 | |
| 150 | // Create the requested communication group |
| 151 | std::shared_ptr<detail::GroupImpl> group{nullptr}; |
| 152 | std::string bk_ = bk; |
| 153 | if (bk == "mpi") { |
| 154 | group = mpi::init(strict); |
| 155 | } else if (bk == "ring") { |
| 156 | group = ring::init(strict); |
| 157 | } else if (bk == "nccl") { |
| 158 | group = nccl::init(strict); |
| 159 | } else if (bk == "jaccl") { |
| 160 | group = jaccl::init(strict); |
| 161 | } else if (bk == "any") { |
| 162 | if (mlx::core::cu::is_available()) { |
| 163 | group = nccl::init(false); |
| 164 | bk_ = "nccl"; |
| 165 | } |
| 166 | if (group == nullptr) { |
| 167 | group = ring::init(false); |
| 168 | bk_ = "ring"; |
| 169 | } |
| 170 | if (group == nullptr) { |
| 171 | group = mpi::init(false); |
| 172 | bk_ = "mpi"; |
| 173 | } |
| 174 | if (group == nullptr) { |
| 175 | group = jaccl::init(false); |
| 176 | bk_ = "jaccl"; |
| 177 | } |
| 178 | if (group == nullptr && strict) { |
| 179 | throw std::runtime_error("[distributed] Couldn't initialize any backend"); |
| 180 | } |
| 181 | } else { |
| 182 | std::ostringstream msg; |
| 183 | msg << "[distributed] The only valid values for backend are 'any', 'mpi', 'nccl', " |
| 184 | << "'jaccl' and 'ring' but '" << bk << "' was provided."; |
| 185 | throw std::invalid_argument(msg.str()); |
| 186 | } |
| 187 | |
| 188 | if (group == nullptr) { |
| 189 | group = std::make_shared<detail::EmptyGroup>(); |
| 190 | } else { |
| 191 | backends.insert({"any", group}); |
| 192 | } |
| 193 | backends.insert({std::move(bk_), group}); |
| 194 | return Group(group); |
| 195 | } |
| 196 | |
| 197 | } // namespace mlx::core::distributed |
nothing calls this directly
no test coverage detected