Skip to content
Snippets Groups Projects
Commit 027538ef authored by David Sehnal's avatar David Sehnal
Browse files

ModelServer: Custom prop support in pre-processer

parent 72b95a7d
No related branches found
No related tags found
No related merge requests found
...@@ -7,8 +7,7 @@ ...@@ -7,8 +7,7 @@
import * as fs from 'fs' import * as fs from 'fs'
import * as path from 'path' import * as path from 'path'
import * as argparse from 'argparse' import * as argparse from 'argparse'
import { preprocessFile } from './preprocess'; import { runMaster, PreprocessEntry } from './parallel';
import { ParallelPreprocessConfig, runMaster } from './parallel';
const cmdParser = new argparse.ArgumentParser({ const cmdParser = new argparse.ArgumentParser({
addHelp: true, addHelp: true,
...@@ -17,14 +16,16 @@ const cmdParser = new argparse.ArgumentParser({ ...@@ -17,14 +16,16 @@ const cmdParser = new argparse.ArgumentParser({
cmdParser.addArgument(['--input', '-i'], { help: 'Input filename', required: false }); cmdParser.addArgument(['--input', '-i'], { help: 'Input filename', required: false });
cmdParser.addArgument(['--outCIF', '-oc'], { help: 'Output CIF filename', required: false }); cmdParser.addArgument(['--outCIF', '-oc'], { help: 'Output CIF filename', required: false });
cmdParser.addArgument(['--outBCIF', '-ob'], { help: 'Output BinaryCIF filename', required: false }); cmdParser.addArgument(['--outBCIF', '-ob'], { help: 'Output BinaryCIF filename', required: false });
cmdParser.addArgument(['--bulk', '-b'], { help: 'Bulk JSON ({ numProcesses?: number, entries: { source: string, cif?: string, bcif?: string }[] })', required: false }); // TODO: add back? cmdParser.addArgument(['--bulk', '-b'], { help: 'Bulk JSON ({ numProcesses?: number, entries: { source: string, cif?: string, bcif?: string }[] })', required: false });
cmdParser.addArgument(['--folderIn', '-f'], { help: 'Convert folder', required: false }); cmdParser.addArgument(['--cfg', '-c'], { help: 'Config file path', required: false });
cmdParser.addArgument(['--folderIn', '-fin'], { help: 'Convert folder', required: false });
cmdParser.addArgument(['--folderOutCIF', '-foc'], { help: 'Convert folder text output', required: false }); cmdParser.addArgument(['--folderOutCIF', '-foc'], { help: 'Convert folder text output', required: false });
cmdParser.addArgument(['--folderOutBCIF', '-fob'], { help: 'Convert folder binary output', required: false }); cmdParser.addArgument(['--folderOutBCIF', '-fob'], { help: 'Convert folder binary output', required: false });
cmdParser.addArgument(['--folderNumProcesses', '-fp'], { help: 'Convert folder num processes', required: false }); cmdParser.addArgument(['--folderNumProcesses', '-fp'], { help: 'Convert folder num processes', required: false });
interface CmdArgs { interface CmdArgs {
bulk?: string, // bulk?: string,
cfg?: string,
input?: string, input?: string,
outCIF?: string, outCIF?: string,
outBCIF?: string, outBCIF?: string,
...@@ -34,31 +35,38 @@ interface CmdArgs { ...@@ -34,31 +35,38 @@ interface CmdArgs {
folderNumProcesses?: string folderNumProcesses?: string
} }
export interface PreprocessConfig {
numProcesses?: number,
customPropertyProviders?: string[]
}
const cmdArgs = cmdParser.parseArgs() as CmdArgs; const cmdArgs = cmdParser.parseArgs() as CmdArgs;
if (cmdArgs.input) preprocessFile(cmdArgs.input, cmdArgs.outCIF, cmdArgs.outBCIF); let entries: PreprocessEntry[] = []
else if (cmdArgs.bulk) runBulk(cmdArgs.bulk); let config: PreprocessConfig = { numProcesses: 1, customPropertyProviders: [] }
else if (cmdArgs.folderIn) runFolder(cmdArgs);
function runBulk(input: string) { if (cmdArgs.input) entries.push({ source: cmdArgs.input, cif: cmdArgs.outCIF, bcif: cmdArgs.outBCIF });
const config = JSON.parse(fs.readFileSync(input, 'utf8')) as ParallelPreprocessConfig; // else if (cmdArgs.bulk) runBulk(cmdArgs.bulk);
runMaster(config); else if (cmdArgs.folderIn) findEntries();
if (cmdArgs.cfg) {
config = JSON.parse(fs.readFileSync(cmdArgs.cfg, 'utf8')) as PreprocessConfig;
} }
function runFolder(args: CmdArgs) { runMaster(config, entries);
const files = fs.readdirSync(args.folderIn!);
const config: ParallelPreprocessConfig = { numProcesses: +args.folderNumProcesses! || 1, entries: [] }; function findEntries() {
const files = fs.readdirSync(cmdArgs.folderIn!);
const cifTest = /\.cif$/; const cifTest = /\.cif$/;
for (const f of files) { for (const f of files) {
if (!cifTest.test(f)) continue; if (!cifTest.test(f)) continue;
config.entries.push({ entries.push({
source: path.join(args.folderIn!, f), source: path.join(cmdArgs.folderIn!, f),
cif: cmdArgs.folderOutCIF ? path.join(args.folderOutCIF!, f) : void 0, cif: cmdArgs.folderOutCIF ? path.join(cmdArgs.folderOutCIF!, f) : void 0,
bcif: cmdArgs.folderOutBCIF ? path.join(args.folderOutBCIF!, path.parse(f).name + '.bcif') : void 0, bcif: cmdArgs.folderOutBCIF ? path.join(cmdArgs.folderOutBCIF!, path.parse(f).name + '.bcif') : void 0,
}); });
} }
runMaster(config);
} }
// example: // example:
......
...@@ -9,6 +9,9 @@ import * as cluster from 'cluster' ...@@ -9,6 +9,9 @@ import * as cluster from 'cluster'
import { now } from 'mol-task'; import { now } from 'mol-task';
import { PerformanceMonitor } from 'mol-util/performance-monitor'; import { PerformanceMonitor } from 'mol-util/performance-monitor';
import { preprocessFile } from './preprocess'; import { preprocessFile } from './preprocess';
import { createModelPropertiesProviderFromSources } from '../property-provider';
type PreprocessConfig = import('./master').PreprocessConfig
export interface PreprocessEntry { export interface PreprocessEntry {
source: string, source: string,
...@@ -16,27 +19,23 @@ export interface PreprocessEntry { ...@@ -16,27 +19,23 @@ export interface PreprocessEntry {
bcif?: string bcif?: string
} }
export interface ParallelPreprocessConfig { export function runMaster(config: PreprocessConfig, entries: PreprocessEntry[]) {
numProcesses?: number,
entries: PreprocessEntry[]
}
export function runMaster(config: ParallelPreprocessConfig) {
const parts = partitionArray(config.entries, config.numProcesses || 1);
// const numForks = Math.min(parts.length, config.numProcesses);
const started = now(); const started = now();
let progress = 0; let progress = 0;
const onMessage = (msg: any) => { const onMessage = (msg: any) => {
if (msg.type === 'tick') { if (msg.type === 'tick') {
progress++; progress++;
const elapsed = now() - started; const elapsed = now() - started;
console.log(`[${progress}/${config.entries.length}] in ${PerformanceMonitor.format(elapsed)} (avg ${PerformanceMonitor.format(elapsed / progress)}).`); console.log(`[${progress}/${entries.length}] in ${PerformanceMonitor.format(elapsed)} (avg ${PerformanceMonitor.format(elapsed / progress)}).`);
} else if (msg.type === 'error') { } else if (msg.type === 'error') {
console.error(`${msg.id}: ${msg.error}`) console.error(`${msg.id}: ${msg.error}`)
} }
} }
if (entries.length === 1) {
runSingle(entries[0], config, onMessage);
} else {
const parts = partitionArray(entries, config.numProcesses || 1);
for (const _ of parts) { for (const _ of parts) {
const worker = cluster.fork(); const worker = cluster.fork();
worker.on('message', onMessage); worker.on('message', onMessage);
...@@ -44,15 +43,17 @@ export function runMaster(config: ParallelPreprocessConfig) { ...@@ -44,15 +43,17 @@ export function runMaster(config: ParallelPreprocessConfig) {
let i = 0; let i = 0;
for (const id in cluster.workers) { for (const id in cluster.workers) {
cluster.workers[id]!.send(parts[i++]); cluster.workers[id]!.send({ entries: parts[i++], config });
}
} }
} }
export function runChild() { export function runChild() {
process.on('message', async (entries: PreprocessEntry[]) => { process.on('message', async ({ entries, config }: { entries: PreprocessEntry[], config: PreprocessConfig }) => {
const props = createModelPropertiesProviderFromSources(config.customPropertyProviders || []);
for (const entry of entries) { for (const entry of entries) {
try { try {
await preprocessFile(entry.source, entry.cif, entry.bcif); await preprocessFile(entry.source, props, entry.cif, entry.bcif);
} catch (e) { } catch (e) {
process.send!({ type: 'error', id: path.parse(entry.source).name, error: '' + e }); process.send!({ type: 'error', id: path.parse(entry.source).name, error: '' + e });
} }
...@@ -62,6 +63,16 @@ export function runChild() { ...@@ -62,6 +63,16 @@ export function runChild() {
}); });
} }
async function runSingle(entry: PreprocessEntry, config: PreprocessConfig, onMessage: (msg: any) => void) {
const props = createModelPropertiesProviderFromSources(config.customPropertyProviders || []);
try {
await preprocessFile(entry.source, props, entry.cif, entry.bcif);
} catch (e) {
onMessage({ type: 'error', id: path.parse(entry.source).name, error: '' + e });
}
onMessage({ type: 'tick' });
}
function partitionArray<T>(xs: T[], count: number): T[][] { function partitionArray<T>(xs: T[], count: number): T[][] {
const ret: T[][] = []; const ret: T[][] = [];
const s = Math.ceil(xs.length / count); const s = Math.ceil(xs.length / count);
......
...@@ -4,66 +4,45 @@ ...@@ -4,66 +4,45 @@
* @author David Sehnal <david.sehnal@gmail.com> * @author David Sehnal <david.sehnal@gmail.com>
*/ */
import { readStructureWrapper, resolveStructure } from '../server/structure-wrapper'; import { readStructureWrapper, resolveStructures } from '../server/structure-wrapper';
import { classifyCif } from './converter'; import { classifyCif } from './converter';
// import { ConsoleLogger } from 'mol-util/console-logger';
import { Structure } from 'mol-model/structure'; import { Structure } from 'mol-model/structure';
import { CifWriter } from 'mol-io/writer/cif'; import { CifWriter } from 'mol-io/writer/cif';
import Writer from 'mol-io/writer/writer'; import Writer from 'mol-io/writer/writer';
import { wrapFileToWriter } from '../server/api-local'; import { wrapFileToWriter } from '../server/api-local';
import { Task/*, now*/ } from 'mol-task';
import { /*showProgress, clearLine */ } from './util';
import { encode_mmCIF_categories, CifExportContext } from 'mol-model/structure/export/mmcif'; import { encode_mmCIF_categories, CifExportContext } from 'mol-model/structure/export/mmcif';
import { ModelPropertiesProvider } from '../property-provider';
// TODO: error handling // TODO: error handling
// let linearId = 0;
export async function preprocessFile(filename: string, outputCif?: string, outputBcif?: string) { export async function preprocessFile(filename: string, propertyProvider?: ModelPropertiesProvider, outputCif?: string, outputBcif?: string) {
// linearId++; const input = await readStructureWrapper('entry', '_local_', filename, propertyProvider);
//const started = now();
//ConsoleLogger.log(`${linearId}`, `Reading '${filename}'...`);
// TODO: support the custom prop provider list here.
const input = await readStructureWrapper('entry', '_local_', filename, void 0);
const categories = await classifyCif(input.cifFrame); const categories = await classifyCif(input.cifFrame);
const inputStructure = (await resolveStructure(input))!; const inputStructures = (await resolveStructures(input))!;
//ConsoleLogger.log(`${linearId}`, `Classifying CIF categories...`); const exportCtx = CifExportContext.create(inputStructures);
//clearLine();
const exportCtx = CifExportContext.create(inputStructure);
if (outputCif) { if (outputCif) {
//ConsoleLogger.log(`${linearId}`, `Encoding CIF...`);
const writer = wrapFileToWriter(outputCif); const writer = wrapFileToWriter(outputCif);
const encoder = CifWriter.createEncoder({ binary: false }); const encoder = CifWriter.createEncoder({ binary: false });
await encode(inputStructure, input.cifFrame.header, categories, encoder, exportCtx, writer); encode(inputStructures[0], input.cifFrame.header, categories, encoder, exportCtx, writer);
// clearLine();
writer.end(); writer.end();
} }
if (outputBcif) { if (outputBcif) {
// ConsoleLogger.log(`${linearId}`, `Encoding BinaryCIF...`);
const writer = wrapFileToWriter(outputBcif); const writer = wrapFileToWriter(outputBcif);
const encoder = CifWriter.createEncoder({ binary: true, binaryAutoClassifyEncoding: true }); const encoder = CifWriter.createEncoder({ binary: true, binaryAutoClassifyEncoding: true });
await encode(inputStructure, input.cifFrame.header, categories, encoder, exportCtx, writer); encode(inputStructures[0], input.cifFrame.header, categories, encoder, exportCtx, writer);
//clearLine();
writer.end(); writer.end();
} }
// ConsoleLogger.log(`${linearId}`, `Finished '${filename}' in ${Math.round(now() - started)}ms`);
} }
function encode(structure: Structure, header: string, categories: CifWriter.Category[], encoder: CifWriter.Encoder, exportCtx: CifExportContext[], writer: Writer) { function encode(structure: Structure, header: string, categories: CifWriter.Category[], encoder: CifWriter.Encoder, exportCtx: CifExportContext[], writer: Writer) {
return Task.create('Encode', async ctx => {
const skipCategoryNames = new Set<string>(categories.map(c => c.name)); const skipCategoryNames = new Set<string>(categories.map(c => c.name));
encoder.startDataBlock(header); encoder.startDataBlock(header);
// let current = 0;
for (const cat of categories) { for (const cat of categories) {
encoder.writeCategory(cat); encoder.writeCategory(cat);
// current++;
// if (ctx.shouldUpdate) await ctx.update({ message: 'Encoding...', current, max: categories.length });
} }
encode_mmCIF_categories(encoder, structure, { skipCategoryNames, exportCtx }); encode_mmCIF_categories(encoder, structure, { skipCategoryNames, exportCtx });
encoder.encode(); encoder.encode();
encoder.writeTo(writer); encoder.writeTo(writer);
}).run();
} }
\ No newline at end of file
...@@ -10,10 +10,14 @@ import Config from './config'; ...@@ -10,10 +10,14 @@ import Config from './config';
export type ModelPropertiesProvider = (model: Model, cache: object) => Promise<any>[] export type ModelPropertiesProvider = (model: Model, cache: object) => Promise<any>[]
export function createModelPropertiesProviderFromConfig(): ModelPropertiesProvider { export function createModelPropertiesProviderFromConfig(): ModelPropertiesProvider {
if (!Config.customPropertyProviders || Config.customPropertyProviders.length === 0) return () => []; return createModelPropertiesProviderFromSources(Config.customPropertyProviders);
}
export function createModelPropertiesProviderFromSources(sources: string[]): ModelPropertiesProvider {
if (!sources || sources.length === 0) return () => [];
const ps: ModelPropertiesProvider[] = []; const ps: ModelPropertiesProvider[] = [];
for (const p of Config.customPropertyProviders) { for (const p of sources) {
ps.push(require(p).attachModelProperties); ps.push(require(p).attachModelProperties);
} }
...@@ -25,4 +29,3 @@ export function createModelPropertiesProviderFromConfig(): ModelPropertiesProvid ...@@ -25,4 +29,3 @@ export function createModelPropertiesProviderFromConfig(): ModelPropertiesProvid
return ret; return ret;
} }
} }
...@@ -16,7 +16,7 @@ import Version from '../version'; ...@@ -16,7 +16,7 @@ import Version from '../version';
import { Job } from './jobs'; import { Job } from './jobs';
import { createStructureWrapperFromJob, StructureWrapper, resolveStructures } from './structure-wrapper'; import { createStructureWrapperFromJob, StructureWrapper, resolveStructures } from './structure-wrapper';
import CifField = CifWriter.Field import CifField = CifWriter.Field
import { createModelPropertiesProviderFromConfig } from '../property-provider'; import { createModelPropertiesProviderFromConfig, ModelPropertiesProvider } from '../property-provider';
export interface Stats { export interface Stats {
structure: StructureWrapper, structure: StructureWrapper,
...@@ -26,12 +26,17 @@ export interface Stats { ...@@ -26,12 +26,17 @@ export interface Stats {
const perf = new PerformanceMonitor(); const perf = new PerformanceMonitor();
const propertyProvider = createModelPropertiesProviderFromConfig(); let _propertyProvider: ModelPropertiesProvider;
function propertyProvider() {
if (_propertyProvider) return _propertyProvider;
_propertyProvider = createModelPropertiesProviderFromConfig();
return _propertyProvider;
}
export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> { export async function resolveJob(job: Job): Promise<CifWriter.Encoder<any>> {
ConsoleLogger.logId(job.id, 'Query', 'Starting.'); ConsoleLogger.logId(job.id, 'Query', 'Starting.');
const wrappedStructure = await createStructureWrapperFromJob(job, propertyProvider); const wrappedStructure = await createStructureWrapperFromJob(job, propertyProvider());
try { try {
perf.start('query'); perf.start('query');
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment