Mschaara/pytorch tutorial

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)"

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_<k>/sample_<i>.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 = "/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())

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()

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()

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()