• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

llnl / dftracer-utils / 30002033718

23 Jul 2026 11:08AM UTC coverage: 52.742% (+0.08%) from 52.66%
30002033718

Pull #99

github

web-flow
Merge df0b9588d into 06bc84ec9
Pull Request #99: Support CM time_metric (NS/MS/SEC/US) across reader and viz

41941 of 102829 branches covered (40.79%)

Branch coverage included in aggregate %.

2221 of 2839 new or added lines in 58 files covered. (78.23%)

137 existing lines in 13 files now uncovered.

37042 of 46924 relevant lines covered (78.94%)

59523.07 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.44
/src/dftracer/utils/python/sst_distribution.cpp
1
#include <dftracer/utils/core/common/constants.h>
2
#include <dftracer/utils/core/runtime.h>
3
#include <dftracer/utils/core/tasks/coro_scope.h>
4
#include <dftracer/utils/python/py_errors.h>
5
#include <dftracer/utils/python/py_runtime_mixin.h>
6
#include <dftracer/utils/python/py_type_helpers.h>
7
#include <dftracer/utils/python/runtime.h>
8
#include <dftracer/utils/python/sst_distribution.h>
9
#include <dftracer/utils/utilities/composites/dft/aggregators/aggregation_config.h>
10
#include <dftracer/utils/utilities/composites/dft/aggregators/aggregation_key.h>
11
#include <dftracer/utils/utilities/composites/dft/aggregators/aggregation_visitor.h>
12
#include <dftracer/utils/utilities/composites/dft/aggregators/association_tracker.h>
13
#include <dftracer/utils/utilities/filesystem/pattern_directory_scanner_utility.h>
14
#include <dftracer/utils/utilities/indexer/file_partition.h>
15
#include <dftracer/utils/utilities/indexer/index_batch_sink.h>
16
#include <dftracer/utils/utilities/indexer/index_builder_utility.h>
17
#include <dftracer/utils/utilities/indexer/index_database_sst_writer_context.h>
18
#include <dftracer/utils/utilities/indexer/internal/common/gzip_member_scanner.h>
19
#include <fcntl.h>
20
#include <sys/stat.h>
21
#include <unistd.h>
22

23
#include <algorithm>
24
#include <atomic>
25
#include <cstdint>
26
#include <memory>
27
#include <mutex>
28
#include <new>
29
#include <optional>
30
#include <string>
31
#include <unordered_set>
32
#include <vector>
33

34
using dftracer::utils::Runtime;
35
using dftracer::utils::utilities::filesystem::FileEntry;
36
using dftracer::utils::utilities::filesystem::PatternDirectoryScannerUtility;
37
using dftracer::utils::utilities::filesystem::
38
    PatternDirectoryScannerUtilityInput;
39
using dftracer::utils::utilities::indexer::IndexBatchBuilderUtility;
40
using dftracer::utils::utilities::indexer::IndexBatchSink;
41
using dftracer::utils::utilities::indexer::IndexBuildBatchConfig;
42
using dftracer::utils::utilities::indexer::IndexBuildBatchResult;
43
using dftracer::utils::utilities::indexer::IndexDatabaseSstWriterContext;
44
using dftracer::utils::utilities::indexer::plan_lpt_partition;
45
using dftracer::utils::utilities::indexer::SstArtifactRegistry;
46
using dftracer::utils::utilities::indexer::internal::
47
    enumerate_gzip_member_candidates;
48
using dftracer::utils::utilities::indexer::internal::GzipMember;
49

50
// ---------------------------------------------------------------------------
51
// SstArtifactRegistry type
52
// ---------------------------------------------------------------------------
53

54
typedef struct {
55
    PyObject_HEAD std::shared_ptr<SstArtifactRegistry> registry;
56
} SstArtifactRegistryObject;
57

58
static void SstArtifactRegistry_dealloc(SstArtifactRegistryObject *self) {
×
59
    self->registry.~shared_ptr<SstArtifactRegistry>();
×
60
    Py_TYPE(self)->tp_free((PyObject *)self);
×
61
}
×
62

63
static PyObject *SstArtifactRegistry_new(PyTypeObject *type,
×
64
                                         PyObject * /*args*/,
65
                                         PyObject * /*kwds*/) {
66
    auto *self = (SstArtifactRegistryObject *)type->tp_alloc(type, 0);
×
67
    if (!self) return NULL;
×
68
    new (&self->registry) std::shared_ptr<SstArtifactRegistry>(
×
69
        std::make_shared<SstArtifactRegistry>());
×
70
    return (PyObject *)self;
×
71
}
72

