Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment