Introduction

Cython is a superset of Python and allows you to statically type variables, resulting in significant performance boosts, especially in numerical for loops. Cython also allows you to wrap C libraries and create Python C extensions. This means you don't have to know much C, if any, to get access to C functions and expose the results to Python.

In this guide, we will use Cython to wrap some C functions in Slurm's sdiag program and make the statistics available to Python.


Prerequisites

This guide assumes you are using Slurm, an HPC workload manager. This guide uses Slurm 15.08.8, Cython 0.23.4 and Python 2.7.5, all on CentOS 7.2.

This should also work in Python 3.5.1.


Looking at sdiag Output

sdiag shows diagnostics for Slurm job scheduling and other information about the Slurm controller. Run the sdiag command and you should see similar output:

*******************************************************
sdiag output at Sun Mar 27 14:21:29 2016
Data since      Sun Mar 27 14:09:09 2016
*******************************************************
Server thread count: 3
Agent queue size:    0

Jobs submitted: 0
Jobs started:   0
Jobs completed: 0
Jobs canceled:  0
Jobs failed:    0

Main schedule statistics (microseconds):
    Last cycle:   20
    Max cycle:    45
    Total cycles: 13
    Mean cycle:   26
    Mean depth cycle:  0
    Cycles per minute: 1
    Last queue length: 0

Backfilling stats
    Total backfilled jobs (since last slurm start): 0
    Total backfilled jobs (since last stats cycle start): 0
    Total cycles: 0
    Last cycle when: Thu Jan  1 00:00:00 1970
    Last cycle: 0
    Max cycle:  0
    Last depth cycle: 0
    Last depth cycle (try sched): 0
    Last queue length: 0
    [...]

We want to use Cython to wrap the necessary Slurm API functions to get access to these statistics via Python.


Looking at the C API

Normally, we would start with Slurm's [API page](http://slurm.schedmd.com/api.html], which lists various functions to get direct access to Slurm data structures. This page, however, does not have any functions related to statistics. Next, we could check the man pages, which has details on the APIs and some examples. Again, no man pages exist at the moment for any statistics APIs. Therefore, we have to dig into the source code for this example.

Go to Slurm's GitHub page and set the tag to your running version of Slurm. Go to the slurm/src/sdiag directory. The sdiag.c file is the source code for the sdiag program. We need to replicate a portion of this program in order to get access to the statistics.

Looking at the program, we need to wrap the slurm_get_statistics, slurm_reset_statistics, and the slurm_free_stats_response_msg functions.


Organizing the Cython code

We need two files for this project: an implementation (.pyx) file and a definition file (.pxd).

We put the C function and variable declarations in the in the definition file. The Cython code for wrapping the functions goes in the implementation file.


Slurm Header Files

We need to declare the C functions and variables, and we get those from the Slurm header files. The header file is in the include directory of your Slurm installation path, /path/to/slurm/include/slurm/slurm.h

Open the slurm.h header file and search for slurm_get_statistics, one of the functions we want to wrap. There, we see its declaration and signature:

extern int slurm_get_statistics PARAMS((stats_info_response_msg_t **buf,
                    stats_info_request_msg_t *req));

The signature for slurm_get_statistics has two structs: stats_info_response_msg_t and stats_info_request_msg_t. We will also have to declare both structs in the definition file.


Wrapping C Code

Let's start by declaring the structs. The C code for the stats_info_request_msg struct looks like:

typedef struct stats_info_request_msg {
    uint16_t command_id;
} stats_info_request_msg_t;

Create a file called sdiag.pxd and start with the following:

from libc.stdint cimport uint16_t

cdef extern from "slurm/slurm.h" nogil:
    ctypedef struct stats_info_request_msg_t:
        uint16_t command_id

The cimport statement is used to access the uint16_t type declared in libc's stdint.h. The cdef extern from syntax tells Cython where the following declarations are located. The nogil annotation declares that this block of code can be used without Python's GIL.

Cython's declaration for the struct is almost cut and paste from the header file. The typedef in C becomes a ctypedef in Cython.

For more information on struct and enum declaration, see Cython's documentation.

Next, add the stats_info_response_msg_t struct to the .pxd file:

    ctypedef struct stats_info_response_msg_t:
        uint32_t parts_packed
        time_t req_time
        time_t req_time_start
        uint32_t server_thread_count
        uint32_t agent_queue_size

        uint32_t schedule_cycle_max
        uint32_t schedule_cycle_last
        uint32_t schedule_cycle_sum
        uint32_t schedule_cycle_counter
        uint32_t schedule_cycle_depth
        uint32_t schedule_queue_len
        [...]

Note some of the types in this struct. We need to import these as well into the .pxd file:

from libc.stdint cimport uint16_t, uint32_t, uint64_t

cdef extern from "sys/types.h" nogil:
    ctypedef long time_t

Now, we can declare the two functions:

    int slurm_get_statistics(stats_info_response_msg_t **buf,
                             stats_info_request_msg_t *req)
    int slurm_reset_statistics(stats_info_request_msg_t *req)

There are some C macros that we also need to declare:

    int STAT_COMMAND_RESET
    int STAT_COMMAND_GET

We now need to declare some functions and macros from slurm_errno.h:

cdef extern from "slurm/slurm_errno.h" nogil:
    int SLURM_SUCCESS

    char *slurm_strerror(int errnum)
    int slurm_get_errno()

The above declarations are for checking return codes and printing errors.