73
namespace {
74

75
// Field names in the Artifacts dict returned by build_sst_batch and
76
// consumed by SstArtifactRegistry.append. Must match the field names on
77
// IndexDatabaseSstWriterContext::Artifacts.
78
constexpr const char *ARTIFACT_FIELDS[] = {
79
    "metadata_sst",        "checkpoints_sst",         "manifest_sst",
80
    "chunk_bloom_sst",     "file_bloom_sst",          "chunk_stats_sst",
81
    "chunk_dim_stats_sst", "dimensions_sst",          "file_scalar_stats_sst",
82
    "file_cat_counts_sst", "file_pid_tid_counts_sst", "file_name_counts_sst",
83
    "name_dictionary_sst", "name_file_postings_sst",  "name_chunk_postings_sst",
84
    "hash_tables_sst",     "aggregation_sst",         "system_metrics_sst",
85
};
86

87
/// Map a slot name to the matching Artifacts member. Kept in one place so
88
/// that adding a new CF requires updating only `ARTIFACT_FIELDS` plus
89
/// `dispatch_*` below.
90
std::optional<std::string> *artifacts_slot(
×
91
    IndexDatabaseSstWriterContext::Artifacts &a, std::string_view name) {
92
    if (name == "metadata_sst") return &a.metadata_sst;
×
93
    if (name == "checkpoints_sst") return &a.checkpoints_sst;
×
94
    if (name == "manifest_sst") return &a.manifest_sst;
×
95
    if (name == "chunk_bloom_sst") return &a.chunk_bloom_sst;
×
96
    if (name == "file_bloom_sst") return &a.file_bloom_sst;
×
97
    if (name == "chunk_stats_sst") return &a.chunk_stats_sst;
×
98
    if (name == "chunk_dim_stats_sst") return &a.chunk_dim_stats_sst;
×
99
    if (name == "dimensions_sst") return &a.dimensions_sst;
×
100
    if (name == "file_scalar_stats_sst") return &a.file_scalar_stats_sst;
×
101
    if (name == "file_cat_counts_sst") return &a.file_cat_counts_sst;
×
102
    if (name == "file_pid_tid_counts_sst") return &a.file_pid_tid_counts_sst;
×
103
    if (name == "file_name_counts_sst") return &a.file_name_counts_sst;
×
104
    if (name == "name_dictionary_sst") return &a.name_dictionary_sst;
×
105
    if (name == "name_file_postings_sst") return &a.name_file_postings_sst;
×
106
    if (name == "name_chunk_postings_sst") return &a.name_chunk_postings_sst;
×
107
    if (name == "hash_tables_sst") return &a.hash_tables_sst;
×
108
    if (name == "aggregation_sst") return &a.aggregation_sst;
×
109
    if (name == "system_metrics_sst") return &a.system_metrics_sst;
×
110
    return nullptr;
×
111
}
112

113
/// Convert a Python artifacts dict to the C++ Artifacts struct. Missing,
114
/// None, or empty-string entries become nullopt. Returns false on type
115
/// errors (exception set).
116
bool artifacts_from_dict(PyObject *dict,
×
117
                         IndexDatabaseSstWriterContext::Artifacts *out) {
118
    if (!PyDict_Check(dict)) {
×
119
        PyErr_SetString(PyExc_TypeError, "artifacts must be a dict");
×
120
        return false;
×
121
    }
122
    for (const char *field : ARTIFACT_FIELDS) {
×
123
        PyObject *val = PyDict_GetItemString(dict, field);  // borrowed
×
124
        if (!val || val == Py_None) continue;
×
125
        if (!PyUnicode_Check(val)) {
×
126
            PyErr_Format(PyExc_TypeError, "artifacts['%s'] must be str or None",
×
127
                         field);
128
            return false;
×
129
        }
130
        const char *s = PyUnicode_AsUTF8(val);
×
131
        if (!s) return false;
×
132
        if (s[0] == '\0') continue;
×
133
        auto *slot = artifacts_slot(*out, field);
×
134
        if (slot) *slot = std::string(s);
×
135
    }
136
    return true;
×
137
}
138

139
PyObject *artifacts_to_dict(const IndexDatabaseSstWriterContext::Artifacts &a) {
×
140
    PyObject *dict = PyDict_New();
×
141
    if (!dict) return NULL;
×
142
    auto set_field = [&](const char *name,
×
143
                         const std::optional<std::string> &slot) -> bool {
144
        PyObject *v = slot.has_value() ? PyUnicode_FromString(slot->c_str())
×
145
                                       : (Py_INCREF(Py_None), Py_None);
×
146
        if (!v) return false;
×
147
        int rc = PyDict_SetItemString(dict, name, v);
×
148
        Py_DECREF(v);
149
        return rc == 0;
×
150
    };
×
151
    if (!set_field("metadata_sst", a.metadata_sst) ||
×
152
        !set_field("checkpoints_sst", a.checkpoints_sst) ||
×
153
        !set_field("manifest_sst", a.manifest_sst) ||
×
154
        !set_field("chunk_bloom_sst", a.chunk_bloom_sst) ||
×
155
        !set_field("file_bloom_sst", a.file_bloom_sst) ||
×
156
        !set_field("chunk_stats_sst", a.chunk_stats_sst) ||
×
157
        !set_field("chunk_dim_stats_sst", a.chunk_dim_stats_sst) ||
×
158
        !set_field("dimensions_sst", a.dimensions_sst) ||
×
159
        !set_field("file_scalar_stats_sst", a.file_scalar_stats_sst) ||
×
160
        !set_field("file_cat_counts_sst", a.file_cat_counts_sst) ||
×
161
        !set_field("file_pid_tid_counts_sst", a.file_pid_tid_counts_sst) ||
×
162
        !set_field("file_name_counts_sst", a.file_name_counts_sst) ||
×
163
        !set_field("name_dictionary_sst", a.name_dictionary_sst) ||
×
164
        !set_field("name_file_postings_sst", a.name_file_postings_sst) ||
×
165
        !set_field("name_chunk_postings_sst", a.name_chunk_postings_sst) ||
×
166
        !set_field("hash_tables_sst", a.hash_tables_sst) ||
×
167
        !set_field("aggregation_sst", a.aggregation_sst) ||
×
168
        !set_field("system_metrics_sst", a.system_metrics_sst)) {
×
169
        Py_DECREF(dict);
×
170
        return NULL;
×
171
    }
172
    return dict;
×
173
}
174

175
}  // namespace
176

177
static PyObject *SstArtifactRegistry_append(SstArtifactRegistryObject *self,
×
178
                                            PyObject *args) {
179
    PyObject *dict;
180
    if (!PyArg_ParseTuple(args, "O", &dict)) return NULL;
×
181
    IndexDatabaseSstWriterContext::Artifacts a;
×
182
    if (!artifacts_from_dict(dict, &a)) return NULL;
×
183
    self->registry->append(std::move(a));
×
184
    Py_RETURN_NONE;
×
185
}
×
186

187
static PyMethodDef SstArtifactRegistry_methods[] = {
188
    {"append", (PyCFunction)SstArtifactRegistry_append, METH_VARARGS,
189
     "append(artifacts_dict) -> None\n"
190
     "Add a per-batch Artifacts dict (as returned by build_sst_batch or "
191
     "IndexDatabaseSstWriterContext.commit) to the registry."},
192
    {NULL}};
193

194
static PyTypeObject SstArtifactRegistryType = {
195
    PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.SstArtifactRegistry",
196
    sizeof(SstArtifactRegistryObject),
197
    0,
198
    (destructor)SstArtifactRegistry_dealloc,
199
    0,
200
    0,
201
    0,
202
    0,
203
    0,
204
    0,
205
    0,
206
    0,
207
    0,
208
    0,
209
    0,
210
    0,
211
    0,
212
    0,
213
    Py_TPFLAGS_DEFAULT,
214
    "Thread-safe collector for SST artifact paths.",
215
    0,
216
    0,
217
    0,
218
    0,
219
    0,
220
    0,
221
    SstArtifactRegistry_methods,
222
    0,
223
    0,
224
    0,
225
    0,
226
    0,
227
    0,
228
    0,
229
    0,
230
    0,
231
    SstArtifactRegistry_new,
232
};
233

234
SstArtifactRegistry *sst_artifact_registry_get(PyObject *obj) {
×
235
    if (!PyObject_TypeCheck(obj, &SstArtifactRegistryType)) return nullptr;
×
236
    return ((SstArtifactRegistryObject *)obj)->registry.get();
×
237
}
238

239
// ---------------------------------------------------------------------------
240
// scan_files: parallel directory scan with size info
241
// ---------------------------------------------------------------------------
242

243
static PyObject *scan_files_fn(PyObject * /*self*/, PyObject *args,
×
244
                               PyObject *kwds) {
245
    static const char *kwlist[] = {"directory", "patterns", "recursive",
246
                                   "runtime", NULL};
247
    const char *directory;
248
    PyObject *patterns_obj = NULL;
×
249
    int recursive = 0;
×
250
    PyObject *runtime_arg = NULL;
×
251
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|OpO", (char **)kwlist,
×
252
                                     &directory, &patterns_obj, &recursive,
253
                                     &runtime_arg)) {
254
        return NULL;
×
255
    }
256

257
    std::vector<std::string> patterns;
×
258
    if (patterns_obj && patterns_obj != Py_None) {
×
259
        PyObject *seq =
260
            PySequence_Fast(patterns_obj, "patterns must be a sequence");
×
261
        if (!seq) return NULL;
×
262
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
263
        patterns.reserve(n);
×
264
        for (Py_ssize_t i = 0; i < n; ++i) {
×
265
            const char *s = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(seq, i));
×
266
            if (!s) {
×
267
                Py_DECREF(seq);
×
268
                return NULL;
×
269
            }
270
            patterns.emplace_back(s);
×
271
        }
272
        Py_DECREF(seq);
×
273
    }
274

275
    Runtime *rt = nullptr;
