script.py

import os
import subprocess
import tempfile
import contextlib
import secrets
import glob
import shutil
import time
from argparse import ArgumentParser
from functools import partial

LIMIT = 32

def call(command, *args):
    kwargs = {}
    if disable_noise:
        kwargs['stdout'] = subprocess.DEVNULL
        kwargs['stderr'] = subprocess.DEVNULL
    subprocess.check_call([command, *args], **kwargs)

dvc = partial(call, "dvc")
git = partial(call, "git")

@contextlib.contextmanager
def temp_dir():
    current_dir = os.getcwd()
    try:
        with tempfile.TemporaryDirectory() as directory:
            yield os.chdir(directory)
    finally:
        os.chdir(current_dir)

@contextlib.contextmanager
def timed(settings):
    start = time.perf_counter()
    try:
        yield
    finally:
        end = time.perf_counter()
    
    print(f"Time [{settings}]: ", end - start)

def create_stages(num_files):
    for index in range(num_files):
        file = f"{index}.txt"
        with open(file, "w") as stream:
            stream.write(str(index))
        
        dvc("add", file)
    
    dvc("push")

def clear_stages():
    for file in glob.iglob("*.txt"):
        os.unlink(file)
    
    shutil.rmtree(".dvc/cache")

@temp_dir()
def new_env(num_files, repeat=1):
    test_id = f"{num_files}-{secrets.token_hex(16)}"

    git("init")
    dvc("init")
    dvc("remote", "add", "-d", "origin", f"s3://dvc-temp/batuhan/test-{test_id}")
    create_stages(num_files)
    clear_stages()

    for num_times in range(repeat):
        with timed(f"{num_files=}"):
            dvc("status", "-c")

def main():
    global disable_noise

    parser = ArgumentParser()
    parser.add_argument("num_files", type=int)
    parser.add_argument("--repeat", type=int, default=1)
    parser.add_argument("--disable-noise", action="store_true")
    parser.add_argument("--enable-ascending", action="store_true")

    options = parser.parse_args()

    disable_noise = options.disable_noise
    
    if options.enable_ascending:
        limit = LIMIT
    else:
        limit = options.num_files + 1

    for num_files in range(options.num_files, limit):
        new_env(num_files, repeat=options.repeat)

if __name__ == "__main__":
    main()