remove getenv(CI) from core tinygrad (#16326)

examples/mixtral.py

@@ -1,6 +1,6 @@
 import functools, argparse, pathlib
 from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
-from tinygrad.helpers import Timing, Profiling, CI, tqdm
+from tinygrad.helpers import Timing, Profiling, tqdm
 from tinygrad.nn.state import torch_load, get_state_dict
 from extra.models.llama import FeedForward, Transformer
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -36,7 +36,7 @@ if __name__ == "__main__":
     model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
     model_state_dict = get_state_dict(model)
 
-    for k in (t := tqdm(state, disable=CI)):
+    for k in (t := tqdm(state, disable=None)):
       if 'feed_forward.experts.' in k:
         expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
         device = Device.DEFAULT + ":" + str((expert_no//2)+1)
@@ -44,7 +44,7 @@ if __name__ == "__main__":
         device = Device.DEFAULT
       t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
       model_state_dict[k].replace(state[k].to(device).half()).realize()
-  if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+  if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
 
   from sentencepiece import SentencePieceProcessor
   spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")

extra/optimization/test_beam_search.py

@@ -1,7 +1,8 @@
 import unittest
 import numpy as np
 
-from tinygrad.helpers import BEAM, Timing, CI, prod
+from test.helpers import CI
+from tinygrad.helpers import BEAM, Timing, prod
 from tinygrad import Variable, Device, Tensor
 from tinygrad.nn import Conv2d
 from tinygrad.uop.ops import AxisType, Ops

extra/training.py

@@ -1,6 +1,6 @@
 import numpy as np
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import CI, trange
+from tinygrad.helpers import trange
 from tinygrad.engine.jit import TinyJit
 
 
@@ -24,7 +24,7 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: ou
 
   with Tensor.train():
     losses, accuracies = [], []
-    for i in (t := trange(steps, disable=CI)):
+    for i in (t := trange(steps, disable=None)):
       samp = np.random.randint(0, X_train.shape[0], size=(BS))
       x = Tensor(transform(X_train[samp]))
       y = Tensor(target_transform(Y_train[samp]))
@@ -43,7 +43,7 @@ def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=Fal
   Tensor.training = False
   def numpy_eval(Y_test, num_classes):
     Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
-    for i in trange((len(Y_test)-1)//BS+1, disable=CI):
+    for i in trange((len(Y_test)-1)//BS+1, disable=None):
       x = Tensor(transform(X_test[i*BS:(i+1)*BS]))
       out = model.forward(x) if hasattr(model, 'forward') else model(x)
       Y_test_preds_out[i*BS:(i+1)*BS] = out.numpy()

test/backend/test_dtype.py

@@ -2,13 +2,13 @@ import contextlib, unittest, math
 import numpy as np
 import torch
 from typing import Any, List
-from tinygrad.helpers import getenv, DEBUG, CI, EMULATED_DTYPES
+from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
-from test.helpers import rand_for_dtype
+from test.helpers import rand_for_dtype, CI
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")

test/backend/test_dtype_alu.py

@@ -1,12 +1,13 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate, fp8_to_float
-from tinygrad.helpers import CI, EMULATED_DTYPES, DEV, getenv
+from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad.uop import Ops
+from test.helpers import CI
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings

test/backend/test_interop.py

@@ -3,10 +3,11 @@ import unittest
 import torch
 import numpy as np
 
-from tinygrad.helpers import CI, DEV
+from tinygrad.helpers import DEV
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
+from test.helpers import CI
 
 MOCKGPU = DEV.interface.startswith("MOCK")
 

test/backend/test_ops.py

@@ -2,11 +2,12 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, CI, DEBUG, DEV, IMAGE, Context
+from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer
+from test.helpers import CI
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:

test/backend/test_profiler.py

@@ -1,10 +1,11 @@
 import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
+from test.helpers import CI
 
 MOCKGPU = DEV.interface.startswith("MOCK")
 def _dev_base(d):

test/backend/test_randomness.py

@@ -2,14 +2,14 @@ import unittest, math
 from functools import partial
 
 from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable
-from tinygrad.helpers import getenv, CI, OSX
+from tinygrad.helpers import getenv, OSX
 from tinygrad.codegen import to_program
 
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad.renderer.isa.x86 import X86Renderer
-from test.helpers import not_support_multi_device, needs_second_gpu
+from test.helpers import not_support_multi_device, needs_second_gpu, CI
 
 import numpy as np
 import torch

test/backend/test_schedule.py

@@ -10,8 +10,9 @@ from hypothesis import assume, given, strategies as strat
 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Ops, UPat
-from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
+from tinygrad.helpers import DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.engine.realize import compile_linear, run_linear
+from test.helpers import CI
 
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
 

test/backend/test_transcendental.py

@@ -1,7 +1,8 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.helpers import Context, getenv, CI, DEV, OSX
+from tinygrad.helpers import Context, getenv, DEV, OSX
+from test.helpers import CI
 from test.backend.test_schedule import check_schedule
 from test.backend.test_dtype_alu import ht, dtypes_float
 import numpy as np

test/backend/test_uops.py

@@ -2,7 +2,7 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import CI, Context
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
@@ -11,7 +11,7 @@ from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import to_uops_list
+from test.helpers import to_uops_list, CI
 
 def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
   buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]

test/external/external_test_example.py

@@ -1,7 +1,8 @@
 import unittest
 from tinygrad import Device
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, CI, OSX
+from tinygrad.helpers import getenv, OSX
+from test.helpers import CI
 
 def multidevice_test(fxn):
   exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")

test/external/external_test_hcq.py

@@ -1,9 +1,10 @@
 import unittest, ctypes, struct, time, array
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import to_mv, CI
+from tinygrad.helpers import to_mv
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
+from test.helpers import CI
 
 def _time_queue(q, d):
   st = time.perf_counter()

test/external/external_test_jit_on_models.py

@@ -3,8 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
 from tinygrad.engine.jit import TinyJit
-from tinygrad.helpers import CI
-from test.helpers import derandomize_model
+from test.helpers import derandomize_model, CI
 
 from examples.llama import Transformer
 

test/helpers.py

@@ -8,11 +8,14 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
 from tinygrad.dtype import DType
 from tinygrad.nn.state import get_parameters
-from tinygrad.helpers import T, CI, Target
+from tinygrad.helpers import T, Target
 from tinygrad.renderer import Renderer
 from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
 from tinygrad.codegen.late.linearizer import linearize
 
+# TODO: remove this everywhere!
+CI = os.getenv("CI", "") != ""
+
 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
 from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler

test/models/test_end2end.py

@@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.nn import optim, Linear, Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
 from extra.datasets import fetch_mnist
-from tinygrad.helpers import CI
+from test.helpers import CI
 
 def compare_tiny_torch(model, model_torch, X, Y):
   with Tensor.train():

test/null/test_device.py

@@ -3,8 +3,9 @@ import unittest, os, subprocess
 from unittest.mock import patch
 from tinygrad import Tensor
 from tinygrad.device import Device, Compiler, enumerate_devices_str
-from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, CI, OSX, DEV
+from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, OSX, DEV
 from tinygrad.runtime.support.c import DLL
+from test.helpers import CI
 
 class TestDevice(unittest.TestCase):
   def test_canonicalize(self):

test/null/test_winograd.py

@@ -1,6 +1,7 @@
 import unittest, sys
 from tinygrad import Tensor, GlobalCounters, dtypes, Context
-from tinygrad.helpers import CI, Profiling, WINO
+from tinygrad.helpers import Profiling, WINO
+from test.helpers import CI
 
 @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinograd(unittest.TestCase):

test/speed/external_test_copy_speed.py

@@ -1,6 +1,7 @@
 import unittest, numpy as np
 from tinygrad import Tensor, Device, TinyJit
-from tinygrad.helpers import Timing, CI, OSX, getenv
+from tinygrad.helpers import Timing, OSX, getenv
+from test.helpers import CI
 import multiprocessing.shared_memory as shared_memory
 
 N = getenv("NSZ", 256)

test/speed/external_test_specific_conv.py

@@ -1,6 +1,6 @@
 import unittest
-from tinygrad.helpers import CI
 from tinygrad import Tensor, Device, dtypes
+from test.helpers import CI
 # similar to test/external/external_test_gpu_ast.py, but universal
 
 @unittest.skipIf(Device.DEFAULT in {"CUDA", "NV"} and CI, "slow on CUDA CI")

test/speed/external_test_speed_v_torch.py

@@ -12,7 +12,8 @@ import numpy as np
 np.set_printoptions(linewidth=160)
 from tinygrad import Tensor, Device, GlobalCounters, TinyJit
 from tinygrad.nn import Conv2d
-from tinygrad.helpers import colorize_float, getenv, CI, DEV
+from tinygrad.helpers import colorize_float, getenv, DEV
+from test.helpers import CI
 
 IN_CHANS = [int(x) for x in getenv("IN_CHANS", "4,16,64").split(",")]
 

test/test_tiny.py

@@ -1,7 +1,8 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
 from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
-from tinygrad.helpers import CI, getenv
+from tinygrad.helpers import getenv
+from test.helpers import CI
 
 class TestTiny(unittest.TestCase):
 

test/testextra/test_bench_log.py

@@ -2,9 +2,10 @@ import unittest, time
 from unittest.case import skipIf
 
 from extra.bench_log import BenchEvent, InstantBenchEvent, WallTimeEvent, KernelTimeEvent, log_event_instant, _events, clear_events
-from tinygrad.helpers import Context, CI
+from tinygrad.helpers import Context
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
+from test.helpers import CI
 
 _SKIP_KERNEL_TIMING = Device.DEFAULT == "WEBGPU"  # WEBGPU kernel timing not supported
 

test/unit/test_assign.py

@@ -3,7 +3,8 @@ import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import temp, CI, DEV, Context
+from tinygrad.helpers import temp, DEV, Context
+from test.helpers import CI
 
 N = 200  # has to be bigger than the cache to fail
 

test/unit/test_shm_tensor.py

@@ -1,7 +1,8 @@
 import unittest
 import multiprocessing.shared_memory as shared_memory
-from tinygrad.helpers import CI, WIN
+from tinygrad.helpers import WIN
 from tinygrad import Tensor, Device
+from test.helpers import CI
 import numpy as np
 
 class TestRawShmBuffer(unittest.TestCase):

tinygrad/helpers.py

@@ -14,7 +14,6 @@ def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)
 
 # NOTE: helpers is not allowed to import from anything else in tinygrad
 OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
-CI = os.getenv("CI", "") != ""
 ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
 BASEDIR = pathlib.Path(__file__).parent
 
@@ -462,7 +461,7 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
       assert r.status in {200, 206}, r.status
       length = int(r.headers.get('content-length', 0)) if not gunzip else None
       readfile = gzip.GzipFile(fileobj=r) if gunzip else r
-      progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}", disable=CI)
+      progress_bar:tqdm = tqdm(total=length, unit='B', unit_scale=True, desc=f"{url}")
       h = hashlib.sha256() if sha256 else None
       with tempfile.NamedTemporaryFile(dir=_dir, delete=False) as f:
         while chunk := readfile.read(16384):
@@ -531,9 +530,10 @@ def flat_mv(mv:memoryview): return mv if len(mv) == 0 else mv.cast("B", shape=(m
 # *** tqdm
 
 class tqdm(Generic[T]):
-  def __init__(self, iterable:Iterable[T]|None=None, desc:str='', disable:bool=False,
+  def __init__(self, iterable:Iterable[T]|None=None, desc:str='', disable:bool|None=False,
                unit:str='it', unit_scale=False, total:int|None=None, rate:int=100):
-    self.iterable, self.disable, self.unit, self.unit_scale, self.rate = iterable, disable, unit, unit_scale, rate
+    self.disable = not sys.stderr.isatty() if disable is None else disable
+    self.iterable, self.unit, self.unit_scale, self.rate = iterable, unit, unit_scale, rate
     self.st, self.i, self.n, self.skip, self.t = time.perf_counter(), -1, 0, 1, getattr(iterable, "__len__", lambda:0)() if total is None else total
     self.set_description(desc)
     self.update(0)

tinygrad/nn/state.py

@@ -3,7 +3,7 @@ from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, GlobalCounters, tqdm, round_up, T, strides_for_shape
+from tinygrad.helpers import prod, argsort, DEBUG, Timing, GlobalCounters, tqdm, round_up, T, strides_for_shape
 
 class TensorIO(io.RawIOBase, BinaryIO):
   def __init__(self, t: Tensor):
@@ -145,7 +145,7 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
     model_state_dict = get_state_dict(model)
     if DEBUG >= 1 and len(state_dict) > len(model_state_dict):
       print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys())))
-    for k,v in (t := tqdm(model_state_dict.items(), disable=CI or not verbose)):
+    for k,v in (t := tqdm(model_state_dict.items(), disable=None if verbose else True)):
       t.desc = f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}: "
       if k not in state_dict and not strict:
         if DEBUG >= 1: print(f"WARNING: not loading {k}")