×
276
    if (runtime_arg && runtime_arg != Py_None) {
×
277
        if (!PyObject_TypeCheck(runtime_arg, &RuntimeType)) {
×
278
            PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");
×
279
            if (!native || !PyObject_TypeCheck(native, &RuntimeType)) {
×
280
                Py_XDECREF(native);
×
281
                PyErr_SetString(PyExc_TypeError,
×
282
                                "runtime must be a Runtime instance or None");
283
                return NULL;
×
284
            }
285
            rt = ((RuntimeObject *)native)->runtime.get();
×
286
            Py_DECREF(native);
×
287
        } else {
288
            rt = ((RuntimeObject *)runtime_arg)->runtime.get();
×
289
        }
290
    } else {
×
291
        rt = get_default_runtime();
×
292
    }
293

294
    PatternDirectoryScannerUtilityInput input(directory, patterns,
×
295
                                              recursive != 0, true);
×
296
    std::vector<FileEntry> entries;
×
297
    if (!run_blocking([&] {
×
298
            rt->submit(dftracer::utils::run_coro_scope(
×
299
                           rt->executor(),
300
                           [](dftracer::utils::CoroScope &scope,
×
301
                              PatternDirectoryScannerUtilityInput in,
302
                              std::vector<FileEntry> *out)
303
                               -> dftracer::utils::coro::CoroTask<void> {
×
304
                               PatternDirectoryScannerUtility scanner;
×
305
                               *out = co_await scope.spawn(scanner, in);
×
306
                           },
×
307
                           std::move(input), &entries),
×
308
                       "scan-files")
×
309
                .get();
×
310
        })) {
×
311
        return NULL;
×
312
    }
313

314
    PyObject *out = PyList_New(static_cast<Py_ssize_t>(entries.size()));
×
315
    if (!out) return NULL;
×
316
    for (std::size_t i = 0; i < entries.size(); ++i) {
×
317
        PyObject *t = Py_BuildValue("(sn)", entries[i].path.c_str(),
×
318
                                    (Py_ssize_t)entries[i].size);
×
319
        if (!t) {
×
320
            Py_DECREF(out);
×
321
            return NULL;
×
322
        }
323
        PyList_SET_ITEM(out, i, t);
×
324
    }
325
    return out;
×
326
}
×
327

328
// ---------------------------------------------------------------------------
329
// plan_lpt_partition: LPT bin-packing of (path, size) pairs
330
// ---------------------------------------------------------------------------
331

332
static PyObject *plan_lpt_partition_fn(PyObject * /*self*/, PyObject *args) {
×
333
    PyObject *entries_obj;
334
    Py_ssize_t num_workers;
335
    if (!PyArg_ParseTuple(args, "On", &entries_obj, &num_workers)) return NULL;
×
336
    if (num_workers <= 0) num_workers = 1;
×
337

338
    std::vector<FileEntry> entries;
×
339
    PyObject *seq = PySequence_Fast(entries_obj,
×
340
                                    "entries must be a sequence of "
341
                                    "(path, size) tuples");
342
    if (!seq) return NULL;
×
343
    Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
344
    entries.reserve(n);
×
345
    for (Py_ssize_t i = 0; i < n; ++i) {
×
346
        PyObject *item = PySequence_Fast_GET_ITEM(seq, i);
×
347
        const char *path = nullptr;
×
348
        Py_ssize_t size = 0;
×
349
        if (!PyArg_ParseTuple(item, "sn", &path, &size)) {
×
350
            Py_DECREF(seq);
×
351
            return NULL;
×
352
        }
353
        FileEntry fe;
×
354
        fe.path = path;
×
355
        fe.size = static_cast<std::size_t>(size);
×
356
        fe.is_regular_file = true;
×
357
        entries.push_back(std::move(fe));
×
358
    }
×
359
    Py_DECREF(seq);
×
360

361
    auto buckets = plan_lpt_partition(std::move(entries),
×
362
                                      static_cast<std::size_t>(num_workers));
×
363

364
    PyObject *out = PyList_New(static_cast<Py_ssize_t>(buckets.size()));
×
365
    if (!out) return NULL;
×
366
    for (std::size_t i = 0; i < buckets.size(); ++i) {
×
367
        PyObject *lst = PyList_New(static_cast<Py_ssize_t>(buckets[i].size()));
×
368
        if (!lst) {
×
369
            Py_DECREF(out);
×
370
            return NULL;
×
371
        }
372
        for (std::size_t j = 0; j < buckets[i].size(); ++j) {
×
373
            PyObject *t = Py_BuildValue("(sn)", buckets[i][j].path.c_str(),
×
374
                                        (Py_ssize_t)buckets[i][j].size);
×
375
            if (!t) {
×
376
                Py_DECREF(lst);
×
377
                Py_DECREF(out);
×
378
                return NULL;
×
379
            }
380
            PyList_SET_ITEM(lst, j, t);
×
381
        }
382
        PyList_SET_ITEM(out, i, lst);
×
383
    }
384
    return out;
×
385
}
×
386

387
// ---------------------------------------------------------------------------
388
// build_sst_batch: run the indexer pipeline with an SST sink and return
389
// the merged Artifacts dict.
390
// ---------------------------------------------------------------------------
391

392
static PyObject *build_sst_batch_fn(PyObject * /*self*/, PyObject *args,
×
393
                                    PyObject *kwds) {
394
    static const char *kwlist[] = {"files",
395
                                   "file_ids",
396
                                   "staging_dir",
397
                                   "batch_id",
398
                                   "index_dir",
399
                                   "checkpoint_size",
400
                                   "build_manifest",
401
                                   "force_rebuild",
402
                                   "build_bloom",
403
                                   "bloom_dimensions",
404
                                   "parallelism",
405
                                   "flush_every_files",
406
                                   "runtime",
407
                                   "aggregation_config",
408
                                   "file_slices",
409
                                   "progress",
410
                                   NULL};
411
    PyObject *files_obj;
412
    PyObject *file_ids_obj;
413
    const char *staging_dir;
414
    const char *batch_id;
415
    const char *index_dir = "";
×
416
    Py_ssize_t checkpoint_size = static_cast<Py_ssize_t>(
×
417
        dftracer::utils::constants::indexer::DEFAULT_CHECKPOINT_SIZE);
418
    int build_manifest = 0;
×
419
    int force_rebuild = 0;
×
NEW
420
    int build_bloom = 1;
×
421
    PyObject *bloom_dims_obj = NULL;
×
422
    Py_ssize_t parallelism = 0;
×
423
    Py_ssize_t flush_every_files = 0;
×
424
    PyObject *runtime_arg = NULL;
×
425
    PyObject *aggregation_config_obj = NULL;
×
426
    PyObject *file_slices_obj = NULL;
×
NEW
427
    PyObject *progress_obj = NULL;
×
428

429
    if (!PyArg_ParseTupleAndKeywords(
×
430
            args, kwds, "OOss|snpppOnnOOOO", (char **)kwlist, &files_obj,
431
            &file_ids_obj, &staging_dir, &batch_id, &index_dir,
432
            &checkpoint_size, &build_manifest, &force_rebuild, &build_bloom,
433
            &bloom_dims_obj, &parallelism, &flush_every_files, &runtime_arg,
434
            &aggregation_config_obj, &file_slices_obj, &progress_obj)) {
UNCOV
435
        return NULL;
×
436
    }
437

438
    // Unpack files.
439
    std::vector<std::string> files;
×
440
    {
441
        PyObject *seq = PySequence_Fast(files_obj, "files must be a sequence");
×
442
        if (!seq) return NULL;
×
443
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
444
        files.reserve(n);
×
445
        for (Py_ssize_t i = 0; i < n; ++i) {
×
446
            const char *s = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(seq, i));
×
447
            if (!s) {
×
448
                Py_DECREF(seq);
×
449
                return NULL;
×
450
            }
451
            files.emplace_back(s);
×
452
        }
453
        Py_DECREF(seq);
×
454
    }
