30002033718

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

Build # 30002033718

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge df0b9588d into 06bc84ec9

Pull Request Pull Request #99: Support CM time_metric (NS/MS/SEC/US) across reader and viz

Coverage Stats

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

49.62

/src/dftracer/utils/python/runtime.cpp

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <dftracer/utils/core/rocksdb/database.h>
#include <dftracer/utils/python/py_dict_helpers.h>
#include <dftracer/utils/python/py_errors.h>
#include <dftracer/utils/python/py_runtime_mixin.h>
#include <dftracer/utils/python/py_type_helpers.h>
#include <dftracer/utils/python/runtime.h>

#include <chrono>
#include <memory>

static std::shared_ptr<dftracer::utils::Runtime> g_default_runtime;

dftracer::utils::Runtime *get_default_runtime() {
    if (!g_default_runtime) {
        g_default_runtime = std::make_shared<dftracer::utils::Runtime>(0);
    }
    return g_default_runtime.get();
}

static void Runtime_dealloc(RuntimeObject *self) {
    self->runtime.reset();
    Py_TYPE(self)->tp_free((PyObject *)self);
}

static PyObject *Runtime_new(PyTypeObject *type, PyObject *args,
                             PyObject *kwds) {
    RuntimeObject *self = (RuntimeObject *)type->tp_alloc(type, 0);
    if (self) {
        // Placement-new the shared_ptr (tp_alloc gives raw memory)
        new (&self->runtime) std::shared_ptr<dftracer::utils::Runtime>(nullptr);
    }
    return (PyObject *)self;
}

static int Runtime_init(RuntimeObject *self, PyObject *args, PyObject *kwds) {
    static const char *kwlist[] = {"threads", "io_threads", NULL};
    Py_ssize_t threads = 0;
    Py_ssize_t io_threads = 0;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|nn", (char **)kwlist,
                                     &threads, &io_threads)) {
        return -1;
    }

    if (threads < 0) {
        PyErr_SetString(PyExc_ValueError, "threads must be >= 0");
        return -1;
    }
    if (io_threads < 0) {
        PyErr_SetString(PyExc_ValueError, "io_threads must be >= 0");
        return -1;
    }

    try {
        dftracer::utils::ExecutorConfig config;
        config.num_threads = static_cast<std::size_t>(threads);
        config.io_pool_size = static_cast<std::size_t>(io_threads);
        self->runtime =
            std::make_shared<dftracer::utils::Runtime>(config, true);
    } catch (const std::exception &e) {
        set_typed_py_error(e);
        return -1;
    }

    return 0;
}

static PyObject *Runtime_shutdown(RuntimeObject *self,
                                  PyObject *Py_UNUSED(ignored)) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }
    Py_BEGIN_ALLOW_THREADS self->runtime->shutdown();
    Py_END_ALLOW_THREADS Py_RETURN_NONE;
}

static bool set_size(PyObject *d, const char *key, std::size_t val) {
    return dict_set_size(d, key, val) == 0;
}

static bool set_double(PyObject *d, const char *key, double val) {
    return dict_set_f64(d, key, val) == 0;
}

static bool set_str(PyObject *d, const char *key, const std::string &val) {
    return dict_set_str(d, key, val.c_str()) == 0;
}

static bool set_bool(PyObject *d, const char *key, bool val) {
    return dict_set_bool(d, key, val) == 0;
}

static PyObject *build_task_progress(const dftracer::utils::TaskProgress &tp) {
    PyObject *td = PyDict_New();
    if (!td) return NULL;

    if (!set_str(td, "name", tp.name) || !set_str(td, "state", tp.state) ||
        !set_double(td, "queued_duration_ms", tp.queued_duration_ms) ||
        !set_double(td, "execution_duration_ms", tp.execution_duration_ms) ||
        !set_size(td, "total_subtasks", tp.total_subtasks) ||
        !set_size(td, "completed_subtasks", tp.completed_subtasks) ||
        !set_double(td, "progress_pct", tp.progress_percentage) ||
        !set_str(td, "location", tp.location)) {
        Py_DECREF(td);
        return NULL;
    }

    PyObject *children =
        PyList_New(static_cast<Py_ssize_t>(tp.children.size()));
    if (!children) {
        Py_DECREF(td);
        return NULL;
    }
    for (std::size_t i = 0; i < tp.children.size(); ++i) {
        PyObject *child = build_task_progress(tp.children[i]);
        if (!child) {
            Py_DECREF(children);
            Py_DECREF(td);
            return NULL;
        }
        PyList_SET_ITEM(children, static_cast<Py_ssize_t>(i), child);
    }
    if (PyDict_SetItemString(td, "children", children) < 0) {
        Py_DECREF(children);
        Py_DECREF(td);
        return NULL;
    }
    Py_DECREF(children);
    return td;
}

