netron/test/worker.js

import * as base from '../source/base.js';
import * as fs from 'fs/promises';
import * as mock from './mock.js';
import * as node from '../source/node.js';
import * as path from 'path';
import * as process from 'process';
import * as python from '../source/python.js';
import * as tar from '../source/tar.js';
import * as url from 'url';
import * as view from '../source/view.js';
import * as worker_threads from 'worker_threads';
import * as zip from '../source/zip.js';

const access = async (path) => {
    try {
        await fs.access(path);
        return true;
    } catch {
        return false;
    }
};

const dirname = (...args) => {
    const file = url.fileURLToPath(import.meta.url);
    const dir = path.dirname(file);
    return path.join(dir, ...args);
};

const decompress = (buffer) => {
    let archive = zip.Archive.open(buffer, 'gzip');
    if (archive && archive.entries.size === 1) {
        const stream = archive.entries.values().next().value;
        buffer = stream.peek();
    }
    const formats = [zip, tar];
    for (const module of formats) {
        archive = module.Archive.open(buffer);
        if (archive) {
            break;
        }
    }
    return archive;
};

export class Target {

    constructor(item) {
        Object.assign(this, item);
        this.events = {};
        this.tags = new Set(this.tags);
        this.folder = item.type ? path.normalize(dirname('..', 'third_party' , 'test', item.type)) : process.cwd();
        this.assert = !this.assert || Array.isArray(this.assert) ? this.assert : [this.assert];
        this.serial = false;
    }

    on(event, callback) {
        this.events[event] = this.events[event] || [];
        this.events[event].push(callback);
    }

    emit(event, data) {
        if (this.events && this.events[event]) {
            for (const callback of this.events[event]) {
                callback(this, data);
            }
        }
    }

    status(message) {
        this.emit('status', message);
    }

    async execute() {
        if (this.measures) {
            this.measures.set('name', this.name);
        }
        await zip.Archive.import();
        const environment = { zoom: 'none', serial: this.serial };
        this.host = await new mock.Host(environment);
        this.view = new view.View(this.host);
        this.view.options.attributes = true;
        this.view.options.initializers = true;
        const time = async (method) => {
            const start = process.hrtime.bigint();
            let err = null;
            try {
                await method.call(this);
            } catch (error) {
                err = error;
            }
            const duration = Number(process.hrtime.bigint() - start) / 1e9;
            if (this.measures) {
                this.measures.set(method.name, duration);
            }
            if (err) {
                throw err;
            }
        };
        this.status({ name: 'name', target: this.name });
        const errors = [];
        try {
            await time(this.download);
            await time(this.load);
            await time(this.validate);
            if (!this.tags.has('skip-render')) {
                await time(this.render);
            }
        } catch (error) {
            errors.push(error);
        }
        errors.push(...this.host.errors);
        if (errors.length === 0 && this.error) {
            throw new Error('Expected error.');
        }
        if (errors.length > 0 && (!this.error || errors.map((error) => error.message).join('\n') !== this.error)) {
            throw errors[0];
        }
        this.view.dispose();
    }

    async request(url, init) {
        const response = await fetch(url, init);
        if (!response.ok) {
            throw new Error(response.status.toString());
        }
        if (response.body) {
            const reader = response.body.getReader();
            const length = response.headers.has('Content-Length') ? parseInt(response.headers.get('Content-Length'), 10) : -1;
            let position = 0;
            /* eslint-disable consistent-this */
            const target = this;
            /* eslint-enable consistent-this */
            const stream = new ReadableStream({
                async start(controller) {
                    const read = async () => {
                        try {
                            const result = await reader.read();
                            if (result.done) {
                                target.status({ name: 'download' });
                                controller.close();
                            } else {
                                position += result.value.length;
                                if (length >= 0) {
                                    const percent = position / length;
                                    target.status({ name: 'download', target: url, percent });
                                } else {
                                    target.status({ name: 'download', target: url, position });
                                }
                                controller.enqueue(result.value);
                                return await read();
                            }
                        } catch (error) {
                            controller.error(error);
                            throw error;
                        }

                        return null;
                    };
                    return read();
                }
            });
            return new Response(stream, {
                status: response.status,
                statusText: response.statusText,
                headers: response.headers
            });
        }
        return response;
    }
    async download(targets, sources) {
        targets = targets || Array.from(this.targets);
        sources = sources || this.source;
        const files = targets.map((file) => path.resolve(this.folder, file));
        const exists = await Promise.all(files.map((file) => access(file)));
        if (exists.every((value) => value)) {
            return;
        }
        if (!sources) {
            throw new Error('Download source not specified.');
        }
        let source = '';
        let sourceFiles = [];
        const match = sources.match(/^(.*?)\[(.*?)\](.*)$/);
        if (match) {
            [, source, sourceFiles, sources] = match;
            sourceFiles = sourceFiles.split(',').map((file) => file.trim());
            sources = sources && sources.startsWith(',') ? sources.substring(1).trim() : '';
        } else {
            const comma = sources.indexOf(',');
            if (comma === -1) {
                source = sources;
                sources = '';
            } else {
                source = sources.substring(0, comma);
                sources = sources.substring(comma + 1);
            }
        }
        await Promise.all(targets.map((target) => {
            const dir = path.dirname(`${this.folder}/${target}`);
            return fs.mkdir(dir, { recursive: true });
        }));
        const response = await this.request(source);
        const buffer = await response.arrayBuffer();
        const data = new Uint8Array(buffer);
        if (sourceFiles.length > 0) {
            this.status({ name: 'decompress' });
            const archive = decompress(data);
            for (const name of sourceFiles) {
                this.status({ name: 'write', target: name });
                if (name === '.') {
                    const target = targets.shift();
                    const dir = path.join(this.folder, target);
                    /* eslint-disable no-await-in-loop */
                    await fs.mkdir(dir, { recursive: true });
                    /* eslint-enable no-await-in-loop */
                } else {
                    const stream = archive.entries.get(name);
                    if (!stream) {
                        throw new Error(`Entry not found '${name}. Archive contains entries: ${JSON.stringify(Array.from(archive.entries.keys()))} .`);
                    }
                    const target = targets.shift();
                    const buffer = stream.peek();
                    const file = path.join(this.folder, target);
                    /* eslint-disable no-await-in-loop */
                    await fs.writeFile(file, buffer, null);
                    /* eslint-enable no-await-in-loop */
                }
            }
        } else {
            const target = targets.shift();
            this.status({ name: 'write', target });
            await fs.writeFile(`${this.folder}/${target}`, data, null);
        }
        if (targets.length > 0 && sources.length > 0) {
            await this.download(targets, sources);
        }
    }

    async load() {
        const target = path.resolve(this.folder, this.targets[0]);
        const identifier = path.basename(target);
        const stat = await fs.stat(target);
        let context = null;
        if (stat.isFile()) {
            const stream = new node.FileStream(target, 0, stat.size, stat.mtimeMs);
            const dirname = path.dirname(target);
            context = new mock.Context(this.host, dirname, identifier, stream, new Map());
        } else if (stat.isDirectory()) {
            const entries = new Map();
            const file = async (pathname) => {
                const stat = await fs.stat(pathname);
                const stream = new node.FileStream(pathname, 0, stat.size, stat.mtimeMs);
                const name = pathname.split(path.sep).join(path.posix.sep);
                entries.set(name, stream);
            };
            const walk = async (dir) => {
                const stats = await fs.readdir(dir, { withFileTypes: true });
                const promises = [];
                for (const stat of stats) {
                    const pathname = path.join(dir, stat.name);
                    if (stat.isDirectory()) {
                        promises.push(walk(pathname));
                    } else if (stat.isFile()) {
                        promises.push(file(pathname));
                    }
                }
                await Promise.all(promises);
            };
            await walk(target);
            context = new mock.Context(this.host, target, identifier, null, entries);
        }
        const modelFactoryService = new view.ModelFactoryService(this.host);
        this.model = await modelFactoryService.open(context);
        this.view.model = this.model;
    }

    async validate() {
        const model = this.model;
        if (!model.format || (this.format && this.format !== model.format)) {
            throw new Error(`Invalid model format '${model.format}'.`);
        }
        if (this.producer && model.producer !== this.producer) {
            throw new Error(`Invalid producer '${model.producer}'.`);
        }
        if (this.runtime && model.runtime !== this.runtime) {
            throw new Error(`Invalid runtime '${model.runtime}'.`);
        }
        if (model.metadata && (!Array.isArray(model.metadata) || !model.metadata.every((argument) => argument.name && (argument.value || argument.value === null || argument.value === '' || argument.value === false || argument.value === 0)))) {
            throw new Error("Invalid model metadata.");
        }
        if (this.assert) {
            for (const assert of this.assert) {
                const parts = assert.split('==').map((item) => item.trim());
                const properties = parts[0].split('.');
                const value = JSON.parse(parts[1].replace(/\s*'|'\s*/g, '"'));
                let context = { model };
                while (properties.length) {
                    const property = properties.shift();
                    if (context[property] !== undefined) {
                        context = context[property];
                        continue;
                    }
                    const match = /(.*)\[(.*)\]/.exec(property);
                    if (match && match.length === 3 && context[match[1]] !== undefined) {
                        const array = context[match[1]];
                        const index = parseInt(match[2], 10);
                        if (array[index] !== undefined) {
                            context = array[index];
                            continue;
                        }
                    }
                    throw new Error(`Invalid property path '${parts[0]}'.`);
                }
                if (context !== value) {
                    throw new Error(`Invalid '${context}' != '${assert}'.`);
                }
            }
        }
        if (model.version || model.description || model.author || model.license) {
            // continue
        }
        const validateGraph = async (graph) => {
            /* eslint-disable no-unused-expressions */
            const values = new Map();
            const validateValue = async (value) => {
                if (value === null) {
                    return;
                }
                value.name.toString();
                value.name.length;
                value.description;
                if (value.quantization) {
                    if (!this.tags.has('quantization')) {
                        throw new Error("Invalid 'quantization' tag.");
                    }
                    const quantization = new view.Quantization(value.quantization);
                    quantization.toString();
                }
                if (value.type) {
                    value.type.toString();
                }
                if (value.initializer) {
                    value.initializer.type.toString();
                    if (value.initializer && value.initializer.peek && !value.initializer.peek()) {
                        await value.initializer.read();
                    }
                    const tensor = new base.Tensor(value.initializer);
                    if (!this.tags.has('skip-tensor-value')) {
                        if (tensor.encoding !== '<' && tensor.encoding !== '>' && tensor.encoding !== '|') {
                            throw new Error(`Tensor encoding '${tensor.encoding}' is not implemented.`);
                        }
                        if (tensor.layout && (tensor.layout !== 'sparse' && tensor.layout !== 'sparse.coo')) {
                            throw new Error(`Tensor layout '${tensor.layout}' is not implemented.`);
                        }
                        if (!tensor.empty) {
                            if (tensor.type && tensor.type.dataType === '?') {
                                throw new Error('Tensor data type is not defined.');
                            } else if (tensor.type && !tensor.type.shape) {
                                throw new Error('Tensor shape is not defined.');
                            } else {
                                tensor.toString();
                                if (this.tags.has('validation')) {
                                    const size = tensor.type.shape.dimensions.reduce((a, b) => a * b, 1);
                                    if (tensor.type && tensor.type.dataType !== '?' && size < 8192) {
                                        let data_type = '?';
                                        switch (tensor.type.dataType) {
                                            case 'boolean': data_type = 'bool'; break;
                                            case 'bfloat16': data_type = 'float32'; break;
                                            case 'float8e5m2': data_type = 'float16'; break;
                                            case 'float8e5m2fnuz': data_type = 'float16'; break;
                                            case 'float8e4m3fn': data_type = 'float16'; break;
                                            case 'float8e4m3fnuz': data_type = 'float16'; break;
                                            case 'int4': data_type = 'int8'; break;
                                            default: data_type = tensor.type.dataType; break;
                                        }
                                        Target.execution = Target.execution || new python.Execution();
                                        const execution = Target.execution;
                                        const io = execution.__import__('io');
                                        const numpy = execution.__import__('numpy');
                                        const bytes = new io.BytesIO();
                                        const dtype = new numpy.dtype(data_type);
                                        const array = numpy.asarray(tensor.value, dtype);
                                        numpy.save(bytes, array);
                                    }
                                }
                            }
                        }
                    }
                } else if (value.name.length === 0) {
                    throw new Error('Empty value name.');
                }
                if (value.name.length > 0 && value.initializer === null) {
                    if (!values.has(value.name)) {
                        values.set(value.name, value);
                    } else if (value !== values.get(value.name)) {
                        throw new Error(`Duplicate value '${value.name}'.`);
                    }
                }
            };
            const signatures = Array.isArray(graph.signatures) ? graph.signatures : [graph];
            for (const signature of signatures) {
                for (const input of signature.inputs) {
                    input.name.toString();
                    input.name.length;
                    for (const value of input.value) {
                        /* eslint-disable no-await-in-loop */
                        await validateValue(value);
                        /* eslint-enable no-await-in-loop */
                    }
                }
                for (const output of signature.outputs) {
                    output.name.toString();
                    output.name.length;
                    if (Array.isArray(output.value)) {
                        for (const value of output.value) {
                            /* eslint-disable no-await-in-loop */
                            await validateValue(value);
                            /* eslint-enable no-await-in-loop */
                        }
                    }
                }
            }
            if (graph.metadata && (!Array.isArray(graph.metadata) || !graph.metadata.every((argument) => argument.name && argument.value !== undefined))) {
                throw new Error("Invalid graph metadata.");
            }
            for (const node of graph.nodes) {
                const type = node.type;
                if (!type || typeof type.name !== 'string') {
                    throw new Error(`Invalid node type '${JSON.stringify(node.type)}'.`);
                }
                if (Array.isArray(type.nodes)) {
                    /* eslint-disable no-await-in-loop */
                    await validateGraph(type);
                    /* eslint-enable no-await-in-loop */
                }
                view.Documentation.open(type);
                node.name.toString();
                node.description;
                if (node.metadata && (!Array.isArray(node.metadata) || !node.metadata.every((argument) => argument.name && argument.value !== undefined))) {
                    throw new Error("Invalid node metadata.");
                }
                const attributes = node.attributes;
                if (attributes) {
                    for (const attribute of attributes) {
                        attribute.name.toString();
                        attribute.name.length;
                        const type = attribute.type;
                        const value = attribute.value;
                        if ((type === 'graph' || type === 'function') && value && Array.isArray(value.nodes)) {
                            /* eslint-disable no-await-in-loop */
                            await validateGraph(value);
                            /* eslint-enable no-await-in-loop */
                        } else {
                            let text = new view.Formatter(attribute.value, attribute.type).toString();
                            if (text && text.length > 1000) {
                                text = `${text.substring(0, 1000)}...`;
                            }
                            /* value = */ text.split('<');
                        }
                    }
                }
                const inputs = node.inputs;
                if (Array.isArray(inputs)) {
                    for (const input of inputs) {
                        input.name.toString();
                        input.name.length;
                        if (!input.type || input.type.endsWith('*')) {
                            for (const value of input.value) {
                                /* eslint-disable no-await-in-loop */
                                await validateValue(value);
                                /* eslint-enable no-await-in-loop */
                            }
                            if (this.tags.has('validation')) {
                                if (input.value.length === 1 && input.value[0].initializer) {
                                    const sidebar = new view.TensorSidebar(this.view, input);
                                    sidebar.render();
                                }
                            }
                        }
                    }
                }
                const outputs = node.outputs;
                if (Array.isArray(outputs)) {
                    for (const output of node.outputs) {
                        output.name.toString();
                        output.name.length;
                        if (!output.type || output.type.endsWith('*')) {
                            for (const value of output.value) {
                                /* eslint-disable no-await-in-loop */
                                await validateValue(value);
                                /* eslint-enable no-await-in-loop */
                            }
                        }
                    }
                }
                if (node.chain) {
                    for (const chain of node.chain) {
                        chain.name.toString();
                        chain.name.length;
                    }
                }
                const sidebar = new view.NodeSidebar(this.view, node);
                sidebar.render();
            }
            const sidebar = new view.ModelSidebar(this.view, this.model, graph);
            sidebar.render();
            /* eslint-enable no-unused-expressions */
        };
        const validateTarget = async (target) => {
            switch (target.type) {
                default: {
                    await validateGraph(target);
                }
            }
        };
        for (const module of model.modules) {
            /* eslint-disable no-await-in-loop */
            await validateTarget(module);
            /* eslint-enable no-await-in-loop */
        }
        const functions = model.functions || [];
        for (const func of functions) {
            /* eslint-disable no-await-in-loop */
            await validateTarget(func);
            /* eslint-enable no-await-in-loop */
        }
    }

    async render() {
        for (const graph of this.model.modules) {
            const signatures = Array.isArray(graph.signatures) && graph.signatures.length > 0 ? graph.signatures : [graph];
            for (const signature of signatures) {
                /* eslint-disable no-await-in-loop */
                await this.view.render(graph, signature);
                /* eslint-enable no-await-in-loop */
            }
        }
    }
}

if (!worker_threads.isMainThread) {
    worker_threads.parentPort.addEventListener('message', async (e) => {
        const message = e.data;
        const response = {};
        try {
            const target = new Target(message);
            response.type = 'complete';
            response.target = target.name;
            target.on('status', (sender, message) => {
                message = { type: 'status', ...message };
                worker_threads.parentPort.postMessage(message);
            });
            if (message.measures) {
                target.measures = new Map();
            }
            await target.execute();
            response.measures = target.measures;
        } catch (error) {
            response.type = 'error';
            response.error = {
                name: error.name,
                message: error.message,
                stack: error.stack
            };
            const cause = error.cause;
            if (cause) {
                response.error.cause = {
                    name: cause.name,
                    message: cause.message
                };
            }
        }
        worker_threads.parentPort.postMessage(response);
    });
}