455
    if (files.empty()) {
×
456
        return PyDict_New();
×
457
    }
458

459
    // Unpack file_ids, parallel to files.
460
    std::vector<int> file_ids;
×
461
    {
462
        PyObject *seq =
463
            PySequence_Fast(file_ids_obj, "file_ids must be a sequence");
×
464
        if (!seq) return NULL;
×
465
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
466
        if (static_cast<std::size_t>(n) != files.size()) {
×
467
            Py_DECREF(seq);
×
468
            PyErr_SetString(PyExc_ValueError,
×
469
                            "file_ids must have the same length as files");
470
            return NULL;
×
471
        }
472
        file_ids.reserve(n);
×
473
        for (Py_ssize_t i = 0; i < n; ++i) {
×
474
            long v = PyLong_AsLong(PySequence_Fast_GET_ITEM(seq, i));
×
475
            if (v == -1 && PyErr_Occurred()) {
×
476
                Py_DECREF(seq);
×
477
                return NULL;
×
478
            }
479
            file_ids.push_back(static_cast<int>(v));
×
480
        }
481
        Py_DECREF(seq);
×
482
    }
483

484
    // Optional bloom dimensions override.
485
    std::vector<std::string> bloom_dims;
×
486
    if (bloom_dims_obj && bloom_dims_obj != Py_None) {
×
487
        PyObject *seq = PySequence_Fast(bloom_dims_obj,
×
488
                                        "bloom_dimensions must be a sequence");
489
        if (!seq) return NULL;
×
490
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
491
        bloom_dims.reserve(n);
×
492
        for (Py_ssize_t i = 0; i < n; ++i) {
×
493
            const char *s = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(seq, i));
×
494
            if (!s) {
×
495
                Py_DECREF(seq);
×
496
                return NULL;
×
497
            }
498
            bloom_dims.emplace_back(s);
×
499
        }
500
        Py_DECREF(seq);
×
501
    }
502

503
    // Resolve Runtime (matching CheckpointIndexer pattern).
504
    Runtime *rt = nullptr;
×
505
    if (runtime_arg && runtime_arg != Py_None) {
×
506
        if (PyObject_TypeCheck(runtime_arg, &RuntimeType)) {
×
507
            rt = ((RuntimeObject *)runtime_arg)->runtime.get();
×
508
        } else {
509
            PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");
×
510
            if (!native || !PyObject_TypeCheck(native, &RuntimeType)) {
×
511
                Py_XDECREF(native);
×
512
                PyErr_SetString(PyExc_TypeError,
×
513
                                "runtime must be a Runtime instance or None");
514
                return NULL;
×
515
            }
516
            rt = ((RuntimeObject *)native)->runtime.get();
×
517
            Py_DECREF(native);
×
518
        }
519
    } else {
×
520
        rt = get_default_runtime();
×
521
    }
522

523
    // Build config + sink factory shared state.
524
    struct SharedArtifacts {
525
        std::mutex mu;
526
        std::vector<IndexDatabaseSstWriterContext::Artifacts> list;
527
    };
528
    auto artifacts = std::make_shared<SharedArtifacts>();
×
529
    auto staging = std::string(staging_dir);
×
530
    auto batch = std::string(batch_id);
×
531

532
    // Optional aggregation config, extracted from the Python dataclass.
533
    std::shared_ptr<dftracer::utils::utilities::composites::dft::aggregators::
534
                        AggregationConfig>
535
        agg_config_ptr;
×
536
    if (aggregation_config_obj && aggregation_config_obj != Py_None) {
×
537
        using dftracer::utils::utilities::composites::dft::aggregators::
538
            AggregationConfig;
539
        auto cfg = std::make_shared<AggregationConfig>();
×
540
        auto pull_double = [&](const char *name, double fallback) -> double {
×
541
            PyObject *v = PyObject_GetAttrString(aggregation_config_obj, name);
×
542
            if (!v || v == Py_None) {
×
543
                Py_XDECREF(v);
×
544
                PyErr_Clear();
×
545
                return fallback;
×
546
            }
547
            double out = PyFloat_AsDouble(v);
×
548
            Py_DECREF(v);
549
            if (out == -1.0 && PyErr_Occurred()) return fallback;
×
550
            return out;
×
551
        };
×
552
        auto pull_bool = [&](const char *name, bool fallback) -> bool {
×
553
            PyObject *v = PyObject_GetAttrString(aggregation_config_obj, name);
×
554
            if (!v || v == Py_None) {
×
555
                Py_XDECREF(v);
×
556
                PyErr_Clear();
×
557
                return fallback;
×
558
            }
559
            int out = PyObject_IsTrue(v);
×
560
            Py_DECREF(v);
561
            return out > 0 ? true : fallback;
×
562
        };
×
563
        auto pull_string_list =
564
            [&](const char *name) -> std::vector<std::string> {
×
565
            std::vector<std::string> out;
×
566
            PyObject *v = PyObject_GetAttrString(aggregation_config_obj, name);
×
567
            if (!v || v == Py_None) {
×
568
                Py_XDECREF(v);
×
569
                PyErr_Clear();
×
570
                return out;
×
571
            }
572
            PyObject *seq = PySequence_Fast(v, "expected list of str");
×
573
            Py_DECREF(v);
×
574
            if (!seq) {
×
575
                PyErr_Clear();
×
576
                return out;
×
577
            }
578
            Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
579
            out.reserve(n);
×
580
            for (Py_ssize_t i = 0; i < n; ++i) {
×
581
                const char *s =
582
                    PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(seq, i));
×
583
                if (s) out.emplace_back(s);
×
584
            }
585
            Py_DECREF(seq);
×
586
            return out;
×
587
        };
×
588
        double time_interval_ms = pull_double("time_interval_ms", 5000.0);
×
589
        cfg->time_interval_us =
×
590
            static_cast<std::uint64_t>(time_interval_ms * 1000.0);
×
591
        cfg->compute_percentiles = pull_bool("compute_percentiles", false);
×
592
        cfg->extra_group_keys = pull_string_list("group_keys");
×
593
        cfg->custom_metric_fields = pull_string_list("custom_metric_fields");
×
594
        agg_config_ptr = std::move(cfg);
×
595
    }
×
596

597
    // owned_member_maps must outlive rt->submit: FileSlice::members is raw.
598
    std::vector<std::vector<GzipMember>> owned_member_maps;
×
599
    std::vector<IndexBuildBatchConfig::FileSlice> parsed_slices;