static PyObject *Runtime_get_progress(RuntimeObject *self,
                                      PyObject *Py_UNUSED(ignored)) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }

    dftracer::utils::ExecutorProgress prog;
    Py_BEGIN_ALLOW_THREADS prog = self->runtime->get_progress();
    Py_END_ALLOW_THREADS

        PyObject *d = PyDict_New();
    if (!d) return NULL;

    if (!set_size(d, "total", prog.total_tasks_submitted) ||
        !set_size(d, "completed", prog.tasks_completed) ||
        !set_size(d, "running", prog.tasks_running) ||
        !set_size(d, "queued", prog.tasks_queued) ||
        !set_size(d, "failed", prog.tasks_failed)) {
        Py_DECREF(d);
        return NULL;
    }

    // Workers
    PyObject *workers =
        PyList_New(static_cast<Py_ssize_t>(prog.workers.size()));
    if (!workers) {
        Py_DECREF(d);
        return NULL;
    }
    for (std::size_t i = 0; i < prog.workers.size(); ++i) {
        const auto &w = prog.workers[i];
        PyObject *wd = PyDict_New();
        if (!wd || !set_size(wd, "id", w.worker_id) ||
            !set_bool(wd, "idle", w.is_idle) ||
            !set_str(wd, "task", w.current_task_name) ||
            !set_size(wd, "queue_depth", w.local_queue_depth)) {
            Py_XDECREF(wd);
            Py_DECREF(workers);
            Py_DECREF(d);
            return NULL;
        }
        PyList_SET_ITEM(workers, static_cast<Py_ssize_t>(i), wd);
    }
    if (PyDict_SetItemString(d, "workers", workers) < 0) {
        Py_DECREF(workers);
        Py_DECREF(d);
        return NULL;
    }
    Py_DECREF(workers);

    // Tasks
    PyObject *tasks =
        PyList_New(static_cast<Py_ssize_t>(prog.root_tasks.size()));
    if (!tasks) {
        Py_DECREF(d);
        return NULL;
    }
    for (std::size_t i = 0; i < prog.root_tasks.size(); ++i) {
        PyObject *tp = build_task_progress(prog.root_tasks[i]);
        if (!tp) {
            Py_DECREF(tasks);
            Py_DECREF(d);
            return NULL;
        }
        PyList_SET_ITEM(tasks, static_cast<Py_ssize_t>(i), tp);
    }
    if (PyDict_SetItemString(d, "tasks", tasks) < 0) {
        Py_DECREF(tasks);
        Py_DECREF(d);
        return NULL;
    }
    Py_DECREF(tasks);

    // Errors
    PyObject *errors =
        PyList_New(static_cast<Py_ssize_t>(prog.recent_errors.size()));
    if (!errors) {
        Py_DECREF(d);
        return NULL;
    }
    for (std::size_t i = 0; i < prog.recent_errors.size(); ++i) {
        const auto &[tid, msg] = prog.recent_errors[i];
        PyObject *ed = PyDict_New();
        if (!ed || !set_size(ed, "task_id", static_cast<std::size_t>(tid)) ||
            !set_str(ed, "message", msg)) {
            Py_XDECREF(ed);
            Py_DECREF(errors);
            Py_DECREF(d);
            return NULL;
        }
        PyList_SET_ITEM(errors, static_cast<Py_ssize_t>(i), ed);
    }
    if (PyDict_SetItemString(d, "errors", errors) < 0) {
        Py_DECREF(errors);
        Py_DECREF(d);
        return NULL;
    }
    Py_DECREF(errors);

    return d;
}

static PyObject *Runtime_is_responsive(RuntimeObject *self,
                                       PyObject *Py_UNUSED(ignored)) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }
    bool resp;
    Py_BEGIN_ALLOW_THREADS resp = self->runtime->is_responsive();
    Py_END_ALLOW_THREADS return PyBool_FromLong(resp ? 1 : 0);
}

static PyObject *Runtime_set_timeout(RuntimeObject *self, PyObject *args,
                                     PyObject *kwds) {
    static const char *kwlist[] = {"global_ms", NULL};
    Py_ssize_t ms = 0;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|n", (char **)kwlist, &ms)) {
        return NULL;
    }

    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS self->runtime->set_global_timeout(
        std::chrono::milliseconds(ms));
    Py_END_ALLOW_THREADS Py_RETURN_NONE;
}

static PyObject *Runtime_set_default_task_timeout(RuntimeObject *self,
                                                  PyObject *args,
                                                  PyObject *kwds) {
    static const char *kwlist[] = {"ms", NULL};
    Py_ssize_t ms = 0;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|n", (char **)kwlist, &ms)) {
        return NULL;
    }

    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }

    Py_BEGIN_ALLOW_THREADS self->runtime->set_default_task_timeout(
        std::chrono::milliseconds(ms));
    Py_END_ALLOW_THREADS Py_RETURN_NONE;
}

static PyObject *Runtime_wait_all(RuntimeObject *self,
                                  PyObject *Py_UNUSED(ignored)) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }
    if (!run_blocking([&] { self->runtime->wait_all(); })) return NULL;
    Py_RETURN_NONE;
}

static PyObject *Runtime_enter(RuntimeObject *self,
                               PyObject *Py_UNUSED(ignored)) {
    Py_INCREF(self);
    return (PyObject *)self;
}

