[curtischong]

While training nequix I noticed that index_queue.get(timeout=0) makes the workers continuously poll (because of the timeout=0). By removing this timeout, the workers block until they receive a message in the queue.

This is a benchmark run by claude:

Workers looped on index_queue.get(timeout=0) / continue, burning ~17% of a
core each while idle (~23k wakeups/s). The parent had the same spin on
output_queue. Blocking get() drops idle CPU to zero and cuts round-trip
latency 5-16x. Also add DataLoader.shutdown() so the None sentinel is
actually sent and workers are joined instead of living until process exit.

Sidenote: we might also want to update https://teddykoker.com/2020/12/dataloader to reflect this change.

This PR also adds explicit shutdown calls to the workers

This PR was tested to train on 1500 matprj structures on 110 training steps

 import itertools
  from pathlib import Path

  import ase.db
  import ase.io
  import yaml

  SRC = Path("/data/curtis/dev/nequix/data/mptrj-gga-ggapu")
  OUT = Path(__file__).parent / "smoke-mptrj.aselmdb"
  N = 1500

  OUT.unlink(missing_ok=True)
  count = 0
  with ase.db.connect(OUT) as db:
      for f in sorted(SRC.glob("*.extxyz")):
          for atoms in ase.io.iread(f, index=":"):
              db.write(atoms)
              count += 1
              if count >= N:
                  break
          if count >= N:
              break
  print(f"wrote {count} structures to {OUT}")

  config = yaml.safe_load(Path("/data/curtis/nequix/configs/nequix-mp-1.yml").read_text())
  for key in ["shift", "scale", "avg_n_neighbors", "max_n_edges", "max_n_nodes", "avg_n_nodes", "avg_n_edges"]:
      config.pop(key, None)
  config.update(
      train_path=str(OUT),
      valid_frac=0.1,
      state_path=str(OUT.parent / "smoke-state.pkl"),
      hidden_irreps="32x0e + 16x1o",
      lmax=2,
      n_layers=2,
      batch_size=16,
      n_epochs=1,
      log_every=10,
      kernel=False,
  )
  smoke_config = OUT.parent / "smoke-mptrj.yml"
  smoke_config.write_text(yaml.dump(config))
  print(f"wrote {smoke_config}")

[curtischong]

from the docs https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Queue.cancel_join_thread

A better name for this method might be allow_exit_without_flush(). It is likely to cause enqueued data to be lost, and you almost certainly will not need to use it. It is really only there if you need the current process to exit immediately without waiting to flush enqueued data to the underlying pipe, and you don’t care about lost data.