×
600
    if (file_slices_obj && file_slices_obj != Py_None) {
×
601
        PyObject *seq =
602
            PySequence_Fast(file_slices_obj, "file_slices must be a sequence");
×
603
        if (!seq) return NULL;
×
604
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
605
        if (static_cast<std::size_t>(n) != files.size()) {
×
606
            Py_DECREF(seq);
×
607
            PyErr_SetString(PyExc_ValueError,
×
608
                            "file_slices must match files length");
609
            return NULL;
×
610
        }
611
        owned_member_maps.resize(n);
×
612
        parsed_slices.resize(n);
×
613
        for (Py_ssize_t i = 0; i < n; ++i) {
×
614
            PyObject *entry = PySequence_Fast_GET_ITEM(seq, i);
×
615
            if (entry == Py_None) {
×
616
                continue;  // leave slice default-constructed (members=null)
×
617
            }
618
            Py_ssize_t mb = 0, me = 0, ckpt_base = 0;
×
619
            int skip_scoped = 0;
×
620
            PyObject *members_obj = nullptr;
×
621
            if (!PyArg_ParseTuple(entry, "nnnpO", &mb, &me, &ckpt_base,
×
622
                                  &skip_scoped, &members_obj)) {
623
                Py_DECREF(seq);
×
624
                return NULL;
×
625
            }
626
            PyObject *mseq = PySequence_Fast(
×
627
                members_obj, "file_slices[i].members must be a sequence");
628
            if (!mseq) {
×
629
                Py_DECREF(seq);
×
630
                return NULL;
×
631
            }
632
            Py_ssize_t mn = PySequence_Fast_GET_SIZE(mseq);
×
633
            auto &mv = owned_member_maps[i];
×
634
            mv.resize(mn);
×
635
            for (Py_ssize_t j = 0; j < mn; ++j) {
×
636
                PyObject *m = PySequence_Fast_GET_ITEM(mseq, j);
×
637
                unsigned long long c_offset = 0, c_size = 0;
×
638
                if (!PyArg_ParseTuple(m, "KK", &c_offset, &c_size)) {
×
639
                    Py_DECREF(mseq);
×
640
                    Py_DECREF(seq);
×
641
                    return NULL;
×
642
                }
643
                mv[j].c_offset = static_cast<std::uint64_t>(c_offset);
×
644
                mv[j].c_size = static_cast<std::uint64_t>(c_size);
×
645
            }
646
            Py_DECREF(mseq);
×
647
            parsed_slices[i].members = &mv;
×
648
            parsed_slices[i].member_begin = static_cast<std::size_t>(mb);
×
649
            parsed_slices[i].member_end = static_cast<std::size_t>(me);
×
650
            parsed_slices[i].checkpoint_idx_base =
×
651
                static_cast<std::uint64_t>(ckpt_base);
×
652
            parsed_slices[i].skip_file_scoped_writes = skip_scoped != 0;
×
653
        }
654
        Py_DECREF(seq);
×
655
    }
656

657
    auto batch_config = std::make_shared<IndexBuildBatchConfig>();
×
658
    batch_config->file_paths = std::move(files);
×
659
    batch_config->preassigned_file_ids = std::move(file_ids);
×
660
    if (!parsed_slices.empty()) {
×
661
        batch_config->file_slices = parsed_slices;
×
662
    }
663
    batch_config->index_dir = index_dir;
×
NEW
664
    batch_config->build_bloom = build_bloom != 0;
×
665
    batch_config->checkpoint_size = static_cast<std::size_t>(checkpoint_size);
×
666
    batch_config->build_manifest = build_manifest != 0;
×
667
    batch_config->force_rebuild = force_rebuild != 0;
×
668
    batch_config->bloom_dimensions = std::move(bloom_dims);
×
669
    batch_config->parallelism =
×
670
        parallelism > 0 ? static_cast<std::size_t>(parallelism)
×
671
                        : (rt ? std::max<std::size_t>(rt->threads(), 1) : 1);
×
672
    batch_config->flush_every_files =
×
673
        static_cast<std::size_t>(flush_every_files);
×
674
    batch_config->rebuild_root_summaries = false;
×
675

676
    // The build runs with the GIL released; re-acquire it per call. The
677
    // GIL-holding deleter drops the ref safely after the build.
NEW
678
    if (progress_obj && progress_obj != Py_None) {
×
NEW
679
        Py_INCREF(progress_obj);
×
NEW
680
        std::shared_ptr<PyObject> cb(progress_obj, [](PyObject *p) {
×
NEW
681
            PyGILState_STATE g = PyGILState_Ensure();
×
682
            Py_DECREF(p);
NEW
683
            PyGILState_Release(g);
×
NEW
684
        });
×
NEW
685
        batch_config->progress = [cb](std::size_t done, std::size_t total) {
×
NEW
686
            PyGILState_STATE g = PyGILState_Ensure();
×
NEW
687
            PyObject *r = PyObject_CallFunction(cb.get(), "nn",
×
688
                                                static_cast<Py_ssize_t>(done),
689
                                                static_cast<Py_ssize_t>(total));
NEW
690
            if (r) {
×
691
                Py_DECREF(r);
692
            } else {
693
                // A failing progress callback must not abort the build.
NEW
694
                PyErr_Clear();
×
695
            }
NEW
696
            PyGILState_Release(g);
×
NEW
697
        };
×
NEW
698
    }
×
699

700
    if (agg_config_ptr) {
×
701
        auto agg_staging = staging;
×
702
        auto agg_prefix = batch + "_agg";
×
703
        auto agg_intern = dftracer::utils::utilities::composites::dft::
×
NEW
704
            aggregators::intern_for_index(index_dir);
×
705
        // Counter keeps per-file SST dirs unique across duplicate file_paths
706
        // when one worker owns multiple slices of the same file.
707
        auto visitor_counter = std::make_shared<std::atomic<std::size_t>>(0);
×
708
        batch_config->dft_visitor_factory =
×
NEW
709
            [agg_staging, agg_prefix, agg_config_ptr, visitor_counter,
×
710
             agg_intern](const std::string &file_path)
711
            -> std::vector<std::unique_ptr<
712
                dftracer::utils::utilities::composites::dft::DftEventVisitor>> {
713
            using dftracer::utils::utilities::composites::dft::DftEventVisitor;
714
            using dftracer::utils::utilities::composites::dft::aggregators::
715
                AggregationVisitor;
716
            const std::size_t idx =
717
                visitor_counter->fetch_add(1, std::memory_order_relaxed);
×
718
            std::string prefix = agg_prefix + "_" + std::to_string(idx);
×
719
            std::vector<std::unique_ptr<DftEventVisitor>> visitors;
×
720
            visitors.push_back(std::make_unique<AggregationVisitor>(
×
721
                agg_staging, prefix, /*config_hash=*/0, *agg_config_ptr,
×
NEW
722
                file_path, agg_intern));
×
723
            return visitors;
×
724
        };
×
725
    }
×
726

727
    // Atomic: write phase calls sink_factory from N coroutines concurrently.
728
    auto batch_counter = std::make_shared<std::atomic<std::size_t>>(0);
×
729
    batch_config->sink_factory =
×
730
        [staging, batch, batch_counter]() -> std::unique_ptr<IndexBatchSink> {
×
731
        const std::size_t idx =
732
            batch_counter->fetch_add(1, std::memory_order_relaxed);
×
733
        std::string sub_batch = batch + "_" + std::to_string(idx);
×
734
        return std::make_unique<IndexDatabaseSstWriterContext>(staging,
×
735
                                                               sub_batch);
736
    };
×
737
    batch_config->sink_commit = [artifacts](IndexBatchSink &sink) {
×
738
        auto &sst = static_cast<IndexDatabaseSstWriterContext &>(sink);
×
739
        auto batch_artifacts = sst.commit();
×
740
        std::lock_guard<std::mutex> lock(artifacts->mu);
×
741
        if (!batch_artifacts.empty()) {
×
742
            artifacts->list.push_back(std::move(batch_artifacts));
×
743
        }
744
    };
×
745