static PyObject *Runtime_exit(RuntimeObject *self, PyObject *args) {
    if (self->runtime) {
        Py_BEGIN_ALLOW_THREADS self->runtime->shutdown();
        Py_END_ALLOW_THREADS
    }
    Py_RETURN_NONE;
}

static PyObject *Runtime_get_threads(RuntimeObject *self, void *closure) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }
    return PyLong_FromSize_t(self->runtime->threads());
}

static PyObject *get_default_runtime_py(PyObject *Py_UNUSED(module),
                                        PyObject *Py_UNUSED(ignored)) {
    dftracer::utils::Runtime *rt = get_default_runtime();
    if (!rt) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to create default runtime");
        return NULL;
    }

    RuntimeObject *obj = (RuntimeObject *)RuntimeType.tp_alloc(&RuntimeType, 0);
    if (!obj) return NULL;

    new (&obj->runtime)
        std::shared_ptr<dftracer::utils::Runtime>(g_default_runtime);
    return (PyObject *)obj;
}

// Return the current default runtime without creating one. Unlike
// get_default_runtime, this never materializes a full-machine-sized runtime as
// a side effect, so callers that only want to save/restore the default (e.g. a
// Dask worker plugin) do not each spin up an unused hardware_concurrency-thread
// runtime.
static PyObject *peek_default_runtime_py(PyObject *Py_UNUSED(module),
                                         PyObject *Py_UNUSED(ignored)) {
    if (!g_default_runtime) Py_RETURN_NONE;

    RuntimeObject *obj = (RuntimeObject *)RuntimeType.tp_alloc(&RuntimeType, 0);
    if (!obj) return NULL;

    new (&obj->runtime)
        std::shared_ptr<dftracer::utils::Runtime>(g_default_runtime);
    return (PyObject *)obj;
}

static PyObject *set_default_runtime_py(PyObject *Py_UNUSED(module),
                                        PyObject *args) {
    PyObject *arg;
    if (!PyArg_ParseTuple(args, "O", &arg)) return NULL;

    if (arg == Py_None) {
        g_default_runtime.reset();
        Py_RETURN_NONE;
    }

    if (!PyObject_TypeCheck(arg, &RuntimeType)) {
        PyErr_SetString(PyExc_TypeError, "Expected Runtime or None");
        return NULL;
    }

    g_default_runtime = ((RuntimeObject *)arg)->runtime;
    Py_RETURN_NONE;
}

static PyMethodDef Runtime_methods[] = {
    {"shutdown", (PyCFunction)Runtime_shutdown, METH_NOARGS,
     "shutdown()\n"
     "--\n"
     "\n"
     "Shut down the runtime.\n"},
    {"get_progress", (PyCFunction)Runtime_get_progress, METH_NOARGS,
     "Return progress dict with keys: total, completed, running,\n"
     "queued, failed."},
    {"is_responsive", (PyCFunction)Runtime_is_responsive, METH_NOARGS,
     "Return True if the runtime is making progress."},
    {"set_timeout", (PyCFunction)Runtime_set_timeout,
     METH_VARARGS | METH_KEYWORDS,
     "Set global timeout in milliseconds.\n"
     "\n"
     "Args:\n"
     "    global_ms (int): Timeout in milliseconds (0 = no timeout).\n"},
    {"set_default_task_timeout", (PyCFunction)Runtime_set_default_task_timeout,
     METH_VARARGS | METH_KEYWORDS,
     "Set default per-task timeout in milliseconds.\n"
     "\n"
     "Args:\n"
     "    ms (int): Timeout in milliseconds (0 = no timeout).\n"},
    {"wait_all", (PyCFunction)Runtime_wait_all, METH_NOARGS,
     "Wait for all outstanding submitted tasks to complete."},
    {"__enter__", (PyCFunction)Runtime_enter, METH_NOARGS,
     "Enter context manager."},
    {"__exit__", (PyCFunction)Runtime_exit, METH_VARARGS,
     "Exit context manager (calls shutdown)."},
    {NULL}};

static PyObject *Runtime_get_io_threads(RuntimeObject *self, void *closure) {
    if (!self->runtime) {
        PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");
        return NULL;
    }
    return PyLong_FromSize_t(self->runtime->io_threads());
}

static PyGetSetDef Runtime_getsetters[] = {
    {"threads", (getter)Runtime_get_threads, NULL, "Number of worker threads",
     NULL},
    {"io_threads", (getter)Runtime_get_io_threads, NULL,
     "Number of I/O threads", NULL},
    {NULL}};

