rank_worker.cpp 15.8 KB
Newer Older
1
2
3
4
#include "rank_worker.hpp"

#include "../models/model_factory.hpp"

5
6
#include "infinicore/ops.hpp"

7
8
9
10
11
12
#include <iostream>
#include <spdlog/spdlog.h>
#include <stdexcept>

namespace infinilm::engine {

13
14
15
16
17
18
19
20
21
22
23
24
/**
 * @deprecated This function is deprecated and will be REMOVED in the next major release (v0.2.0).
 *
 * ⚠️ DEVELOPMENT POLICY:
 *   - NO new development or feature additions permitted on this interface
 *   - Only critical bug fixes (security/stability) allowed until removal
 *   - All new code MUST migrate to the polymorphic overload below
 *
 * Replacement: Use the polymorphic overload of this same function name with updated signature
 * Reason: Legacy signature lacks support for dynamic quantization modes.
 * Removal target: v0.2.0 (Q2 2026)
 */
Jiacheng Huang's avatar
Jiacheng Huang committed
25
RankWorker::RankWorker(const InfinilmModel::Config &model_config,
26
                       const distributed::RankInfo &rank_info,
27
                       const cache::CacheConfig *cache_config,
28
                       RankBarrier *barrier,
29
30
                       bool enable_graph_compiling,
                       backends::AttentionBackend attention_backend)
31
    : legacy_model_config_(model_config),
32
      rank_info_(rank_info),
33
      attention_backend_(attention_backend),
34
      enable_graph_compiling_(enable_graph_compiling),
35
36
37
38
      job_cmd_(Command::INIT),
      has_job_(false),
      job_done_(false),
      should_exit_(false),
39
      init_done_(false),
40
41
      rng_(std::random_device{}()),
      barrier_(barrier) {
PanZezhong's avatar
PanZezhong committed
42
43
44
    if (cache_config != nullptr) {
        pending_cache_config_ = cache_config->unique_copy();
    }
45
46
47
48
49
50
51
52
    // start the thread
    thread_ = std::thread(&RankWorker::thread_loop, this);

    // Wait until the worker thread finishes initialization (model created)
    std::unique_lock<std::mutex> lk(mutex_);
    cv_.wait(lk, [&] { return init_done_; });
}

53
54
55
56
57
RankWorker::RankWorker(
    std::shared_ptr<infinilm::config::ModelConfig> model_config,
    const distributed::RankInfo &rank_info,
    const cache::CacheConfig *cache_config,
    RankBarrier *barrier,
58
59
    bool enable_graph_compiling,
    backends::AttentionBackend attention_backend)
60
61
    : model_config_(model_config),
      rank_info_(rank_info),
62
      attention_backend_(attention_backend),
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
      enable_graph_compiling_(enable_graph_compiling),
      job_cmd_(Command::INIT),
      has_job_(false),
      job_done_(false),
      should_exit_(false),
      init_done_(false),
      rng_(std::random_device{}()),
      barrier_(barrier) {
    if (cache_config != nullptr) {
        pending_cache_config_ = cache_config->unique_copy();
    }
    // start the thread
    thread_ = std::thread(&RankWorker::thread_loop, this);
    // Wait until the worker thread finishes initialization (model created)
    std::unique_lock<std::mutex> lk(mutex_);
    cv_.wait(lk, [&] { return init_done_; });
}

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
std::string RankWorker::info() const {
    std::stringstream ss;

    ss << "RankWorker{";

    // Rank related
    ss << rank_info_.to_string() << " ";

    // Flags
    ss << "| init_done: " << (init_done_ ? "true" : "false") << " ";
    ss << "| should_exit: " << (should_exit_ ? "true" : "false") << " ";
    ss << "| has_job: " << (has_job_ ? "true" : "false") << " ";
    ss << "| job_done: " << (job_done_ ? "true" : "false") << " ";

    ss << "}";

    return ss.str();
}

//------------------------------------------------------
// load_param -- synchronous (blocks until worker finishes loading)
//------------------------------------------------------
void RankWorker::load_param(const std::string &name,
                            const infinicore::Tensor &param) {
    {
        std::lock_guard<std::mutex> lock(mutex_);
        // If the worker is stopping, don't submit new jobs.
        if (should_exit_) {
            throw std::runtime_error("RankWorker is closing; cannot load_param");
        }

        pending_param_name_ = name;
        pending_param_ = param;

        job_cmd_ = Command::LOAD;
        has_job_ = true;
        job_done_ = false;
    }
    cv_.notify_all();

    // Wait for job completion
    std::unique_lock<std::mutex> lk(mutex_);
    cv_.wait(lk, [&] { return job_done_ || should_exit_; });

    if (should_exit_) {
        throw std::runtime_error("RankWorker stopped while loading parameter");
    }
}

130
131
132
133
//------------------------------------------------------
// state_dict --
//------------------------------------------------------
std::unordered_map<std::string, infinicore::nn::Parameter> RankWorker::state_dict() {
PanZezhong's avatar
PanZezhong committed
134
135
136
137
138
139
140
141
    std::unique_lock<std::mutex> lk(mutex_);
    cv_.wait(lk, [&] { return init_done_ || should_exit_; });

    if (!model_) {
        throw std::runtime_error("state_dict called before model initialization");
    }

    return model_->state_dict();
142
143
}

144
//------------------------------------------------------
PanZezhong's avatar
PanZezhong committed
145
// run -- asynchronous
146
//------------------------------------------------------
147
void RankWorker::run(const Input &args) {
PanZezhong's avatar
PanZezhong committed
148
    std::lock_guard<std::mutex> lock(mutex_);
149

PanZezhong's avatar
PanZezhong committed
150
151
    if (should_exit_) {
        throw std::runtime_error("RankWorker is closing; cannot run");
152
    }
PanZezhong's avatar
PanZezhong committed
153
154
155
156
157
158

    pending_args_ = args;
    job_cmd_ = Command::RUN;
    has_job_ = true;
    job_done_ = false;

159
    cv_.notify_all();
PanZezhong's avatar
PanZezhong committed
160
}
161

162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
//------------------------------------------------------
// compile -- asynchronous
//------------------------------------------------------
void RankWorker::compile() {
    std::lock_guard<std::mutex> lock(mutex_);
    if (should_exit_) {
        throw std::runtime_error("RankWorker is closing; cannot run");
    }

    job_cmd_ = Command::COMPILE;
    has_job_ = true;
    job_done_ = false;
    cv_.notify_all();
}

PanZezhong's avatar
PanZezhong committed
177
178
179
180
//------------------------------------------------------
// wait -- asynchronous
//------------------------------------------------------
void RankWorker::wait() {
181
182
183
184
    std::unique_lock<std::mutex> lk(mutex_);
    cv_.wait(lk, [&] { return job_done_ || should_exit_; });

    if (should_exit_) {
PanZezhong's avatar
PanZezhong committed
185
        throw std::runtime_error("RankWorker stopped during run");
186
187
188
    }
}

PanZezhong's avatar
PanZezhong committed
189
void RankWorker::reset_cache(const cache::CacheConfig *new_config) {
190
191
192
    std::lock_guard<std::mutex> lock(mutex_);
    if (should_exit_) {
        throw std::runtime_error("RankWorker is closing; cannot reset_cache");
Ceng's avatar
Ceng committed
193
    }
194
195

    // Store both the position and the new config
PanZezhong's avatar
PanZezhong committed
196
197
    pending_cache_config_ = new_config->unique_copy();
    job_cmd_ = Command::RESET_CACHE;
198
199
200
    has_job_ = true;
    job_done_ = false;
    cv_.notify_all();
Ceng's avatar
Ceng committed
201
202
}

203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
//------------------------------------------------------
// close -- request shutdown and join thread
//------------------------------------------------------
void RankWorker::close() {
    {
        std::lock_guard<std::mutex> lock(mutex_);
        should_exit_ = true;
        has_job_ = false; // don't keep old jobs pending
        job_cmd_ = Command::STOP;
    }
    cv_.notify_all();

    if (thread_.joinable()) {
        thread_.join();
    }
}

//------------------------------------------------------
// get_output (thread safe)
//------------------------------------------------------
223
RankWorker::Output RankWorker::get_output() {
224
225
226
227
228
229
230
231
232
233
234
    std::lock_guard<std::mutex> lock(mutex_);
    return output_;
}

//------------------------------------------------------
// thread_loop
//------------------------------------------------------
void RankWorker::thread_loop() {
    try {
        {
            std::lock_guard<std::mutex> lk(mutex_);
PanZezhong's avatar
PanZezhong committed
235
236
237
238
239

            // Initialize device & model outside of holding the main mutex to avoid blocking callers.
            infinicore::context::setDevice(rank_info_.device);

            // Create model using factory (may be expensive)
240
            if (model_config_ == nullptr) {
241
242
243
244
245
                model_ = InfinilmModelFactory::createModel(
                    legacy_model_config_,
                    rank_info_,
                    pending_cache_config_ != nullptr ? pending_cache_config_.get() : nullptr,
                    attention_backend_);
246
247

            } else {
248
249
250
251
252
                model_ = InfinilmModelFactory::createModel(
                    model_config_,
                    rank_info_,
                    pending_cache_config_ != nullptr ? pending_cache_config_.get() : nullptr,
                    attention_backend_);
253
254
            }

PanZezhong's avatar
PanZezhong committed
255
256
257
            if (!model_) {
                throw std::runtime_error("Failed to create model");
            }
258
            if (enable_graph_compiling_) {
259
                compiler_ = std::make_unique<GeneralCompiler>(model_, barrier_);
260
261
            }

262
263
264
265
266
267
268
269
270
            init_done_ = true;
        }
        cv_.notify_all();

        // Main loop: wait for jobs or exit
        while (true) {
            Command local_cmd = Command::INIT;
            std::string local_param_name;
            infinicore::Tensor local_param;
PanZezhong's avatar
PanZezhong committed
271
            Input local_args;
PanZezhong's avatar
PanZezhong committed
272
            std::unique_ptr<cache::CacheConfig> local_cache_config;
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288

            // Wait for a job or exit
            {
                std::unique_lock<std::mutex> lk(mutex_);
                cv_.wait(lk, [&] { return has_job_ || should_exit_; });

                if (should_exit_) {
                    break;
                }

                // capture job data and clear has_job_
                local_cmd = job_cmd_;
                if (local_cmd == Command::LOAD) {
                    local_param_name = pending_param_name_;
                    local_param = pending_param_;
                } else if (local_cmd == Command::RUN) {
PanZezhong's avatar
PanZezhong committed
289
                    local_args = pending_args_;
Ceng's avatar
Ceng committed
290
                } else if (local_cmd == Command::RESET_CACHE) {
PanZezhong's avatar
PanZezhong committed
291
292
293
                    if (pending_cache_config_ != nullptr) {
                        local_cache_config = pending_cache_config_->unique_copy();
                    }
294
295
296
297
298
299
300
301
302
303
304
                }
                // mark job as being processed
                has_job_ = false;
                job_done_ = false;
            } // unlock mutex while executing the job

            // Execute job outside the lock
            if (local_cmd == Command::LOAD) {
                try {
                    model_->load_parameter(local_param_name, local_param);
                } catch (const std::exception &e) {
305
306
307
308
309
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        should_exit_ = true;
                        job_done_ = true;
                    }
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
                    cv_.notify_all();
                    spdlog::error("[{}] exception during load_parameter_: {}\n", info(), e.what());
                    break;
                }

                // signal completion
                {
                    std::lock_guard<std::mutex> lk(mutex_);
                    job_done_ = true;
                }
                cv_.notify_all();

            } else if (local_cmd == Command::RUN) {
                try {
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
326

327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
                        infinicore::Tensor logits;
                        // Try to get compiled graph
                        if (compiler_ != nullptr) {
                            auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu()));
                            if (graph != nullptr && output != nullptr) {
                                graph->run();
                                logits = output->logits;
                            }
                        }
                        // Fall back to eager mode
                        if (!logits) {
                            auto model_args = local_args.to_model_input(rank_info_.device);
                            logits = model_->forward(model_args).logits;
                        }

PanZezhong's avatar
PanZezhong committed
342
                        // Random sampling (rank 0 only)
343
                        if (rank_info_.tp_rank == 0) {
PanZezhong's avatar
PanZezhong committed
344
345
346
                            auto temperature{local_args.temperature};
                            auto top_p{local_args.top_p};
                            auto top_k{local_args.top_k};
347
348
349

                            const auto &logits_shape{logits->shape()};
                            const auto &vocab_size{logits_shape[2]};
350
351
352
                            const auto &total_len{logits_shape[1]};
                            const auto &batch_size{logits_shape[0]};

PanZezhong's avatar
PanZezhong committed
353
                            auto n_req = local_args.input_offsets.value()->size(0) - 1;
354
                            int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data();
355

356
                            auto output_ids{infinicore::Tensor::empty({n_req}, infinicore::DataType::I64, rank_info_.device)};
357

358
                            for (auto i{decltype(n_req)(0)}; i < n_req; ++i) {
PanZezhong's avatar
PanZezhong committed
359
                                auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, size_t(input_offsets[i + 1] - 1), 1}})->view({vocab_size})};
360
                                auto out{output_ids->narrow({{0, i, 1}})->view({})};
361
                                float random_val = std::uniform_real_distribution<float>(0, 1)(rng_);
362
363
364
365
366
367
368
369
370
371
372
373
374
                                infinicore::op::random_sample_(
                                    out, score, random_val, top_p, top_k, temperature);
                            }

                            output_ids = output_ids->to(infinicore::Device::cpu());

                            infinicore::context::syncStream();

                            auto out{Output{output_ids}};

                            output_ = std::move(out);
                        }

375
376
377
378
379
                        job_done_ = true;
                    }
                    cv_.notify_all();

                } catch (const std::exception &e) {
380
381
382
383
384
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        should_exit_ = true;
                        job_done_ = true;
                    }
385
386
387
388
                    cv_.notify_all();
                    spdlog::error("[{}] exception during forward: {}\n", info(), e.what());
                    break;
                }
Ceng's avatar
Ceng committed
389
390
            } else if (local_cmd == Command::RESET_CACHE) {
                try {
PanZezhong's avatar
PanZezhong committed
391
                    model_->reset_cache(local_cache_config != nullptr ? local_cache_config.get() : nullptr);
392
393
394
395
396
397
398
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        job_done_ = true;
                    }
                    cv_.notify_all();

                } catch (const std::exception &e) {
399
400
401
402
403
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        should_exit_ = true;
                        job_done_ = true;
                    }
404
405
406
407
408
409
                    cv_.notify_all();
                    spdlog::error("[{}] exception during reset_cache: {}\n", info(), e.what());
                    break;
                }
            } else if (local_cmd == Command::COMPILE) {
                try {
410
411
412
                    if (compiler_ != nullptr) {
                        compiler_->compile();
                    }
Ceng's avatar
Ceng committed
413
414
415
416
417
418
419
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        job_done_ = true;
                    }
                    cv_.notify_all();

                } catch (const std::exception &e) {
420
421
422
423
424
                    {
                        std::lock_guard<std::mutex> lk(mutex_);
                        should_exit_ = true;
                        job_done_ = true;
                    }
Ceng's avatar
Ceng committed
425
                    cv_.notify_all();
426
                    spdlog::error("[{}] exception during compile: {}\n", info(), e.what());
Ceng's avatar
Ceng committed
427
428
                    break;
                }
429

430
431
432
433
            } else {
                // Shouldn't reach here (no-op)
            }
        } // while
434
435
436

        // Some clean up should be done before exiting the thread
        compiler_.reset();
437
438
439
440
441
442
443
444
445
446
447
448
449
    } catch (const std::exception &e) {
        // Top-level exception: ensure any waiters are woken and the thread exits cleanly.
        {
            std::lock_guard<std::mutex> lk(mutex_);
            should_exit_ = true;
            job_done_ = true;
        }
        cv_.notify_all();
        spdlog::error("[{}] fatal exception in thread_loop: {} \n", info(), e.what());
    }
}

} // namespace infinilm::engine