746
    IndexBuildBatchResult result;
×
747
    if (!run_blocking([&] {
×
748
            rt->submit(dftracer::utils::run_coro_scope(
×
749
                           rt->executor(),
750
                           [](dftracer::utils::CoroScope &scope,
×
751
                              std::shared_ptr<IndexBuildBatchConfig> cfg,
752
                              IndexBuildBatchResult *out)
753
                               -> dftracer::utils::coro::CoroTask<void> {
×
754
                               *out =
×
755
                                   co_await IndexBatchBuilderUtility::process(
×
756
                                       &scope, std::move(cfg));
757
                           },
×
758
                           batch_config, &result),
×
759
                       "build-sst-batch")
×
760
                .get();
×
761
        })) {
×
762
        return NULL;
×
763
    }
764

765
    // If any file failed, surface the first error.
766
    if (result.failed > 0) {
×
767
        for (const auto &r : result.results) {
×
768
            if (!r.success) {
×
769
                PyErr_SetString(PyExc_RuntimeError, r.error_message.c_str());
×
770
                return NULL;
×
771
            }
772
        }
773
    }
774

775
    // One dict per committed sink + per-file aggregation below.
776
    PyObject *out_list = PyList_New(0);
×
777
    if (!out_list) return NULL;
×
778
    {
779
        std::lock_guard<std::mutex> lock(artifacts->mu);
×
780
        for (const auto &a : artifacts->list) {
×
781
            PyObject *main_dict = artifacts_to_dict(a);
×
782
            if (!main_dict || PyList_Append(out_list, main_dict) < 0) {
×
783
                Py_XDECREF(main_dict);
×
784
                Py_DECREF(out_list);
×
785
                return NULL;
×
786
            }
787
            Py_DECREF(main_dict);
×
788
        }
789
    }
×
790
    // Harvest per-file aggregation SSTs from extra visitors. Each visitor
791
    // holds a vector of Artifacts (one per FLUSH_THRESHOLD flush + the
792
    // file-complete flush) because SstFileWriter requires strictly
793
    // ascending keys per SST and cross-flush merge operands would collide.
794
    //
795
    // `extra_visitors` is indexed per input file, but a single
796
    // AggregationVisitor instance is typically shared across every file in
797
    // the batch (one flush at end-of-batch). Without dedup we would emit
798
    // that visitor's artifact dict N_files times, producing a manifest with
799
    // N copies of the same SST path. Dedup by visitor pointer so each
800
    // unique flush-sequence is emitted exactly once.
801
    using dftracer::utils::utilities::composites::dft::aggregators::
802
        AggregationVisitor;
803
    using dftracer::utils::utilities::composites::dft::aggregators::
804
        AssociationTracker;
805
    std::unordered_set<AggregationVisitor *> seen_visitors;
×
806
    for (auto &file_visitors : result.extra_visitors) {
×
807
        for (auto &visitor : file_visitors) {
×
808
            auto *agg = dynamic_cast<AggregationVisitor *>(visitor.get());
×
809
            if (!agg) continue;
×
810
            if (!seen_visitors.insert(agg).second) continue;
×
811
            for (auto &a : agg->aggregation_artifacts()) {
×
812
                if (a.empty()) continue;
×
813
                PyObject *agg_dict = artifacts_to_dict(a);
×
814
                if (!agg_dict || PyList_Append(out_list, agg_dict) < 0) {
×
815
                    Py_XDECREF(agg_dict);
×
816
                    Py_DECREF(out_list);
×
817
                    return NULL;
×
818
                }
819
                Py_DECREF(agg_dict);
×
820
            }
821
        }
822
    }
823

824
    AssociationTracker combined;
×
825
    bool any_tracker = false;
×
826
    for (auto *agg : seen_visitors) {
×
827
        auto out = agg->take_output();
×
828
        if (out.local_tracker) {
×
829
            combined.merge(*out.local_tracker);
×
830
            any_tracker = true;
×
831
        }
832
    }
×
833
    PyObject *tracker_bytes = nullptr;
×
834
    if (any_tracker) {
×
835
        combined.finalize();
×
836
        std::string blob = combined.serialize();
×
837
        tracker_bytes = PyBytes_FromStringAndSize(
×
838
            blob.data(), static_cast<Py_ssize_t>(blob.size()));
×
839
    } else {
×
840
        tracker_bytes = PyBytes_FromStringAndSize(nullptr, 0);
×
841
    }
842
    if (!tracker_bytes) {
×
843
        Py_DECREF(out_list);
×
844
        return NULL;
×
845
    }
846
    PyObject *ret = PyTuple_Pack(2, out_list, tracker_bytes);
×
847
    Py_DECREF(out_list);
×
848
    Py_DECREF(tracker_bytes);
×
849
    return ret;
×
850
}
×
851

852
static PyObject *enable_aggregation_deterministic_ids_fn(PyObject * /*self*/,
×
853
                                                         PyObject * /*args*/) {
854
    dftracer::utils::utilities::composites::dft::aggregators::
855
        enable_deterministic_intern_ids();
856
    Py_RETURN_NONE;
×
857
}
858

859
static PyObject *move_artifacts_fn(PyObject * /*self*/, PyObject *args,
×
860
                                   PyObject *kwds) {
861
    static const char *kwlist[] = {"artifacts", "dest_dir", NULL};
862
    PyObject *dict = NULL;
×
863
    const char *dest_dir = NULL;
×
864
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "Os", (char **)kwlist, &dict,
×
865
                                     &dest_dir)) {
866
        return NULL;
×
867
    }
868
    IndexDatabaseSstWriterContext::Artifacts a;
×
869
    if (!artifacts_from_dict(dict, &a)) return NULL;
×
870
    IndexDatabaseSstWriterContext::Artifacts moved;
×
871
    if (!run_blocking_r([&] { return std::move(a).move_to(dest_dir); }, moved))
×
872
        return NULL;
×
873
    return artifacts_to_dict(moved);
×
874
}
×
875

876
namespace {
877

878
dftracer::utils::coro::CoroTask<void> scan_one_gzip_file(
×
879
    std::string path, std::vector<GzipMember> *out) {
×
880
    out->clear();
881
    int fd = ::open(path.c_str(), O_RDONLY);
×
882
    if (fd < 0) co_return;
×
883
    struct stat st;
884
    if (::fstat(fd, &st) == 0 && st.st_size >= 18) {
×
885
        co_await enumerate_gzip_member_candidates(
×
886
            fd, static_cast<std::uint64_t>(st.st_size), *out);
887
    }
888
    ::close(fd);
×
889
}
×
890

891
}  // namespace
892

893
static PyObject *enumerate_gzip_members_fn(PyObject * /*self*/, PyObject *args,
×
894
                                           PyObject *kwds) {
895
    static const char *kwlist[] = {"files", "runtime", NULL};
896
    PyObject *files_obj = NULL;
×
897
    PyObject *runtime_arg = NULL;
×
898
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|O", (char **)kwlist,
×
899
                                     &files_obj, &runtime_arg)) {
900
        return NULL;
×
901
    }
902

903
    std::vector<std::string> files;
×
904
    {
905
        PyObject *seq = PySequence_Fast(files_obj, "files must be a sequence");
×
906
        if (!seq) return NULL;
×
907
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
908
        files.reserve(n);
×
909
        for (Py_ssize_t i = 0; i < n; ++i) {
×
910
            const char *s = PyUnicode_AsUTF8(PySequence_Fast_GET_ITEM(seq, i));
×
911
            if (!s) {
×
912
                Py_DECREF(seq);
×
913
                return NULL;
×
914
            }
915
            files.emplace_back(s);
×
916
        }
917
        Py_DECREF(seq);
×
918
    }