PyTypeObject RuntimeType = {
    PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.Runtime",
    sizeof(RuntimeObject),                    /* tp_basicsize */
    0,                                        /* tp_itemsize */
    (destructor)Runtime_dealloc,              /* tp_dealloc */
    0,                                        /* tp_vectorcall_offset */
    0,                                        /* tp_getattr */
    0,                                        /* tp_setattr */
    0,                                        /* tp_as_async */
    0,                                        /* tp_repr */
    0,                                        /* tp_as_number */
    0,                                        /* tp_as_sequence */
    0,                                        /* tp_as_mapping */
    0,                                        /* tp_hash */
    0,                                        /* tp_call */
    0,                                        /* tp_str */
    0,                                        /* tp_getattro */
    0,                                        /* tp_setattro */
    0,                                        /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
    "Runtime(threads: int = 0, io_threads: int = 0)\n"
    "--\n"
    "\n"
    "Coroutine runtime backed by a thread pool.\n"
    "\n"
    "Args:\n"
    "    threads (int): Number of worker threads. 0 (default) uses\n"
    "        the hardware concurrency.\n"
    "    io_threads (int): Number of I/O threads. 0 (default) uses\n"
    "        the hardware concurrency.\n", /* tp_doc */
    0,                                     /* tp_traverse */
    0,                                     /* tp_clear */
    0,                                     /* tp_richcompare */
    0,                                     /* tp_weaklistoffset */
    0,                                     /* tp_iter */
    0,                                     /* tp_iternext */
    Runtime_methods,                       /* tp_methods */
    0,                                     /* tp_members */
    Runtime_getsetters,                    /* tp_getset */
    0,                                     /* tp_base */
    0,                                     /* tp_dict */
    0,                                     /* tp_descr_get */
    0,                                     /* tp_descr_set */
    0,                                     /* tp_dictoffset */
    (initproc)Runtime_init,                /* tp_init */
    0,                                     /* tp_alloc */
    Runtime_new,                           /* tp_new */
};

// Module-level function table (registered via PyModule_AddFunctions or
// appended to the module's method table in init_runtime).
static PyMethodDef runtime_module_methods[] = {
    {"get_default_runtime", get_default_runtime_py, METH_NOARGS,
     "Return the module-level default Runtime (lazy-created)."},
    {"peek_default_runtime", peek_default_runtime_py, METH_NOARGS,
     "Return the current default Runtime, or None if none exists yet "
     "(never creates one)."},
    {"set_default_runtime", set_default_runtime_py, METH_VARARGS,
     "Replace the module-level default Runtime (pass None to clear).\n"
     "\n"
     "Args:\n"
     "    runtime (Runtime or None): New default runtime.\n"},
    {NULL}};

// Runs during interpreter finalization. Stop the default runtime's worker and
// I/O threads while the process is still healthy, then tell RocksDB we are
// exiting so cached DB handles skip closing every open SST on teardown. A read
// scan can leave hundreds of SSTs open (more so on a networked filesystem),
// and without this the process can appear to hang after the work is done.
static void dftracer_utils_atexit_cleanup() {
    if (g_default_runtime) {
        g_default_runtime->shutdown();
    }
    dftracer::utils::rocksdb::mark_process_exiting_for_rocksdb();
}

int init_runtime(PyObject *m) {
    if (register_type(m, &RuntimeType, "Runtime") < 0) return -1;

    Py_AtExit(dftracer_utils_atexit_cleanup);

    for (PyMethodDef *def = runtime_module_methods; def->ml_name; ++def) {
        PyObject *fn = PyCFunction_New(def, NULL);
        if (!fn) return -1;
        if (PyModule_AddObject(m, def->ml_name, fn) < 0) {
            Py_DECREF(fn);
            return -1;
        }
    }

    return 0;
}

1	#define PY_SSIZE_T_CLEAN
2	#include <Python.h>
3	#include <dftracer/utils/core/rocksdb/database.h>
4	#include <dftracer/utils/python/py_dict_helpers.h>
5	#include <dftracer/utils/python/py_errors.h>
6	#include <dftracer/utils/python/py_runtime_mixin.h>
7	#include <dftracer/utils/python/py_type_helpers.h>
8	#include <dftracer/utils/python/runtime.h>
9
10	#include <chrono>
11	#include <memory>
12
13	static std::shared_ptr<dftracer::utils::Runtime> g_default_runtime;
14
15	dftracer::utils::Runtime *get_default_runtime() {	744✔
16	if (!g_default_runtime) {	744✔
17	g_default_runtime = std::make_shared<dftracer::utils::Runtime>(0);	2!
18	}	1✔
19	return g_default_runtime.get();	744✔
20	}
21
22	static void Runtime_dealloc(RuntimeObject *self) {	158✔
23	self->runtime.reset();	158✔
24	Py_TYPE(self)->tp_free((PyObject *)self);	158✔
25	}	158✔
26
27	static PyObject Runtime_new(PyTypeObject type, PyObject *args,	156✔
28	PyObject *kwds) {
29	RuntimeObject self = (RuntimeObject )type->tp_alloc(type, 0);	156✔
30	if (self) {	156✔
31	// Placement-new the shared_ptr (tp_alloc gives raw memory)
32	new (&self->runtime) std::shared_ptr<dftracer::utils::Runtime>(nullptr);	156✔
33	}	78✔
34	return (PyObject *)self;	156✔
35	}
36
37	static int Runtime_init(RuntimeObject self, PyObject args, PyObject *kwds) {	156✔
38	static const char *kwlist[] = {"threads", "io_threads", NULL};
39	Py_ssize_t threads = 0;	156✔
40	Py_ssize_t io_threads = 0;	156✔
41
42	if (!PyArg_ParseTupleAndKeywords(args, kwds, "\|nn", (char **)kwlist,	156!
43	&threads, &io_threads)) {
44	return -1;	×
45	}
46
47	if (threads < 0) {	156✔
48	PyErr_SetString(PyExc_ValueError, "threads must be >= 0");	×
49	return -1;	×
50	}
51	if (io_threads < 0) {	156✔
52	PyErr_SetString(PyExc_ValueError, "io_threads must be >= 0");	×
53	return -1;	×
54	}
55
56	try {
57	dftracer::utils::ExecutorConfig config;	156!
58	config.num_threads = static_cast<std::size_t>(threads);	156✔
59	config.io_pool_size = static_cast<std::size_t>(io_threads);	156✔
60	self->runtime =	78✔
61	std::make_shared<dftracer::utils::Runtime>(config, true);	156!
62	} catch (const std::exception &e) {	78!
63	set_typed_py_error(e);	×
64	return -1;	×
65	}	×
66
67	return 0;	156✔
68	}	78✔
69
70	static PyObject Runtime_shutdown(RuntimeObject self,	156✔
71	PyObject *Py_UNUSED(ignored)) {
72	if (!self->runtime) {	156!
73	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
74	return NULL;	×
75	}
76	Py_BEGIN_ALLOW_THREADS self->runtime->shutdown();	156✔
77	Py_END_ALLOW_THREADS Py_RETURN_NONE;	156✔
78	}	78✔
79
80	static bool set_size(PyObject d, const char key, std::size_t val) {	150✔
81	return dict_set_size(d, key, val) == 0;	150✔
82	}
83
84	static bool set_double(PyObject d, const char key, double val) {	42✔
85	return dict_set_f64(d, key, val) == 0;	42✔
86	}
87
88	static bool set_str(PyObject d, const char key, const std::string &val) {	58✔
89	return dict_set_str(d, key, val.c_str()) == 0;	58✔
90	}
91
92	static bool set_bool(PyObject d, const char key, bool val) {	16✔
93	return dict_set_bool(d, key, val) == 0;	16✔
94	}
95
96	static PyObject *build_task_progress(const dftracer::utils::TaskProgress &tp) {	14✔
97	PyObject *td = PyDict_New();	14✔
98	if (!td) return NULL;	14✔
99
100	if (!set_str(td, "name", tp.name) \|\| !set_str(td, "state", tp.state) \|\|	21!
101	!set_double(td, "queued_duration_ms", tp.queued_duration_ms) \|\|	14!
102	!set_double(td, "execution_duration_ms", tp.execution_duration_ms) \|\|	14!
103	!set_size(td, "total_subtasks", tp.total_subtasks) \|\|	14!
104	!set_size(td, "completed_subtasks", tp.completed_subtasks) \|\|	14!
105	!set_double(td, "progress_pct", tp.progress_percentage) \|\|	28!
106	!set_str(td, "location", tp.location)) {	14!
107	Py_DECREF(td);
108	return NULL;	×
109	}
110
111	PyObject *children =	7✔
112	PyList_New(static_cast<Py_ssize_t>(tp.children.size()));	14✔
113	if (!children) {	14✔
114	Py_DECREF(td);
115	return NULL;	×
116	}
117	for (std::size_t i = 0; i < tp.children.size(); ++i) {	14!
118	PyObject *child = build_task_progress(tp.children[i]);	×
119	if (!child) {	×
120	Py_DECREF(children);
121	Py_DECREF(td);
122	return NULL;	×
123	}
124	PyList_SET_ITEM(children, static_cast<Py_ssize_t>(i), child);	×
125	}
126	if (PyDict_SetItemString(td, "children", children) < 0) {	14✔
127	Py_DECREF(children);
128	Py_DECREF(td);
129	return NULL;	×
130	}
131	Py_DECREF(children);	7✔
132	return td;	14✔
133	}	7✔
134
135	static PyObject Runtime_get_progress(RuntimeObject self,	18✔
136	PyObject *Py_UNUSED(ignored)) {
137	if (!self->runtime) {	18!
138	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
139	return NULL;	×
140	}
141
142	dftracer::utils::ExecutorProgress prog;	18✔
143	Py_BEGIN_ALLOW_THREADS prog = self->runtime->get_progress();	18!
144	Py_END_ALLOW_THREADS	18!
145
146	PyObject *d = PyDict_New();	18!
147	if (!d) return NULL;	18✔
148
149	if (!set_size(d, "total", prog.total_tasks_submitted) \|\|	18!
150	!set_size(d, "completed", prog.tasks_completed) \|\|	18!
151	!set_size(d, "running", prog.tasks_running) \|\|	18!
152	!set_size(d, "queued", prog.tasks_queued) \|\|	45!
153	!set_size(d, "failed", prog.tasks_failed)) {	18!
154	Py_DECREF(d);	×
155	return NULL;	×
156	}
157
158	// Workers
159	PyObject *workers =	9✔
160	PyList_New(static_cast<Py_ssize_t>(prog.workers.size()));	18!
161	if (!workers) {	18!
162	Py_DECREF(d);	×
163	return NULL;	×
164	}
165	for (std::size_t i = 0; i < prog.workers.size(); ++i) {	34✔
166	const auto &w = prog.workers[i];	16✔
167	PyObject *wd = PyDict_New();	16!
168	if (!wd \|\| !set_size(wd, "id", w.worker_id) \|\|	16!
169	!set_bool(wd, "idle", w.is_idle) \|\|	16!
170	!set_str(wd, "task", w.current_task_name) \|\|	40!
171	!set_size(wd, "queue_depth", w.local_queue_depth)) {	16!
172	Py_XDECREF(wd);	×
173	Py_DECREF(workers);	×
174	Py_DECREF(d);	×
175	return NULL;	×
176	}
177	PyList_SET_ITEM(workers, static_cast<Py_ssize_t>(i), wd);	16!
178	}	8✔
179	if (PyDict_SetItemString(d, "workers", workers) < 0) {	18!
180	Py_DECREF(workers);	×
181	Py_DECREF(d);	×
182	return NULL;	×
183	}
184	Py_DECREF(workers);	9!
185
186	// Tasks
187	PyObject *tasks =	9✔
188	PyList_New(static_cast<Py_ssize_t>(prog.root_tasks.size()));	18!
189	if (!tasks) {	18!
190	Py_DECREF(d);	×
191	return NULL;	×
192	}
193	for (std::size_t i = 0; i < prog.root_tasks.size(); ++i) {	32✔
194	PyObject *tp = build_task_progress(prog.root_tasks[i]);	14!
195	if (!tp) {	14!
196	Py_DECREF(tasks);	×
197	Py_DECREF(d);	×
198	return NULL;	×
199	}
200	PyList_SET_ITEM(tasks, static_cast<Py_ssize_t>(i), tp);	14!
201	}	7✔
202	if (PyDict_SetItemString(d, "tasks", tasks) < 0) {	18!
203	Py_DECREF(tasks);	×
204	Py_DECREF(d);	×
205	return NULL;	×
206	}
207	Py_DECREF(tasks);	9!
208
209	// Errors
210	PyObject *errors =	9✔
211	PyList_New(static_cast<Py_ssize_t>(prog.recent_errors.size()));	18!
212	if (!errors) {	18!
213	Py_DECREF(d);	×
214	return NULL;	×
215	}
216	for (std::size_t i = 0; i < prog.recent_errors.size(); ++i) {	18✔
217	const auto &[tid, msg] = prog.recent_errors[i];	×
218	PyObject *ed = PyDict_New();	×
219	if (!ed \|\| !set_size(ed, "task_id", static_cast<std::size_t>(tid)) \|\|	×
220	!set_str(ed, "message", msg)) {	×
221	Py_XDECREF(ed);	×
222	Py_DECREF(errors);	×
223	Py_DECREF(d);	×
224	return NULL;	×
225	}
226	PyList_SET_ITEM(errors, static_cast<Py_ssize_t>(i), ed);	×
227	}
228	if (PyDict_SetItemString(d, "errors", errors) < 0) {	18!
229	Py_DECREF(errors);	×
230	Py_DECREF(d);	×
231	return NULL;	×
232	}
233	Py_DECREF(errors);	9!
234
235	return d;	18✔
236	}	18✔
237
238	static PyObject Runtime_is_responsive(RuntimeObject self,	2✔
239	PyObject *Py_UNUSED(ignored)) {
240	if (!self->runtime) {	2!
241	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
242	return NULL;	×
243	}
244	bool resp;
245	Py_BEGIN_ALLOW_THREADS resp = self->runtime->is_responsive();	2✔
246	Py_END_ALLOW_THREADS return PyBool_FromLong(resp ? 1 : 0);	2!
247	}	1✔
248
249	static PyObject Runtime_set_timeout(RuntimeObject self, PyObject *args,	×
250	PyObject *kwds) {
251	static const char *kwlist[] = {"global_ms", NULL};
252	Py_ssize_t ms = 0;	×
253
254	if (!PyArg_ParseTupleAndKeywords(args, kwds, "\|n", (char **)kwlist, &ms)) {	×
255	return NULL;	×
256	}
257
258	if (!self->runtime) {	×
259	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
260	return NULL;	×
261	}
262
263	Py_BEGIN_ALLOW_THREADS self->runtime->set_global_timeout(	×
264	std::chrono::milliseconds(ms));
265	Py_END_ALLOW_THREADS Py_RETURN_NONE;	×
266	}
267
268	static PyObject Runtime_set_default_task_timeout(RuntimeObject self,	×
269	PyObject *args,
270	PyObject *kwds) {
271	static const char *kwlist[] = {"ms", NULL};
272	Py_ssize_t ms = 0;	×
273
274	if (!PyArg_ParseTupleAndKeywords(args, kwds, "\|n", (char **)kwlist, &ms)) {	×
275	return NULL;	×
276	}
277
278	if (!self->runtime) {	×
279	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
280	return NULL;	×
281	}
282
283	Py_BEGIN_ALLOW_THREADS self->runtime->set_default_task_timeout(	×
284	std::chrono::milliseconds(ms));
285	Py_END_ALLOW_THREADS Py_RETURN_NONE;	×
286	}
287
288	static PyObject Runtime_wait_all(RuntimeObject self,	188✔
289	PyObject *Py_UNUSED(ignored)) {
290	if (!self->runtime) {	188!
291	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
292	return NULL;	×
293	}
294	if (!run_blocking([&] { self->runtime->wait_all(); })) return NULL;	376!
295	Py_RETURN_NONE;	188✔
296	}	94✔
297
298	static PyObject Runtime_enter(RuntimeObject self,	×
299	PyObject *Py_UNUSED(ignored)) {
300	Py_INCREF(self);
301	return (PyObject *)self;	×
302	}
303
304	static PyObject Runtime_exit(RuntimeObject self, PyObject *args) {	×
305	if (self->runtime) {	×
306	Py_BEGIN_ALLOW_THREADS self->runtime->shutdown();	×
307	Py_END_ALLOW_THREADS	×
308	}
309	Py_RETURN_NONE;	×
310	}
311
312	static PyObject Runtime_get_threads(RuntimeObject self, void *closure) {	14✔
313	if (!self->runtime) {	14!
314	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
315	return NULL;	×
316	}
317	return PyLong_FromSize_t(self->runtime->threads());	14✔
318	}	7✔
319
320	static PyObject get_default_runtime_py(PyObject Py_UNUSED(module),	2✔
321	PyObject *Py_UNUSED(ignored)) {
322	dftracer::utils::Runtime *rt = get_default_runtime();	2✔
323	if (!rt) {	2✔
324	PyErr_SetString(PyExc_RuntimeError, "Failed to create default runtime");	×
325	return NULL;	×
326	}
327
328	RuntimeObject obj = (RuntimeObject )RuntimeType.tp_alloc(&RuntimeType, 0);	2✔
329	if (!obj) return NULL;	2✔
330
331	new (&obj->runtime)	2✔
332	std::shared_ptr<dftracer::utils::Runtime>(g_default_runtime);	2✔
333	return (PyObject *)obj;	2✔
334	}	1✔
335
336	// Return the current default runtime without creating one. Unlike
337	// get_default_runtime, this never materializes a full-machine-sized runtime as
338	// a side effect, so callers that only want to save/restore the default (e.g. a
339	// Dask worker plugin) do not each spin up an unused hardware_concurrency-thread
340	// runtime.
NEW 341	static PyObject peek_default_runtime_py(PyObject Py_UNUSED(module),	×
342	PyObject *Py_UNUSED(ignored)) {
NEW 343	if (!g_default_runtime) Py_RETURN_NONE;	×
344
NEW 345	RuntimeObject obj = (RuntimeObject )RuntimeType.tp_alloc(&RuntimeType, 0);	×
NEW 346	if (!obj) return NULL;	×
347
NEW 348	new (&obj->runtime)	×
NEW 349	std::shared_ptr<dftracer::utils::Runtime>(g_default_runtime);	×
NEW 350	return (PyObject *)obj;	×
351	}
352
353	static PyObject set_default_runtime_py(PyObject Py_UNUSED(module),	4✔
354	PyObject *args) {
355	PyObject *arg;
356	if (!PyArg_ParseTuple(args, "O", &arg)) return NULL;	4!
357
358	if (arg == Py_None) {	4✔
359	g_default_runtime.reset();	×
360	Py_RETURN_NONE;	×
361	}
362
363	if (!PyObject_TypeCheck(arg, &RuntimeType)) {	4!
364	PyErr_SetString(PyExc_TypeError, "Expected Runtime or None");	×
365	return NULL;	×
366	}
367
368	g_default_runtime = ((RuntimeObject *)arg)->runtime;	4✔
369	Py_RETURN_NONE;	4✔
370	}	2✔
371
372	static PyMethodDef Runtime_methods[] = {
373	{"shutdown", (PyCFunction)Runtime_shutdown, METH_NOARGS,
374	"shutdown()\n"
375	"--\n"
376	"\n"
377	"Shut down the runtime.\n"},
378	{"get_progress", (PyCFunction)Runtime_get_progress, METH_NOARGS,
379	"Return progress dict with keys: total, completed, running,\n"
380	"queued, failed."},
381	{"is_responsive", (PyCFunction)Runtime_is_responsive, METH_NOARGS,
382	"Return True if the runtime is making progress."},
383	{"set_timeout", (PyCFunction)Runtime_set_timeout,
384	METH_VARARGS \| METH_KEYWORDS,
385	"Set global timeout in milliseconds.\n"
386	"\n"
387	"Args:\n"
388	" global_ms (int): Timeout in milliseconds (0 = no timeout).\n"},
389	{"set_default_task_timeout", (PyCFunction)Runtime_set_default_task_timeout,
390	METH_VARARGS \| METH_KEYWORDS,
391	"Set default per-task timeout in milliseconds.\n"
392	"\n"
393	"Args:\n"
394	" ms (int): Timeout in milliseconds (0 = no timeout).\n"},
395	{"wait_all", (PyCFunction)Runtime_wait_all, METH_NOARGS,
396	"Wait for all outstanding submitted tasks to complete."},
397	{"__enter__", (PyCFunction)Runtime_enter, METH_NOARGS,
398	"Enter context manager."},
399	{"__exit__", (PyCFunction)Runtime_exit, METH_VARARGS,
400	"Exit context manager (calls shutdown)."},
401	{NULL}};
402
403	static PyObject Runtime_get_io_threads(RuntimeObject self, void *closure) {	×
404	if (!self->runtime) {	×
405	PyErr_SetString(PyExc_RuntimeError, "Runtime not initialized");	×
406	return NULL;	×
407	}
408	return PyLong_FromSize_t(self->runtime->io_threads());	×
409	}
410
411	static PyGetSetDef Runtime_getsetters[] = {
412	{"threads", (getter)Runtime_get_threads, NULL, "Number of worker threads",
413	NULL},
414	{"io_threads", (getter)Runtime_get_io_threads, NULL,
415	"Number of I/O threads", NULL},
416	{NULL}};
417
418	PyTypeObject RuntimeType = {
419	PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.Runtime",
420	sizeof(RuntimeObject), /* tp_basicsize */
421	0, /* tp_itemsize */
422	(destructor)Runtime_dealloc, /* tp_dealloc */
423	0, /* tp_vectorcall_offset */
424	0, /* tp_getattr */
425	0, /* tp_setattr */
426	0, /* tp_as_async */
427	0, /* tp_repr */
428	0, /* tp_as_number */
429	0, /* tp_as_sequence */
430	0, /* tp_as_mapping */
431	0, /* tp_hash */
432	0, /* tp_call */
433	0, /* tp_str */
434	0, /* tp_getattro */
435	0, /* tp_setattro */
436	0, /* tp_as_buffer */
437	Py_TPFLAGS_DEFAULT \| Py_TPFLAGS_BASETYPE, /* tp_flags */
438	"Runtime(threads: int = 0, io_threads: int = 0)\n"
439	"--\n"
440	"\n"
441	"Coroutine runtime backed by a thread pool.\n"
442	"\n"
443	"Args:\n"
444	" threads (int): Number of worker threads. 0 (default) uses\n"
445	" the hardware concurrency.\n"
446	" io_threads (int): Number of I/O threads. 0 (default) uses\n"
447	" the hardware concurrency.\n", /* tp_doc */
448	0, /* tp_traverse */
449	0, /* tp_clear */
450	0, /* tp_richcompare */
451	0, /* tp_weaklistoffset */
452	0, /* tp_iter */
453	0, /* tp_iternext */
454	Runtime_methods, /* tp_methods */
455	0, /* tp_members */
456	Runtime_getsetters, /* tp_getset */
457	0, /* tp_base */
458	0, /* tp_dict */
459	0, /* tp_descr_get */
460	0, /* tp_descr_set */
461	0, /* tp_dictoffset */
462	(initproc)Runtime_init, /* tp_init */
463	0, /* tp_alloc */
464	Runtime_new, /* tp_new */
465	};
466
467	// Module-level function table (registered via PyModule_AddFunctions or
468	// appended to the module's method table in init_runtime).
469	static PyMethodDef runtime_module_methods[] = {
470	{"get_default_runtime", get_default_runtime_py, METH_NOARGS,
471	"Return the module-level default Runtime (lazy-created)."},
472	{"peek_default_runtime", peek_default_runtime_py, METH_NOARGS,
473	"Return the current default Runtime, or None if none exists yet "
474	"(never creates one)."},
475	{"set_default_runtime", set_default_runtime_py, METH_VARARGS,
476	"Replace the module-level default Runtime (pass None to clear).\n"
477	"\n"
478	"Args:\n"
479	" runtime (Runtime or None): New default runtime.\n"},
480	{NULL}};
481
482	// Runs during interpreter finalization. Stop the default runtime's worker and
483	// I/O threads while the process is still healthy, then tell RocksDB we are
484	// exiting so cached DB handles skip closing every open SST on teardown. A read
485	// scan can leave hundreds of SSTs open (more so on a networked filesystem),
486	// and without this the process can appear to hang after the work is done.
487	static void dftracer_utils_atexit_cleanup() {	2✔
488	if (g_default_runtime) {	2✔
489	g_default_runtime->shutdown();	2✔
490	}	1✔
491	dftracer::utils::rocksdb::mark_process_exiting_for_rocksdb();	2✔
492	}	2✔
493
494	int init_runtime(PyObject *m) {	2✔
495	if (register_type(m, &RuntimeType, "Runtime") < 0) return -1;	2✔
496
497	Py_AtExit(dftracer_utils_atexit_cleanup);	2✔
498
499	for (PyMethodDef *def = runtime_module_methods; def->ml_name; ++def) {	8✔
500	PyObject *fn = PyCFunction_New(def, NULL);	6✔
501	if (!fn) return -1;	6!
502	if (PyModule_AddObject(m, def->ml_name, fn) < 0) {	6!
503	Py_DECREF(fn);
504	return -1;	×
505	}
506	}	3✔
507
508	return 0;	2✔
509	}	1✔

llnl / dftracer-utils / 30002033718

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

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