pipes_test.py

#!/usr/bin/env python
"""
Pipes experiment.
Refs: https://gist.github.com/mottosso/dd59f815025e4d76f453f3555e39558e
"""

import os 
import sys
import time
import argparse

KB = 1 << 10
MB = 1 << 20

BUFSIZE = 64 * KB

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--size", help="size in kb", default=32, type=int)
    parser.add_argument("-i", "--iterations", default=100, type=int)
    opt = parser.parse_args()

    parentR, childW = os.pipe()
    childR, parentW = os.pipe()
    
    pid = os.fork()

    if not pid:
        os.close(parentR)
        os.close(parentW)
        child(childR, childW)
        return

    # In the parent process
    os.close(childR)
    os.close(childW)
    r = os.fdopen(parentR, 'r', BUFSIZE) # unused for now
    w = os.fdopen(parentW, 'w', BUFSIZE)

    size = int(opt.size * KB)
    data = 'x' * size 
    size = sys.getsizeof(data)

    sent = 0
    totalTime = 0

    for i in range(opt.iterations):

        t0 = time.clock()
        w.write(data)
        t1 = time.clock()

        sent += size
        totalTime += (t1 - t0)

        if i % 100 == 0:
            sys.stderr.write("%d/%d..." % (i, opt.iterations))

    w.flush()

    sys.stderr.write("done\n")

    rate = float(sent) / MB / totalTime

    sys.stderr.write('Total Size: %dKB\n' % (sent / KB))
    sys.stderr.write(' Send Size: %dKB\n' % (size / KB))
    sys.stderr.write('  Buf Size: %dKB\n' % (BUFSIZE / KB))
    sys.stderr.write('     Write: %0.6f MB/s\n' % (rate))

    r.close()
    w.close()
    os.waitpid(pid, 0)


def child(rd, wd):
    r = os.fdopen(rd, 'r', BUFSIZE)
    w = os.fdopen(wd, 'w', BUFSIZE) # unused for now

    total = 0
    totalTime = 0
    maxRead = 32 * KB

    while True:
        t0 = time.clock()
        data = r.read(maxRead)
        t1 = time.clock()
        if not data:
            break
        total += sys.getsizeof(data)
        totalTime += (t1 - t0)

    r.close()
    w.close()

    rate = 0
    if totalTime > 0:
        rate = float(total) / MB / totalTime
    sys.stderr.write('      Read: %0.6f MB/s (total: %dKB)\n' % (rate, total / KB))


if __name__ == '__main__':
    main()