919

920
    Runtime *rt = nullptr;
×
921
    if (runtime_arg && runtime_arg != Py_None) {
×
922
        if (PyObject_TypeCheck(runtime_arg, &RuntimeType)) {
×
923
            rt = ((RuntimeObject *)runtime_arg)->runtime.get();
×
924
        } else {
925
            PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");
×
926
            if (!native || !PyObject_TypeCheck(native, &RuntimeType)) {
×
927
                Py_XDECREF(native);
×
928
                PyErr_SetString(PyExc_TypeError,
×
929
                                "runtime must be a Runtime instance or None");
930
                return NULL;
×
931
            }
932
            rt = ((RuntimeObject *)native)->runtime.get();
×
933
            Py_DECREF(native);
×
934
        }
935
    } else {
×
936
        rt = get_default_runtime();
×
937
    }
938

939
    std::vector<std::vector<GzipMember>> results(files.size());
×
940
    if (!run_blocking([&] {
×
941
            rt->submit(
×
942
                  dftracer::utils::run_coro_scope(
×
943
                      rt->executor(),
944
                      [](dftracer::utils::CoroScope &scope,
×
945
                         const std::vector<std::string> *paths,
946
                         std::vector<std::vector<GzipMember>> *out)
947
                          -> dftracer::utils::coro::CoroTask<void> {
×
948
                          co_await scope.scope(
×
949
                              [paths, out](dftracer::utils::CoroScope &child)
×
950
                                  -> dftracer::utils::coro::CoroTask<void> {
×
951
                                  for (std::size_t i = 0; i < paths->size();
×
952
                                       ++i) {
953
                                      const std::string &path = (*paths)[i];
954
                                      auto *slot = &(*out)[i];
955
                                      child.spawn(
×
956
                                          [path,
×
957
                                           slot](dftracer::utils::CoroScope &)
958
                                              -> dftracer::utils::coro::CoroTask<
959
                                                  void> {
×
960
                                              co_await scan_one_gzip_file(path,
×
961
                                                                          slot);
962
                                          });
×
963
                                  }
964
                                  co_return;
965
                              });
×
966
                          co_return;
967
                      },
×
968
                      &files, &results),
×
969
                  "enumerate-gzip-members")
×
970
                .get();
×
971
        })) {
×
972
        return NULL;
×
973
    }
974

975
    PyObject *out_list = PyList_New(static_cast<Py_ssize_t>(results.size()));
×
976
    if (!out_list) return NULL;
×
977
    for (std::size_t i = 0; i < results.size(); ++i) {
×
978
        const auto &mv = results[i];
×
979
        PyObject *inner = PyList_New(static_cast<Py_ssize_t>(mv.size()));
×
980
        if (!inner) {
×
981
            Py_DECREF(out_list);
×
982
            return NULL;
×
983
        }
984
        for (std::size_t j = 0; j < mv.size(); ++j) {
×
985
            PyObject *t =
986
                Py_BuildValue("(KK)", (unsigned long long)mv[j].c_offset,
×
987
                              (unsigned long long)mv[j].c_size);
×
988
            if (!t) {
×
989
                Py_DECREF(inner);
×
990
                Py_DECREF(out_list);
×
991
                return NULL;
×
992
            }
993
            PyList_SET_ITEM(inner, j, t);
×
994
        }
995
        PyList_SET_ITEM(out_list, i, inner);
×
996
    }
997
    return out_list;
×
998
}
×
999

1000
// Mirrors build_work_units + lpt_assign_units in dftracer_aggregator_mpi.cpp
1001
// so the Dask backend produces identical work distribution to MPI.
1002
static PyObject *plan_work_units_fn(PyObject * /*self*/, PyObject *args,
×
1003
                                    PyObject *kwds) {
1004
    static const char *kwlist[] = {"member_map", "num_workers", "target_c_size",
1005
                                   NULL};
1006
    PyObject *map_obj = NULL;
×
1007
    Py_ssize_t num_workers = 0;
×
1008
    unsigned long long target_c_size = 0;
×
1009
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "On|K", (char **)kwlist,
×
1010
                                     &map_obj, &num_workers, &target_c_size)) {
1011
        return NULL;
×
1012
    }
1013
    if (num_workers <= 0) num_workers = 1;
×
1014

1015
    std::vector<std::vector<GzipMember>> member_map;
×
1016
    {
1017
        PyObject *seq =
1018
            PySequence_Fast(map_obj, "member_map must be a sequence");
×
1019
        if (!seq) return NULL;
×
1020
        Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
×
1021
        member_map.resize(n);
×
1022
        for (Py_ssize_t i = 0; i < n; ++i) {
×
1023
            PyObject *inner = PySequence_Fast_GET_ITEM(seq, i);
×
1024
            PyObject *iseq =
1025
                PySequence_Fast(inner, "member_map[i] must be a sequence");
×
1026
            if (!iseq) {
×
1027
                Py_DECREF(seq);
×
1028
                return NULL;
×
1029
            }
1030
            Py_ssize_t ni = PySequence_Fast_GET_SIZE(iseq);
×
1031
            member_map[i].resize(ni);
×
1032
            for (Py_ssize_t j = 0; j < ni; ++j) {
×
1033
                PyObject *t = PySequence_Fast_GET_ITEM(iseq, j);
×
1034
                unsigned long long c_offset = 0, c_size = 0;
×
1035
                if (!PyArg_ParseTuple(t, "KK", &c_offset, &c_size)) {
×
1036
                    Py_DECREF(iseq);
×
1037
                    Py_DECREF(seq);
×
1038
                    return NULL;
×
1039
                }
1040
                member_map[i][j].c_offset =
×
1041
                    static_cast<std::uint64_t>(c_offset);
1042
                member_map[i][j].c_size = static_cast<std::uint64_t>(c_size);
×
1043
            }
1044
            Py_DECREF(iseq);
×
1045
        }
1046
        Py_DECREF(seq);
×
1047
    }
1048

1049
    // Fallback: treat empty/non-gzip files as a single whole-file member.
1050
    std::uint64_t total_c = 0;
×
1051
    for (auto &mv : member_map) {
×
1052
        if (mv.empty()) mv.push_back({0, 0});
×
1053
        for (const auto &m : mv) total_c += m.c_size;
×
1054
    }
1055

1056
    if (target_c_size == 0) {
×
1057
        target_c_size =
×
1058
            (total_c + static_cast<std::uint64_t>(num_workers) - 1) /
×
1059
            std::max<std::uint64_t>(static_cast<std::uint64_t>(num_workers), 1);
×
1060
    }
1061

1062
    struct Unit {
1063
        std::size_t file_idx;
1064
        std::size_t member_begin;
1065
        std::size_t member_end;
1066
        std::uint64_t c_size;
1067
    };
1068
    std::vector<Unit> units;
