"""Read an XProf trace: worked experiments and reference solutions. CPU checks.""" # Prepare a traceable workload # Step 1 — Prepare a traceable workload: Inputs and compilation precede the measured steps. import gzip import json import os from pathlib import Path import tempfile import time import numpy as np import jax import jax.numpy as jnp # Read `trace_root` from environment configuration (with a default fallback). trace_root=Path(os.environ.get('JAX_COURSE_TRACE_DIR') or tempfile.mkdtemp(prefix='jax-course-trace-')).resolve() trace_root.mkdir(parents=True,exist_ok=True) # Draw pseudorandom samples for `rng` using the explicit RNG state. rng=np.random.default_rng(31) host=rng.normal(size=(48,32)).astype(np.float32) weights=rng.normal(size=(32,16)).astype(np.float32) # Create device-backed JAX array `(x, w)`. x,w=jnp.asarray(host),jnp.asarray(weights) # Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function. compiled=jax.jit(lambda a,b:jnp.tanh(a@b)) # Synchronize host execution until asynchronous device computation completes. compiled(x,w).block_until_ready() # Capture named intervals # Step 2 — Capture named intervals: The outer annotation provides context, while two nonoverlapping... # Enter `jax.profiler.trace(trace_root, create_perfett` context block: with jax.profiler.trace(trace_root,create_perfetto_trace=True): for index in range(4): with jax.profiler.StepTraceAnnotation('learner_step',step_num=index): with jax.profiler.TraceAnnotation('input_wait'): time.sleep(.001) # injected host delay; not a measured real input pipeline with jax.profiler.TraceAnnotation('model_and_wait'): result=compiled(x,w) # Synchronize host execution until asynchronous device computation completes. result.block_until_ready() np.testing.assert_allclose(result,np.tanh(host@weights),rtol=2e-5,atol=2e-5) traces=sorted(trace_root.rglob('perfetto_trace.json.gz'),key=lambda p:p.stat().st_mtime_ns) assert traces and list(trace_root.rglob('*.xplane.pb')) events=json.loads(gzip.decompress(traces[-1].read_bytes()))['traceEvents'] def complete(name): return sorted([e for e in events if e.get('ph')=='X' and e.get('name')==name],key=lambda e:e['ts']) steps=complete('learner_step') input_events=complete('input_wait') model_events=complete('model_and_wait') assert len(steps)==len(input_events)==len(model_events)==4 for outer,waiting,compute in zip(steps,input_events,model_events): for inner in (waiting,compute): assert outer['ts']<=inner['ts']<=inner['ts']+inner['dur']<=outer['ts']+outer['dur']+1 assert waiting['ts']+waiting['dur']<=compute['ts']+1 input_us=[e['dur'] for e in input_events] model_us=[e['dur'] for e in model_events] step_us=[e['dur'] for e in steps] assert all(t>=0 for t in input_us+model_us+step_us) print('Trace directory:',trace_root) print(json.dumps({'input_wait_us':input_us,'model_and_wait_us':model_us,'step_us':step_us},indent=2)) print('Four named steps verified; host annotations are not isolated device kernels.') # Step 3: Verify invariants on the completed state step_us=[e['dur'] for e in steps] assert all(t>=0 for t in input_us+model_us+step_us) # Step 1 — Prepare a traceable workload: Inputs and compilation precede the measured steps. import gzip import json import os from pathlib import Path import tempfile import time import numpy as np import jax import jax.numpy as jnp # Read `trace_root` from environment configuration (with a default fallback). trace_root=Path(os.environ.get('JAX_COURSE_TRACE_DIR') or tempfile.mkdtemp(prefix='jax-course-trace-')).resolve() trace_root.mkdir(parents=True,exist_ok=True) # Draw pseudorandom samples for `rng` using the explicit RNG state. rng=np.random.default_rng(31) host=rng.normal(size=(48,32)).astype(np.float32) weights=rng.normal(size=(32,16)).astype(np.float32) # Create device-backed JAX array `(x, w)`. x,w=jnp.asarray(host),jnp.asarray(weights) # Wrap with `jax.jit` (`compiled`) so XLA traces and compiles the function. compiled=jax.jit(lambda a,b:jnp.tanh(a@b)) # Synchronize host execution until asynchronous device computation completes. compiled(x,w).block_until_ready() # Step 2 — Capture named intervals: The outer annotation provides context, while two nonoverlapping... # Enter `jax.profiler.trace(trace_root, create_perfett` context block: with jax.profiler.trace(trace_root,create_perfetto_trace=True): for index in range(4): with jax.profiler.StepTraceAnnotation('learner_step',step_num=index): with jax.profiler.TraceAnnotation('input_wait'): time.sleep(.001) # injected host delay; not a measured real input pipeline with jax.profiler.TraceAnnotation('model_and_wait'): result=compiled(x,w) # Synchronize host execution until asynchronous device computation completes. result.block_until_ready() np.testing.assert_allclose(result,np.tanh(host@weights),rtol=2e-5,atol=2e-5) traces=sorted(trace_root.rglob('perfetto_trace.json.gz'),key=lambda p:p.stat().st_mtime_ns) assert traces and list(trace_root.rglob('*.xplane.pb')) events=json.loads(gzip.decompress(traces[-1].read_bytes()))['traceEvents'] def complete(name): return sorted([e for e in events if e.get('ph')=='X' and e.get('name')==name],key=lambda e:e['ts']) steps=complete('learner_step') input_events=complete('input_wait') model_events=complete('model_and_wait') assert len(steps)==len(input_events)==len(model_events)==4 for outer,waiting,compute in zip(steps,input_events,model_events): for inner in (waiting,compute): assert outer['ts']<=inner['ts']<=inner['ts']+inner['dur']<=outer['ts']+outer['dur']+1 assert waiting['ts']+waiting['dur']<=compute['ts']+1 input_us=[e['dur'] for e in input_events] model_us=[e['dur'] for e in model_events] step_us=[e['dur'] for e in steps] assert all(t>=0 for t in input_us+model_us+step_us) print('Trace directory:',trace_root) print(json.dumps({'input_wait_us':input_us,'model_and_wait_us':model_us,'step_us':step_us},indent=2)) print('Four named steps verified; host annotations are not isolated device kernels.') # Figure data experiment # Compute figure data for: Where four annotated steps spent host time visual_data={'kind':'line','x':[0,1,2,3],'xlabel':'captured step','ylabel':'host annotation duration (microseconds)','series':[{'label':'injected input wait','y':input_us},{'label':'model call and completion wait','y':model_us}]} # Experiment: Account for the enclosing step # Experiment — Account for the enclosing step: Outer and child intervals describe nested boundaries. overhead=np.asarray(step_us)-np.asarray(input_us)-np.asarray(model_us) assert np.all(overhead>=-2) # Convert `wrong` to a host NumPy array for inspection or verification. wrong=np.asarray(step_us)+np.asarray(input_us)+np.asarray(model_us) assert np.all(wrong>=np.asarray(step_us)) print('Unattributed outer time, microseconds:',overhead.tolist()) print('Wrong double-counted totals:',wrong.tolist()) # Reference solution. Try the exercise before reading this. # Exercise solution: Convert every model interval to milliseconds and summarize its... model_ms=np.asarray(model_us)/1000 # Aggregate array values to compute `summary`. summary={'minimum_ms':float(model_ms.min()),'median_ms':float(np.median(model_ms)),'maximum_ms':float(model_ms.max())} assert summary['minimum_ms']<=summary['median_ms']<=summary['maximum_ms'] # Execute `np.testing.assert_allclose(model_ms*1000,model_us)` np.testing.assert_allclose(model_ms*1000,model_us) print(summary) # Reference practice: Compute a union instead of summing overlap # Compute a union instead of summing overlap (Transfer): A union describes time covered by any interval. def union_duration(intervals): total=0 left=right=None for start,end in sorted(intervals): # Guard input contract (`end < start`) and fail fast if violated. if end