diff --git a/Exercises/DAOS/PyTorch/00_env.sh b/Exercises/DAOS/PyTorch/00_env.sh new file mode 100644 index 0000000..c1f2622 --- /dev/null +++ b/Exercises/DAOS/PyTorch/00_env.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Source this file: source 00_env.sh +# Sets up the DAOS client environment + Python venv for the pydaos.torch tutorial. + +INSTALL=/home/chaarawi/install +DAOS=$INSTALL/daos + +# Python: venv with torch + numpy (activate first: it rewrites PATH) +source /home/chaarawi/my_venv/bin/activate + +# DAOS client libraries and tools, plus the pydaos.torch module installed by the DAOS build +export PATH=$DAOS/bin:$PATH +export LD_LIBRARY_PATH=$DAOS/lib64:$DAOS/lib:$LD_LIBRARY_PATH +export DAOS_AGENT_DRPC_DIR=/tmp/chaarawi/daos_agent +export PYTHONPATH=$DAOS/lib64/python3.11/site-packages:$PYTHONPATH + +# Tutorial parameters (used by all the example scripts) +export TUT_POOL=mypool +export TUT_CONT=pytorch_tut +export TUT_NSAMPLES=10000 + +# Quieter DAOS client logging for demos; raise to DEBUG when troubleshooting +export D_LOG_MASK=ERR +export D_LOG_FILE=/tmp/chaarawi/pytorch_tut.log + +echo "pool=$TUT_POOL cont=$TUT_CONT python=$(python --version 2>&1)" diff --git a/Exercises/DAOS/PyTorch/01_populate.py b/Exercises/DAOS/PyTorch/01_populate.py new file mode 100644 index 0000000..4c93208 --- /dev/null +++ b/Exercises/DAOS/PyTorch/01_populate.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +""" +01_populate.py - create a synthetic dataset inside the DAOS POSIX container. + +Layout (ImageFolder-like): + /data/train/class_/sample_.npz x: float32[3,32,32] y: int64 label + +Each sample is a noisy version of a per-class pattern so a tiny model can learn it. +Files are written with the pydaos.torch Checkpoint writer (plain libdfs writes, no dfuse). +Run once before the tutorial: python 01_populate.py +""" +import io +import os +import sys +from multiprocessing import get_context + +import numpy as np + +POOL = os.environ.get("TUT_POOL", "mypool") +CONT = os.environ.get("TUT_CONT", "pytorch_tut") +NSAMPLES = int(os.environ.get("TUT_NSAMPLES", "10000")) +NCLASSES = 10 +PREFIX = "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/data/train" +WRITERS = 16 + + +def make_sample(rng, idx): + label = idx % NCLASSES + pattern = np.zeros((3, 32, 32), dtype=np.float32) + pattern[label % 3, (label * 3) % 32:(label * 3) % 32 + 8, :] = 1.0 + x = pattern + rng.normal(0, 0.5, size=pattern.shape).astype(np.float32) + buf = io.BytesIO() + np.savez(buf, x=x, y=np.int64(label)) + return f"class_{label}/sample_{idx:06d}.npz", buf.getvalue() + + +def writer_proc(worker_id, indices): + # Each worker owns its own DAOS connection (DaosClient is a per-process singleton). + from pydaos.torch import Checkpoint + + rng = np.random.default_rng(worker_id) + # transfer_chunk_size=0: buffer in memory, single dfs write on close() + with Checkpoint(POOL, CONT, prefix=PREFIX, transfer_chunk_size=0) as chkp: + for idx in indices: + name, data = make_sample(rng, idx) + with chkp.writer(name) as w: + w.write(data) + return len(indices) + + +def main(): + print(f"Populating {POOL}:{CONT}{PREFIX} with {NSAMPLES} samples, {NCLASSES} classes, " + f"{WRITERS} writer processes") + shards = [list(range(w, NSAMPLES, WRITERS)) for w in range(WRITERS)] + + # 'spawn' so that no DAOS state is inherited across fork from this process + with get_context("spawn").Pool(WRITERS) as pool: + done = pool.starmap(writer_proc, enumerate(shards)) + + print(f"wrote {sum(done)} files under {PREFIX}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/Exercises/DAOS/PyTorch/02_map_dataset.py b/Exercises/DAOS/PyTorch/02_map_dataset.py new file mode 100644 index 0000000..f50af10 --- /dev/null +++ b/Exercises/DAOS/PyTorch/02_map_dataset.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +""" +02_map_dataset.py - map-style pydaos.torch.Dataset + +Demonstrates: + 1. constructing a Dataset (namespace scan happens here) + 2. indexing: raw bytes vs. a transform_fn that decodes samples into tensors + 3. feeding a torch DataLoader with worker processes + (worker_init_fn=dataset.worker_init is REQUIRED with num_workers > 0) + +Usage: python 02_map_dataset.py [--workers 4] [--batch-size 64] [--no-worker-init] +""" +import argparse + +import torch +from torch.utils.data import DataLoader + +from pydaos.torch import Dataset +from tut_common import CONT, POOL, TRAIN_PATH, decode_npz + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--batch-size", type=int, default=64) + ap.add_argument("--no-worker-init", action="store_true", + help="skip worker_init_fn to show what happens after fork()") + args = ap.parse_args() + + # --- 1. construct: scans /data/train recursively, in parallel ------------------------- + raw = Dataset(pool=POOL, cont=CONT, path=TRAIN_PATH) + print(f"len(dataset) = {len(raw)}") + path, size = raw.objects[0] + print(f"objects[0] = ({path}, {size} bytes)") + + # --- 2a. default transform_fn returns the raw file content --------------------------- + sample = raw[0] + print(f"raw[0] -> {type(sample).__name__} of {len(sample)} bytes") + raw.close() + + # --- 2b. transform_fn decodes bytes -> (tensor, label) -------------------------------- + with Dataset(pool=POOL, cont=CONT, path=TRAIN_PATH, transform_fn=decode_npz) as ds: + x, y = ds[0] + print(f"ds[0] -> tensor {tuple(x.shape)} {x.dtype}, label {y}") + + # __getitems__: one call reads a whole batch of files in parallel through libdfs + batch = ds.__getitems__([0, 1, 2, 3]) + print(f"ds.__getitems__([0..3]) -> {len(batch)} samples, labels {[b[1] for b in batch]}") + + # --- 3. DataLoader ------------------------------------------------------------------ + loader = DataLoader(ds, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.workers, + # re-initializes the DAOS/DFS handles in each forked worker + worker_init_fn=None if args.no_worker_init else ds.worker_init) + nbatches = 0 + nsamples = 0 + shape = None + for x, y in loader: + shape = shape or tuple(x.shape) + nbatches += 1 + nsamples += x.shape[0] + print(f"DataLoader(num_workers={args.workers}): {nbatches} batches, {nsamples} samples, " + f"batch tensor {shape}") + + +if __name__ == "__main__": + torch.set_num_threads(4) + main() diff --git a/Exercises/DAOS/PyTorch/03_iterable_dataset.py b/Exercises/DAOS/PyTorch/03_iterable_dataset.py new file mode 100644 index 0000000..152712c --- /dev/null +++ b/Exercises/DAOS/PyTorch/03_iterable_dataset.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +""" +03_iterable_dataset.py - pydaos.torch.IterableDataset + +Differences from the map-style Dataset: + * no random access / no shuffling by the DataLoader: samples arrive in scan order + * the dataset itself reads `batch_size` files per libdfs call (lazy batching) + * worker_init(worker_id) both re-inits DAOS after fork AND shards the file list, + so each DataLoader worker streams a disjoint slice of the namespace + +Usage: python 03_iterable_dataset.py [--workers 4] [--read-batch 32] [--batch-size 64] +""" +import argparse + +import torch +from torch.utils.data import DataLoader + +from pydaos.torch import IterableDataset +from tut_common import CONT, POOL, TRAIN_PATH, decode_npz + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--read-batch", type=int, default=32, + help="files read per libdfs batch_read inside the dataset") + ap.add_argument("--batch-size", type=int, default=64, help="DataLoader batch size") + args = ap.parse_args() + + with IterableDataset(pool=POOL, cont=CONT, path=TRAIN_PATH, + transform_fn=decode_npz, + batch_size=args.read_batch) as ds: + print(f"{len(ds.objects)} objects discovered") + + # plain python iteration, no DataLoader: yields one sample at a time + x, y = next(iter(ds)) + print(f"next(iter(ds)) -> tensor {tuple(x.shape)}, label {y}") + + loader = DataLoader(ds, + batch_size=args.batch_size, # DataLoader's collate batch + num_workers=args.workers, + worker_init_fn=ds.worker_init) # re-init DAOS + shard the file list + nbatches = 0 + nsamples = 0 + labels = set() + for x, y in loader: + nbatches += 1 + nsamples += x.shape[0] + labels.update(y.tolist()) + # nsamples == len(ds.objects) shows the workers covered disjoint, complete shards + print(f"DataLoader(num_workers={args.workers}, dataset batch={args.read_batch}, " + f"loader batch={args.batch_size}): {nbatches} batches, {nsamples} samples, " + f"{len(labels)} distinct labels") + + +if __name__ == "__main__": + torch.set_num_threads(4) + main() diff --git a/Exercises/DAOS/PyTorch/04_checkpoint.py b/Exercises/DAOS/PyTorch/04_checkpoint.py new file mode 100644 index 0000000..138af14 --- /dev/null +++ b/Exercises/DAOS/PyTorch/04_checkpoint.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +""" +04_checkpoint.py - saving / loading PyTorch checkpoints on DAOS with pydaos.torch.Checkpoint + + chkp.writer(name) -> file-like object for torch.save() + chkp.reader(name) -> BytesIO (or a user-provided stream) for torch.load() + +Two write modes: + transfer_chunk_size=0 buffer everything in memory, one dfs write on close() + transfer_chunk_size=N cut the stream into N-byte chunks, `workers` processes write + them in parallel; at most `chunks_limit` chunks are queued + (bounds memory: ~ chunks_limit * N bytes in flight) + +class_name takes the names used by `daos cont create --file-oclass` ("SX", "S1", "EC_2P1GX" +...), NOT the C enum spelling: "OC_SX" is unknown and silently falls back to the container +default. + +Pitfall: the default open_flags are O_CREAT|O_RDWR (no O_TRUNC) and in chunked mode every +chunk re-opens the file with those flags, so O_TRUNC cannot be used there. Overwriting an +existing checkpoint with a *smaller* one leaves the old tail in place and torch.load() fails. +Use unique file names (epoch/step in the name) or remove the old file first. + +Usage: python 04_checkpoint.py [--size-mb 256] +""" +import argparse +import time + +import torch + +from pydaos.torch import Checkpoint +from tut_common import CKPT_PATH, CONT, POOL + +MiB = 2**20 + + +def make_state(size_mb): + """A fake model state_dict of roughly size_mb MiB (fp32).""" + n = size_mb * MiB // 4 + return {"layer.weight": torch.arange(n, dtype=torch.float32), "epoch": 7} + + +def round_trip(chkp, name, state): + with chkp.writer(name) as w: + torch.save(state, w) + with chkp.reader(name) as r: + loaded = torch.load(r) + assert loaded["epoch"] == state["epoch"] + assert torch.equal(loaded["layer.weight"], state["layer.weight"]) + print(f" {name}: saved and loaded back identically") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--size-mb", type=int, default=256) + args = ap.parse_args() + + state = make_state(args.size_mb) + print(f"state_dict of {args.size_mb} MiB") + # per-run directory: see the O_TRUNC pitfall in the module docstring + run = f"{CKPT_PATH}/run_{int(time.time())}" + + # --- minimal usage ------------------------------------------------------------------- + with Checkpoint(POOL, CONT, prefix=run) as chkp: + # parent directories under prefix are created on demand (ensure_path=True) + with chkp.writer("model/epoch_007.pt") as w: + torch.save(state, w) + with chkp.reader("model/epoch_007.pt") as r: + print(f"loaded epoch = {torch.load(r)['epoch']}") + + # reader(stream=...) streams into any writable object, e.g. a local file, + # for checkpoints that should not be materialised in RAM + with open("/tmp/epoch_007.pt", "wb") as f: + chkp.reader("model/epoch_007.pt", stream=f) + print("streamed a copy to /tmp/epoch_007.pt") + + # --- write modes ----------------------------------------------------------------------- + print("in-memory buffer, single write on close():") + with Checkpoint(POOL, CONT, prefix=run, transfer_chunk_size=0) as chkp: + round_trip(chkp, "inmem.pt", state) + + print("chunked: 64 MiB chunks, up to 4 queued, 4 writer processes:") + with Checkpoint(POOL, CONT, prefix=run, + transfer_chunk_size=64 * MiB, chunks_limit=4, workers=4) as chkp: + round_trip(chkp, "chunked.pt", state) + + # --- per-file object class / chunk size ----------------------------------------------- + # The container default here is S1 (good for 12 KB samples). Large checkpoint files + # should be striped across engines: pick the class per Checkpoint instance. + print("chunked + striped object class (SX) and 4 MiB DFS chunk size for the file:") + with Checkpoint(POOL, CONT, prefix=run, + class_name="SX", file_chunk_size=4 * MiB, + transfer_chunk_size=64 * MiB, chunks_limit=4, workers=4) as chkp: + round_trip(chkp, "chunked_SX.pt", state) + print(f"inspect with: daos fs get-attr {POOL} {CONT} --dfs-path {run}/chunked_SX.pt") + + +if __name__ == "__main__": + main() diff --git a/Exercises/DAOS/PyTorch/05_train_e2e.py b/Exercises/DAOS/PyTorch/05_train_e2e.py new file mode 100644 index 0000000..facfde3 --- /dev/null +++ b/Exercises/DAOS/PyTorch/05_train_e2e.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +""" +05_train_e2e.py - end to end: train from a DAOS dataset, checkpoint to DAOS, resume. + + python 05_train_e2e.py # train epochs 1..2, checkpoint after each + python 05_train_e2e.py --resume # load latest checkpoint from DAOS, continue to --epochs + +Everything the training loop touches on storage goes through pydaos.torch: + samples : Dataset + DataLoader(num_workers, worker_init_fn=ds.worker_init) + checkpoints: Checkpoint.writer()/reader() with torch.save()/torch.load() +""" +import argparse + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader + +from pydaos.torch import Checkpoint, Dataset +from tut_common import CKPT_PATH, CONT, POOL, TRAIN_PATH, decode_npz + +RUN = "e2e" + + +class TinyCNN(nn.Module): + def __init__(self, nclasses=10): + super().__init__() + self.conv = nn.Conv2d(3, 8, 3, padding=1) + self.fc = nn.Linear(8 * 8 * 8, nclasses) + + def forward(self, x): + x = F.max_pool2d(F.relu(self.conv(x)), 4) # 32x32 -> 8x8 + return self.fc(x.flatten(1)) + + +def save_ckpt(chkp, model, opt, epoch): + name = f"{RUN}/epoch_{epoch:03d}.pt" + with chkp.writer(name) as w: + torch.save({"epoch": epoch, + "model": model.state_dict(), + "optimizer": opt.state_dict()}, w) + # 'latest' is rewritten each epoch; same size every time so no-truncate is harmless here + with chkp.writer(f"{RUN}/latest.pt") as w: + torch.save({"epoch": epoch, "file": name}, w) + print(f" checkpoint -> {CKPT_PATH}/{name}") + + +def load_latest(chkp, model, opt): + with chkp.reader(f"{RUN}/latest.pt") as r: + latest = torch.load(r) + with chkp.reader(latest["file"]) as r: + state = torch.load(r) + model.load_state_dict(state["model"]) + opt.load_state_dict(state["optimizer"]) + print(f" resumed from {CKPT_PATH}/{latest['file']} (epoch {state['epoch']})") + return state["epoch"] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--epochs", type=int, default=2) + ap.add_argument("--resume", action="store_true") + ap.add_argument("--workers", type=int, default=8) + ap.add_argument("--batch-size", type=int, default=64) + args = ap.parse_args() + + torch.manual_seed(0) + model = TinyCNN() + opt = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9) + start_epoch = 0 + + # Open the Checkpoint (read-write) BEFORE the Dataset (read-only): libdfs caches the + # pool/container handles per process, keyed by label, and the first open's access mode + # wins. Dataset first -> RO handles cached -> checkpoint writes fail with EPERM. + with Checkpoint(POOL, CONT, prefix=CKPT_PATH, class_name="SX") as chkp, \ + Dataset(pool=POOL, cont=CONT, path=TRAIN_PATH, transform_fn=decode_npz) as ds: + + if args.resume: + start_epoch = load_latest(chkp, model, opt) + + loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True, + num_workers=args.workers, worker_init_fn=ds.worker_init) + + for epoch in range(start_epoch + 1, args.epochs + 1): + total, correct, loss_sum = 0, 0, 0.0 + for x, y in loader: + opt.zero_grad() + out = model(x) + loss = F.cross_entropy(out, y) + loss.backward() + opt.step() + loss_sum += loss.item() * y.numel() + correct += (out.argmax(1) == y).sum().item() + total += y.numel() + print(f"epoch {epoch}: loss {loss_sum / total:.3f} acc {correct / total:.3f} " + f"({total} samples)") + save_ckpt(chkp, model, opt, epoch) + + +if __name__ == "__main__": + torch.set_num_threads(8) + main() diff --git a/Exercises/DAOS/PyTorch/README.md b/Exercises/DAOS/PyTorch/README.md new file mode 100644 index 0000000..7065da5 --- /dev/null +++ b/Exercises/DAOS/PyTorch/README.md @@ -0,0 +1,37 @@ +# pydaos.torch tutorial (SC26) — material + +Working examples for "Feeding PyTorch from DAOS". Slides: `pydaos_torch_tutorial.pptx`. + +## Run + +```bash +cd Exercises/DAOS/PyTorch +source 00_env.sh # venv + DAOS paths; sets TUT_POOL / TUT_CONT +python 01_populate.py # once: 10k .npz samples into /data/train +python 02_map_dataset.py # Dataset + DataLoader +python 02_map_dataset.py --workers 2 --no-worker-init # the fork failure +python 03_iterable_dataset.py # IterableDataset +python 04_checkpoint.py # Checkpoint write modes, object class +python 05_train_e2e.py # train 2 epochs, checkpoint each +python 05_train_e2e.py --resume --epochs 3 +``` + +Container (already created): + +```bash +daos cont create mypool pytorch_tut --type POSIX --dir-oclass SX --file-oclass S1 --properties rd_fac:0 +``` + +## Files + +| File | Purpose | +|---|---| +| `00_env.sh` | environment: `my_venv` (torch 2.14 CPU), `PYTHONPATH` to DAOS site-packages, pool/cont names | +| `tut_common.py` | `decode_npz` transform, pool/cont/path constants | +| `01_populate.py` | synthetic dataset writer (uses `Checkpoint.writer`, 16 spawned processes) | +| `02_map_dataset.py` | map-style `Dataset`: scan, indexing, `transform_fn`, `__getitems__`, `DataLoader`, `worker_init_fn` | +| `03_iterable_dataset.py` | `IterableDataset`: lazy batching, worker sharding | +| `04_checkpoint.py` | `Checkpoint`: in-memory vs chunked writes, `reader(stream=)`, `class_name`, pitfalls | +| `05_train_e2e.py` | tiny CNN trained from DAOS, checkpoints to DAOS, `--resume` | +| `daos_overview.pptx` | DAOS introduction: concepts, API basics, POSIX on DAOS | +| `pydaos_torch_tutorial.pptx` | the PyTorch data-loader / checkpoint deck | diff --git a/Exercises/DAOS/PyTorch/tut_common.py b/Exercises/DAOS/PyTorch/tut_common.py new file mode 100644 index 0000000..876048b --- /dev/null +++ b/Exercises/DAOS/PyTorch/tut_common.py @@ -0,0 +1,23 @@ +""" +tut_common.py - helpers shared by the tutorial scripts. +""" +import io +import os + +import numpy as np +import torch + +POOL = os.environ.get("TUT_POOL", "mypool") +CONT = os.environ.get("TUT_CONT", "pytorch_tut") +TRAIN_PATH = "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/data/train" +CKPT_PATH = "/checkpoints" + + +def decode_npz(data): + """transform_fn: raw bytes read from DFS -> (tensor[3,32,32], label). + + pydaos.torch hands the sample content to transform_fn as a bytearray; whatever + it returns is what DataLoader collates into a batch. + """ + npz = np.load(io.BytesIO(data)) + return torch.from_numpy(npz["x"]), int(npz["y"])