×
1069
    for (std::size_t fi = 0; fi < member_map.size(); ++fi) {
×
1070
        const auto &members = member_map[fi];
×
1071
        if (members.empty()) continue;
×
1072
        std::size_t begin = 0;
×
1073
        std::uint64_t accum = 0;
×
1074
        for (std::size_t i = 0; i < members.size(); ++i) {
×
1075
            accum += members[i].c_size;
×
1076
            const bool is_last = (i + 1 == members.size());
×
1077
            if ((target_c_size > 0 && accum >= target_c_size) || is_last) {
×
1078
                units.push_back({fi, begin, i + 1, accum});
×
1079
                begin = i + 1;
×
1080
                accum = 0;
×
1081
            }
1082
        }
1083
    }
1084

1085
    std::vector<std::size_t> order(units.size());
×
1086
    for (std::size_t i = 0; i < order.size(); ++i) order[i] = i;
×
1087
    std::sort(order.begin(), order.end(), [&](std::size_t a, std::size_t b) {
×
1088
        if (units[a].c_size != units[b].c_size)
×
1089
            return units[a].c_size > units[b].c_size;
×
1090
        if (units[a].file_idx != units[b].file_idx)
×
1091
            return units[a].file_idx < units[b].file_idx;
×
1092
        return units[a].member_begin < units[b].member_begin;
×
1093
    });
1094
    const std::size_t nw = static_cast<std::size_t>(num_workers);
×
1095
    std::vector<std::uint64_t> loads(nw, 0);
×
1096
    std::vector<std::vector<std::size_t>> per_worker(nw);
×
1097
    for (std::size_t ord : order) {
×
1098
        std::size_t best = 0;
×
1099
        for (std::size_t r = 1; r < nw; ++r)
×
1100
            if (loads[r] < loads[best]) best = r;
×
1101
        per_worker[best].push_back(ord);
×
1102
        loads[best] += std::max<std::uint64_t>(units[ord].c_size, 1);
×
1103
    }
1104

1105
    PyObject *out = PyList_New(static_cast<Py_ssize_t>(nw));
×
1106
    if (!out) return NULL;
×
1107
    for (std::size_t w = 0; w < nw; ++w) {
×
1108
        // Keep per-worker slices sorted by (file_idx, member_begin) for
1109
        // deterministic, file-group-friendly iteration downstream.
1110
        auto &lst = per_worker[w];
×
1111
        std::sort(lst.begin(), lst.end(), [&](std::size_t a, std::size_t b) {
×
1112
            if (units[a].file_idx != units[b].file_idx)
×
1113
                return units[a].file_idx < units[b].file_idx;
×
1114
            return units[a].member_begin < units[b].member_begin;
×
1115
        });
1116
        PyObject *inner = PyList_New(static_cast<Py_ssize_t>(lst.size()));
×
1117
        if (!inner) {
×
1118
            Py_DECREF(out);
×
1119
            return NULL;
×
1120
        }
1121
        for (std::size_t k = 0; k < lst.size(); ++k) {
×
1122
            const auto &u = units[lst[k]];
×
1123
            PyObject *t = Py_BuildValue(
×
1124
                "(nnnK)", (Py_ssize_t)u.file_idx, (Py_ssize_t)u.member_begin,
×
1125
                (Py_ssize_t)u.member_end, (unsigned long long)u.c_size);
×
1126
            if (!t) {
×
1127
                Py_DECREF(inner);
×
1128
                Py_DECREF(out);
×
1129
                return NULL;
×
1130
            }
1131
            PyList_SET_ITEM(inner, k, t);
×
1132
        }
1133
        PyList_SET_ITEM(out, w, inner);
×
1134
    }
1135
    return out;
×
1136
}
×
1137

1138
// ---------------------------------------------------------------------------
1139
// Module registration
1140
// ---------------------------------------------------------------------------
1141

1142
static PyMethodDef SstDistributionMethods[] = {
1143
    {"build_sst_batch", (PyCFunction)build_sst_batch_fn,
1144
     METH_VARARGS | METH_KEYWORDS,
1145
     "build_sst_batch(files, file_ids, staging_dir, batch_id, ...) "
1146
     "-> (list[dict], bytes)\n"
1147
     "Run the indexer pipeline with an SST sink and return "
1148
     "(artifact_dicts, tracker_blob). The tracker blob is the serialized "
1149
     "merged AssociationTracker from this batch's aggregation visitors "
1150
     "(empty bytes when no aggregation_config was passed)."},
1151
    {"plan_lpt_partition", (PyCFunction)plan_lpt_partition_fn, METH_VARARGS,
1152
     "plan_lpt_partition(entries, num_workers) -> list[list[(path, size)]]\n"
1153
     "Greedy Longest-Processing-Time-first bin-packing of (path, size) "
1154
     "tuples across num_workers buckets. Minimises the maximum per-worker "
1155
     "total size."},
1156
    {"scan_files", (PyCFunction)scan_files_fn, METH_VARARGS | METH_KEYWORDS,
1157
     "scan_files(directory, patterns=None, recursive=False, runtime=None) "
1158
     "-> list[(path, size)]\n"
1159
     "Parallel directory scan returning (path, size) tuples for regular "
1160
     "files matching the patterns."},
1161
    {"enable_aggregation_deterministic_ids",
1162
     (PyCFunction)enable_aggregation_deterministic_ids_fn, METH_NOARGS,
1163
     "enable_aggregation_deterministic_ids() -> None\n"
1164
     "Flip the global aggregation StringIntern into deterministic-id mode "
1165
     "so the same string maps to the same 32-bit id in every worker "
1166
     "process. Call once at worker startup BEFORE any aggregation work."},
1167
    {"move_artifacts", (PyCFunction)move_artifacts_fn,
1168
     METH_VARARGS | METH_KEYWORDS,
1169
     "move_artifacts(artifacts, dest_dir) -> dict\n"
1170
     "Move every populated SST in `artifacts` (as returned by "
1171
     "`build_sst_batch`) into `dest_dir` via the C++ rename/copy helper, "
1172
     "returning a fresh dict with the new paths. Single GIL release, no "
1173
     "per-file Python shutil.move overhead."},
1174
    {"enumerate_gzip_members", (PyCFunction)enumerate_gzip_members_fn,
1175
     METH_VARARGS | METH_KEYWORDS,
1176
     "enumerate_gzip_members(files, runtime=None) -> list[list[(c_offset, "
1177
     "c_size)]]\n"
1178
     "Cooperative async scan of gzip member offsets across `files`. "
1179
     "Returns lists of (c_offset, c_size) parallel to `files`; empty for "
1180
     "non-gzip / unreadable files."},
1181
    {"plan_work_units", (PyCFunction)plan_work_units_fn,
1182
     METH_VARARGS | METH_KEYWORDS,
1183
     "plan_work_units(member_map, num_workers, target_c_size=0) "
1184
     "-> list[list[(file_idx, member_begin, member_end, c_size)]]\n"
1185
     "Deterministic LPT assignment of intra-file gzip-member slices "
1186
     "across workers. Each worker's list contains (file_idx, "
1187
     "member_begin, member_end, c_size) tuples; a file sliced across "
1188
     "multiple workers appears in each owner's list with disjoint "
1189
     "[member_begin, member_end) ranges."},
1190
    {NULL, NULL, 0, NULL}};
1191

1192
int init_sst_distribution(PyObject *m) {
2✔
1193
    if (register_type(m, &SstArtifactRegistryType, "SstArtifactRegistry") < 0)
2✔
1194
        return -1;
×
1195
    if (PyModule_AddFunctions(m, SstDistributionMethods) < 0) return -1;
2✔
1196
    return 0;
2✔
1197
}
1✔
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE TRIAL · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc