code/gen_data.cjs

#!/usr/bin/env node
/**
 * Training-data generator using cubejs, as an alternative to gen_data.py.
 *
 * The Python path depends on kociemba's C extension, which repeatedly refused to
 * build on rented boxes while installing "successfully" as a pure-Python fallback
 * ~50x slower. This path has no compiled dependency at all: cubejs is pure
 * JavaScript, is already a dependency of packages/solver, and is the very solver
 * the harness scores with -- so labels here come from the same implementation
 * that grades the benchmark.
 *
 * Measured ~59 ms/solve single-threaded against native kociemba's ~23 ms. Slower
 * per core, but it actually runs, and the boxes have enough cores that generation
 * is still well under an hour.
 *
 *   node gen_data.cjs --count 10000000 --workers $(nproc) --out data/train.jsonl
 */
const fs = require('fs');
const os = require('os');
const path = require('path');
const { fork } = require('child_process');

function loadCube() {
  // Resolve from the local install or from a sibling node_modules, so the script
  // works both inside this repo and standing alone on a fresh box.
  for (const p of ['cubejs', path.join(__dirname, 'node_modules', 'cubejs')]) {
    try { return require(p); } catch { /* try next */ }
  }
  throw new Error('cubejs not found -- run: npm install cubejs');
}

/** Deterministic PRNG so a seed reproduces a dataset; cubejs uses Math.random. */
function mulberry32(seed) {
  let a = seed >>> 0;
  return function () {
    a |= 0; a = (a + 0x6D2B79F5) | 0;
    let t = Math.imul(a ^ (a >>> 15), 1 | a);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

const FACES = ['U', 'R', 'F', 'D', 'L', 'B'];
const SUFFIX = ['', "'", '2'];

function parseArgs(argv) {
  const get = (f, d) => { const i = argv.indexOf(f); return i === -1 ? d : argv[i + 1]; };
  return {
    count: Number(get('--count', 1000000)),
    workers: Number(get('--workers', os.cpus().length)),
    seed: Number(get('--seed', 0)),
    shallowFrac: Number(get('--shallow-frac', 0.15)),
    maxShallow: Number(get('--max-shallow', 8)),
    out: get('--out', 'data/train.jsonl'),
  };
}

function runWorker() {
  const { count, seed, shallowFrac, maxShallow, out } = JSON.parse(process.env.JOB);
  const Cube = loadCube();
  Cube.initSolver();
  const rand = mulberry32(seed);
  Math.random = rand;  // cubejs's Cube.random() draws from this

  const stream = fs.createWriteStream(out);
  let written = 0;
  while (written < count) {
    let cube, depth;
    if (rand() < shallowFrac) {
      depth = 1 + Math.floor(rand() * maxShallow);
      cube = new Cube();
      let prev = null;
      for (let i = 0; i < depth; i++) {
        let f;
        do { f = FACES[Math.floor(rand() * 6)]; } while (f === prev);
        prev = f;
        cube.move(f + SUFFIX[Math.floor(rand() * 3)]);
      }
    } else {
      // Uniform over the whole group -- the fully-mixed distribution that every
      // scramble past ~15 moves samples from anyway.
      cube = Cube.random();
      depth = 200;
    }
    if (cube.isSolved()) continue;
    const state = cube.asString();
    const solution = cube.solve();
    if (!solution) continue;
    stream.write(JSON.stringify({ state, solution, depth }) + '\n');
    written++;
    if (written % 20000 === 0) process.send({ written });
  }
  stream.end(() => process.exit(0));
}

function runMaster() {
  const args = parseArgs(process.argv.slice(2));
  fs.mkdirSync(path.dirname(args.out) || '.', { recursive: true });
  const per = Math.ceil(args.count / args.workers);
  const start = Date.now();
  let done = 0, total = 0;
  const parts = [];

  for (let i = 0; i < args.workers; i++) {
    const partOut = `${args.out}.part${i}`;
    parts.push(partOut);
    const job = { count: per, seed: args.seed + 1000 * i + 1, shallowFrac: args.shallowFrac,
                  maxShallow: args.maxShallow, out: partOut };
    const child = fork(__filename, [], { env: { ...process.env, JOB: JSON.stringify(job), IS_WORKER: '1' } });
    child.on('message', (m) => {
      total += 20000;
      const rate = total / ((Date.now() - start) / 1000);
      process.stderr.write(`  ${total}/${args.count} (${rate.toFixed(0)}/s)\n`);
    });
    child.on('exit', () => {
      if (++done === args.workers) {
        const outStream = fs.createWriteStream(args.out);
        let n = 0;
        for (const p of parts) {
          for (const line of fs.readFileSync(p, 'utf8').split('\n')) {
            if (line) { outStream.write(line + '\n'); n++; }
          }
          fs.unlinkSync(p);
        }
        outStream.end(() => {
          const secs = (Date.now() - start) / 1000;
          process.stderr.write(`wrote ${n} pairs in ${secs.toFixed(0)}s (${(n / secs).toFixed(0)}/s, ${args.workers} workers)\n`);
        });
      }
    });
  }
}

if (process.env.IS_WORKER) runWorker(); else runMaster();