There's one more declaration we need, but it is not in the slurm.h header file. In this case, since the function is externalized, we can declare it directly:

cdef extern void slurm_free_stats_response_msg(stats_info_response_msg_t *msg)

This API function is for freeing memory allocated by the stats response buffer.


Writing the Implementation File

The implementation file (.pyx) is where the Cython code goes. This is where we need to look at the sdiag.c file and recreate the needed portions in Cython.

We will expose two functions to Python: get_stats() and reset_stats(). Therefore, these need to use the cpdef function declaration, since cpdef functions have access to both C and Python.

cpdef dict get_stats():
    pass

cpdef int reset_stats():
    pass

In the get_stats function, we need to statically type some variables that will get used in the function. This is done with cdef:

cpdef dict get_stats():
    cdef:
        int rc
        uint32_t i
        dict rpc_type_stats
        dict rpc_user_stats
        dict stat_dict
        stats_info_request_msg_t req
        stats_info_response_msg_t *buf

Next, follow through sdiag.c and write the necessary Cython to call the slurm_get_statistics function, populate the buffer, and copy the buffer into a dictionary.

    [...]
    req.command_id = STAT_COMMAND_GET
    rc = slurm_get_statistics(&buf, <stats_info_request_msg_t*>&req)

    if rc == SLURM_SUCCESS:
        stat_dict = {}
        stat_dict["parts_packed"] = buf.parts_packed
        stat_dict["req_time"] = buf.req_time
        stat_dict["req_time_start"] = buf.req_time_start
        stat_dict["server_thread_count"] = buf.server_thread_count
        stat_dict["agent_queue_size"] = buf.agent_queue_size
        [...]

Part of this block of code requires the rpc_num2string function. Unfortunately, it is not externalized and we cannot simply declare it in the .pxd file and use it. We have to recreate it in the implementation file as a separate function. Since this function does not need to be exposed to Python, we can use a cdef function:

cdef rpc_num2string(uint16_t opcode):
    num2string = {
        1001: "REQUEST_NODE_REGISTRATION_STATUS",
        1002: "MESSAGE_NODE_REGISTRATION_STATUS",
        1003: "REQUEST_RECONFIGURE",
        1004: "RESPONSE_RECONFIGURE",
        1005: "REQUEST_SHUTDOWN",
        1006: "REQUEST_SHUTDOWN_IMMEDIATE",
    [...]

At the end of the get_stats function, we need to free the response message buffer and return the dictionary back to the caller:

    slurm_free_stats_response_msg(buf)
    buf = NULL
    return stat_dict

Pay attention to the types required in the implementation file, they need to be imported, similarly to the definition file:

from libc.stdint cimport uint16_t, uint32_t
from pwd import getpwuid

Lastly, the functions we created in the implementation file can be declared in the definition file, using the same signature as in the implementation file:

cpdef dict get_stats()
cpdef int reset_stats()
cdef rpc_num2string(uint16_t opcode)

The end result should look something like this gist.


Compiling with Cython

Once we have the .pyx file ready, we can use the Cython compiler to convert it to a C file.

cython sdiag.pyx

This process will not produce any output to stdout or stderr if there are no errors or warnings. This is the beauty of Cython. It did all the hard work of compiling the Cython source to C code. The resulting file will be sdiag.c. Have a look and see what Cython had done for us.


Creating the Python Extension Module

The resulting C output file can now be converted to a Python extension module. The standard way of creating an extension module is through the use of distutils in a setup.py file.

We can, as an exercise, compile by the extension module by hand to get an understanding of how the module is generated with the compiler and linker.

Our next step is to compile sdiag.c into an object file. Use Python's python-config to get the appropriate CFLAGS to pass to the compiler:

gcc -c sdiag.c $(python-config --cflags) -fPIC

This step produces the sdiag.o object file.

Now, compile sdiag.o into a shared library. We will need to add the linker flags that point to Slurm's libslurm.so:

gcc sdiag.o -o sdiag.so -shared $(python-config --ldflags) \
  -L/usr/local/slurm/15.08.8/lib -lslurm


Testing the module

Create a directory for the module:

mkdir ${HOME}/local
cd ${HOME}/local
mkdir sdiag

Copy the sdiag.so module to the ${HOME}/local/sdiag directory. Next, create an __init__.py file. Since Slurm makes heavy use of dlopen() calls, we need to set the flags in __init__.py:

from __future__ import absolute_import

import sys
import ctypes

sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)

from .sdiag import *                                         

Now, we can test the module. Make sure you are in the ${HOME}/local directory before running the Python interpreter. Otherwise, be sure to adjust your PYTHONPATH.

>>> import sdiag
>>> from pprint import pprint
>>> pprint(sdiag.get_stats())
{u'agent_queue_size': 0,
 u'bf_active': 0,
 u'bf_backfilled_jobs': 0,
 u'bf_cycle_counter': 0,
 u'bf_cycle_last': 0,
 u'bf_cycle_max': 0,
 u'bf_cycle_sum': 0,
 u'bf_depth_sum': 0,
 u'bf_depth_try_sum': 0,
 u'bf_last_backfilled_jobs': 0,
 [...]
}

The get_stats() function returns a dictionary of key value pairs of information produced by the sdiag command.


Conclusion

With Cython, we are able to easily wrap C code and make it available to Python. We now have tighter integration with Slurm and don't have to run Slurm commands and parse the output. For more examples of using Cython to wrap the Slurm C API, have a look at the PySlurm project on GitHub.


Comments

comments powered by Disqus