Benchmarking spectral/spatial profile reads using S3 signed URLs

test_s3.ts

import DataView from "https://deno.land/x/[email protected]/_DataView.js";

async function getPartialFile(url: string, byteOffset: number = 0, byteLength: number = 4096) {
    try {
        const headers = new Headers();
        headers.append("Range", `bytes=${byteOffset}-${byteOffset+byteLength - 1}`);
        const res = await fetch(url, { mode: "no-cors", headers });
        return await res.blob();
    } catch (err) {
        console.log(err);
    }
}

async function readFitsHeader(url: string, maxHeaderCount = 8192): Promise<{headers: string[], dataOffset: number}> {
    const headers = new Array<string>();
    let offset = 0;
    const lineSize = 80;
    const chunkSize = 2880;
    const linesPerChunk = chunkSize / lineSize;
    while (headers.length < maxHeaderCount) {
        const data = await getPartialFile(url, offset, chunkSize);
        const text = await data?.text();
        offset += chunkSize;
        if (text?.length === chunkSize) {
            for (let i = 0; i < linesPerChunk; i++) {
                const entry = text.slice(i * lineSize, (i + 1) * lineSize);
                headers.push(entry);
                if (entry.startsWith("END")) {
                    return {headers, dataOffset: offset};
                }
            }
        }
    }

    return {headers, dataOffset: -1};
}

type FitsSize = {
    bitPix: number;
    bytesPerPixel: number;
    dims: number[];
}
function parseFitsHeader(headers: string[]): FitsSize | undefined {
    const bitPixRegex = /^BITPIX  =\s*(-?\d+)/;
    const naxisRegex = /^NAXIS   =\s*(\d+)/;
    let bitPix;
    let naxis;

    for (const header of headers) {
        const result = header.match(bitPixRegex);
        if (result?.length === 2) {
            bitPix = Number.parseInt(result[1]);
            break;
        }
    }

    for (const header of headers) {
        const result = header.match(naxisRegex);
        if (result?.length === 2) {
            naxis = Number.parseInt(result[1]);
            break;
        }
    }

    if (bitPix && naxis) {
        const dimsRegex = /^NAXIS(\d+)\s*=\s*(\d+)/;
        const dims = new Array<number>(naxis);
        for (const header of headers) {
            const result = header.match(dimsRegex);
            if (result?.length === 3) {
                const index = Number.parseInt(result[1]) - 1;
                const size = Number.parseInt(result[2]);
                if (size && index >= 0 && index < naxis) {
                    dims[index] = size;
                }
            }
        }
        return {dims, bitPix, bytesPerPixel: Math.abs(bitPix) / 8};
    }

    return undefined;
}

async function getXProfile(url: string, offset: number, fitsSize: FitsSize, y: number, channel: number): Promise<Float32Array | Float64Array | Int32Array | undefined> {
    // Only support Float32 images with 2+ dimensions
    if (fitsSize.bitPix !== -32) {
        return undefined;
    }
    if (fitsSize.dims.length < 2) {
        return undefined;
    }

    const lineLength = fitsSize.bytesPerPixel * fitsSize.dims[0];
    const sliceLength = lineLength * fitsSize.dims[1];
    if (y) {
        offset += y * lineLength;
    }
    if (channel) {
        offset += channel * sliceLength;
    }
    const res = await getPartialFile(url, offset, lineLength);
    const buffer = await res?.arrayBuffer();
    const data = new Float32Array(fitsSize.dims[0]);
    const view = new DataView(buffer);
    for (let i = 0; i < data.length; i++) {
        data[i] = view.getFloat32(i * 4, false);
    }
    return data;
}

async function fillPixel(url: string, offset: number, data: Float32Array, index: number) {
    const res = await getPartialFile(url, offset, 4);
    const buffer = await res?.arrayBuffer();
    const view = new DataView(buffer);
    data[index] = view.getFloat32(0, false);
    return;
}

async function getYProfile(url: string, offset: number, fitsSize: FitsSize, x: number, channel: number): Promise<Float32Array | Float64Array | Int32Array | undefined> {
    // Only support Float32 images with 2+ dimensions
    if (fitsSize.bitPix !== -32) {
        return undefined;
    }
    if (fitsSize.dims.length < 2) {
        return undefined;
    }

    const lineLength = fitsSize.bytesPerPixel * fitsSize.dims[0];
    const sliceLength = lineLength * fitsSize.dims[1];
    if (x) {
        offset += x * fitsSize.bytesPerPixel;
    }
    if (channel) {
        offset += channel * sliceLength;
    }

    const data = new Float32Array(fitsSize.dims[1]);
    const promises = [];
    for (let y = 0; y < fitsSize.dims[1]; y++) {
        promises.push(fillPixel(url, offset, data, y));
        offset += lineLength;
    }
    await Promise.all(promises);
    return data;
}

async function getZProfile(url: string, offset: number, fitsSize: FitsSize, x: number, y: number): Promise<Float32Array | Float64Array | Int32Array | undefined> {
    // Only support Float32 images with 3+ dimensions
    if (fitsSize.bitPix !== -32) {
        return undefined;
    }
    if (fitsSize.dims.length < 3) {
        return undefined;
    }

    const lineLength = fitsSize.bytesPerPixel * fitsSize.dims[0];
    const sliceLength = lineLength * fitsSize.dims[1];
    if (x) {
        offset += x * fitsSize.bytesPerPixel;
    }
    if (y) {
        offset += y * lineLength;
    }

    const data = new Float32Array(fitsSize.dims[2]);
    const promises = [];
    for (let z = 0; z < fitsSize.dims[2]; z++) {
        promises.push(fillPixel(url, offset, data, z));
        offset += sliceLength;
    }
    await Promise.all(promises);
    return data;
}


const signedUrl =
    "https://dashboard2.ilifu.ac.za:6780/test/t10_firstpass_contsub.fits?AWSAccessKeyId=f9d7174d82e44b50bcd7cb7304213dae&Expires=1666229443&Signature=HeJD0o5Dq3pUcczje5hjtVgIYhY%3D";
const fileSize = 10.1 * 1024**3;
const numTests = 8192;
const numStreams = 4096;
const chunkSize = 4096;
const queueDepth = numTests / numStreams;
let testsCompleted = 0;

async function getPartialFileStream() {
    for (let i = 0; i < queueDepth; i++) {
        const startOffset = Math.floor(Math.random() * (fileSize - chunkSize));
        await getPartialFile(signedUrl, startOffset, chunkSize);
        testsCompleted++;
    }
}

async function benchIops() {
    console.log("Warming up");
    await getPartialFile(signedUrl);
    console.log("Starting benchmark");
    const tStart = performance.now();
    const promises = [];

    for (let i = 0; i < numStreams; i++) {
        promises.push(getPartialFileStream());
    }
    await Promise.all(promises);

    const tEnd = performance.now();
    const dt = tEnd - tStart;
    const iops = numTests / dt * 1000;
    const throughput = iops * chunkSize / (1024**2);
    console.log(`${testsCompleted} requests completed in ${dt.toFixed(1)} ms (${iops.toFixed(0)} IOPS, ${throughput.toFixed(1)} MB/s)`);
}

async function benchProfiles() {
    const res = await readFitsHeader(signedUrl);
    console.log(`Data starts at offset ${res.dataOffset} bytes`);
    const fitsSize = parseFitsHeader(res.headers);
    if (fitsSize && fitsSize.dims.length >= 3) {
        const lineLength = fitsSize.bytesPerPixel * fitsSize.dims[0];
        const sliceLength = lineLength * fitsSize.dims[1];
        const cubeLength = sliceLength * fitsSize.dims[2];
        console.log(fitsSize, {lineLength, sliceLength, cubeLength});

        const numTests = 25;
        let sumX = 0;
        let sumX2 = 0;
        for (let i = 0; i < numTests; i++) {
            const pixelLocation = {x: Math.floor(Math.random() * fitsSize.dims[0]), y: Math.floor(Math.random() * fitsSize.dims[1]), z: Math.floor(Math.random() * fitsSize.dims[2])};

            const tStart = performance.now();
            const profile = await getYProfile(signedUrl, res.dataOffset, fitsSize, pixelLocation.x, pixelLocation.z);
            //const profile = await getZProfile(signedUrl, res.dataOffset, fitsSize, pixelLocation.x, pixelLocation.y);
            const tEnd = performance.now();
            const dt = tEnd - tStart;
            sumX += dt;
            sumX2 += dt*dt;
            const numPixels = profile?.length ?? 0;
            const iops = numPixels / dt * 1000;
            console.log(`Spatial profile for (x,z)=(${pixelLocation.x}, ${pixelLocation.z}) fetched in ${dt.toFixed(1)} ms (${numPixels} pixels @ ${iops.toFixed(0)} IOPS)`);
            //console.log(`Spectral profile for (x,y)=(${pixelLocation.x}, ${pixelLocation.y}) fetched in ${dt.toFixed(1)} ms (${numPixels} pixels @ ${iops.toFixed(0)} IOPS)`);
        }

        const meanTime = sumX / numTests;
        const stdDev = Math.sqrt(sumX2 / numTests - meanTime * meanTime);
        console.log(`${meanTime.toFixed(1)} +- ${stdDev.toFixed(1)} ms`);
    }
}

benchProfiles();
//benchIops();