36502607218

Committed 29 Sep 2026 12:20AM UTC coverage: 57.996% (+5.4%) from 52.58%

Build # 36502607218

Build Type

push

github

Committed by

rayandrew

Commit Message

Merge branch 'feat/view-as-lazyframe' into 'develop'

feat: view as lazyframe

See merge request dftracer/dftracer-utils!25

Coverage Stats

74077 of 162805 branches covered (45.5%)

Branch coverage included in aggregate %.

3652 of 3978 new or added lines in 49 files covered. (91.8%)

3464 existing lines in 85 files now uncovered.

65902 of 78556 relevant lines covered (83.89%)

179668.29 hits per line

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

48.6

/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_method.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>

dftracer::utils::Runtime *dftracer::utils::python::get_default_runtime() {
    return &dftracer::utils::default_runtime();
}

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

static PyObject *Runtime_new(PyTypeObject *type, PyObject *, PyObject *) {
    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",
                                     const_cast<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) {
        dftracer::utils::python::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;
    }

    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);

    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);

    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",
                                     const_cast<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",
                                     const_cast<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 *) {
    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 *) {
    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 =
        dftracer::utils::python::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>(
        dftracer::utils::default_runtime_shared());
    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)) {
    auto rt = dftracer::utils::peek_default_runtime();
    if (!rt) Py_RETURN_NONE;

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

    new (&obj->runtime) std::shared_ptr<dftracer::utils::Runtime>(rt);
    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) {
        dftracer::utils::set_default_runtime(nullptr);
        Py_RETURN_NONE;
    }

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

    dftracer::utils::set_default_runtime(((RuntimeObject *)arg)->runtime);
    Py_RETURN_NONE;
}

static PyMethodDef Runtime_methods[] = {
    {"shutdown", DFTU_PYCFUNCTION(Runtime_shutdown), METH_NOARGS,
     "shutdown()\n"
     "--\n"
     "\n"
     "Shut down the runtime.\n"},
    {"get_progress", DFTU_PYCFUNCTION(Runtime_get_progress), METH_NOARGS,
     "Return progress dict with keys: total, completed, running,\n"
     "queued, failed."},
    {"is_responsive", DFTU_PYCFUNCTION(Runtime_is_responsive), METH_NOARGS,
     "Return True if the runtime is making progress."},
    {"set_timeout", DFTU_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",
     DFTU_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", DFTU_PYCFUNCTION(Runtime_wait_all), METH_NOARGS,
     "Wait for all outstanding submitted tasks to complete."},
    {"__enter__", DFTU_PYCFUNCTION(Runtime_enter), METH_NOARGS,
     "Enter context manager."},
    {"__exit__", DFTU_PYCFUNCTION(Runtime_exit), METH_VARARGS,
     "Exit context manager (calls shutdown)."},
    {NULL}};

static PyObject *Runtime_get_io_threads(RuntimeObject *self, void *) {
    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 (auto rt = dftracer::utils::peek_default_runtime()) {
        rt->shutdown();
    }
    dftracer::utils::rocksdb::mark_process_exiting_for_rocksdb();
}

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

llnl-asr / dftracer-utils / 36502607218

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