diff --git a/examples/mlperf/dataloader.py b/examples/mlperf/dataloader.py index 278b7ec6ab..275bcd4aef 100644 --- a/examples/mlperf/dataloader.py +++ b/examples/mlperf/dataloader.py @@ -1,11 +1,11 @@ import os, random, pickle, queue, struct, math, functools, hashlib, time from typing import List from pathlib import Path -from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count +from multiprocessing import Queue, Process, shared_memory, connection, Lock import numpy as np from tinygrad import dtypes, Tensor -from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX +from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS from tinygrad.nn.state import TensorIO ### ResNet @@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}") Y = [None] * (batch_size*BATCH_COUNT) - for _ in range(cpu_count()): + for _ in range(NUM_CPU_THREADS.value): p = Process(target=loader_process, args=(q_in, q_out, X, seed)) p.daemon = True p.start() @@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None): rng.shuffle(fs) train_files.append(fs.pop(0)) - cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files)) + cycle_length = min(NUM_CPU_THREADS.value, len(train_files)) assert cycle_length > 0, "cycle_length must be greater than 0" dataset = InterleavedDataset(train_files, cycle_length) @@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool= X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}") Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}") - for _ in range(cpu_count()): + for _ in range(NUM_CPU_THREADS.value): proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y)) proc.daemon = True proc.start() @@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh dataset_iter = iter(image_ids) try: - for _ in range(cpu_count()): + for _ in range(NUM_CPU_THREADS.value): proc = Process( target=load_retinanet_data, args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels), diff --git a/extra/export_model.py b/extra/export_model.py index 4b9e4a6c09..a9050b77e3 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -241,8 +241,8 @@ export default {model_name}; def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False): assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported" - # NOTE: CPU_COUNT=1, since export does not support threading - with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs) + # NOTE: NUM_CPU_THREADS=1, since export does not support threading + with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs) functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs) state = get_state_dict(model) weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None} diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index b4bd8b99f2..5b7e9891a8 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -248,7 +248,18 @@ ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), Conte MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0) EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "") CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0) -CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))) +def _get_cpu_count() -> int: + # os.process_cpu_count (3.13+) respects cgroup limits + if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count()) + # cgroup v2 (containers with --cpus=N) + try: + with open("/sys/fs/cgroup/cpu.max") as f: + quota, period = f.read().strip().split() + if quota != "max": return max(1, int(quota) // int(period)) + except (FileNotFoundError, ValueError, ZeroDivisionError): pass + # fall back to affinity (respects taskset but not cgroup quota) + return max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)) +NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count()) NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0) # VIZ implies PROFILE, but you can run PROFILE without VIZ VIZ = ContextVar("VIZ", 0) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 110b210c6f..0adde67841 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -3,7 +3,7 @@ import math, sys, struct from collections import defaultdict, Counter from tinygrad.codegen.opt import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters -from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, CPU_COUNT, IMAGE, FLOAT16, is_image_shape +from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_THREADS, IMAGE, FLOAT16, is_image_shape from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16 from tinygrad.renderer import Renderer @@ -258,7 +258,7 @@ class ClangRenderer(CStyleLanguage): gep_arr_threshold = 0 has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (CPU_COUNT.value, 0, 0) + global_max = (NUM_CPU_THREADS.value, 0, 0) infinity = "__builtin_inff()" nan = '__builtin_nanf("")' diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 6ded880f02..5f3597e18d 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, DType, truncate, AddrSpace from tinygrad.uop import FastEnum, auto, Ops, GroupOp from tinygrad.uop.ops import UOp, UPat, PatternMatcher from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg -from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target +from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target # ***** X86 Ops ***** @@ -804,7 +804,7 @@ class X86Renderer(ISARenderer): device = "CPU" has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (CPU_COUNT.value, 0, 0) + global_max = (NUM_CPU_THREADS.value, 0, 0) extra_matcher = extra_matcher pre_isel_matcher = pre_isel_matcher isel_matcher = isel_matcher diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 7bb6c23496..81eb6453d0 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -5,7 +5,7 @@ from tinygrad.renderer.cstyle import HIPRenderer, create_non_native_float_pats, from tinygrad.codegen.decomp.transcendental import xexp2, xlog2 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str from tinygrad.dtype import dtypes, float_to_fp8, DType, truncate, AddrSpace -from tinygrad.helpers import prod, Target, CPU_COUNT, getenv, OSX +from tinygrad.helpers import prod, Target, NUM_CPU_THREADS, getenv, OSX def is_volatile(u:UOp) -> bool: return (buf:=u.buf_uop).op is Ops.PARAM and buf.arg.volatile @@ -189,7 +189,7 @@ class LLVMRenderer(Renderer): class CPULLVMRenderer(LLVMRenderer): has_local = False has_threads = bool(getenv("THREADS", 1)) - global_max = (CPU_COUNT.value, 0, 0) + global_max = (NUM_CPU_THREADS.value, 0, 0) abi = 'win64cc' if sys.platform == 'win32' else None string_rewrite = base_rewrite def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))