shinka-backup / ccevolve /baselines /ttt-discover /examples /gpu_mode /lib /libkernelbot /run_eval.py
| import base64 | |
| import copy | |
| import dataclasses | |
| import datetime | |
| import functools | |
| import json | |
| import os | |
| import shlex | |
| import shutil | |
| import socket | |
| import subprocess | |
| import tempfile | |
| import time | |
| from pathlib import Path | |
| from types import NoneType | |
| from typing import Optional, Protocol, Union | |
| from libkernelbot.consts import CUDA_FLAGS, ExitCode, Timeout | |
| class ProfileResult: | |
| # fmt: off | |
| profiler: str # The profiler used to gather this data | |
| # Profiler trace. May be empty, in which case `download_url` | |
| # should point to the trace file. | |
| trace: str | |
| # Public download URL of all files created by the profiler | |
| # This may also be configured later | |
| download_url: Optional[str] | |
| # fmt: on | |
| class CompileResult: | |
| # fmt: off | |
| nvcc_found: bool # did we find nvcc? | |
| nvcc_version: str # the result of nvcc --version | |
| success: bool # did it compile successfully | |
| command: str # the command that was run to compile the code | |
| stdout: str # standard output produced by the compiler | |
| stderr: str # standard error produced by the compiler | |
| exit_code: int # exit code produced by the compiler | |
| # fmt: on | |
| class RunResult: | |
| # fmt: off | |
| success: bool # did the compiled executable run successfully | |
| passed: bool # did it pass all tests | |
| command: str # the command that was run to compile the code | |
| stdout: str # standard output produced by the compiler | |
| stderr: str # standard error produced by the compiler | |
| exit_code: int # exit code produced by the compiler | |
| duration: float # execution time (NOT kernel duration) | |
| result: dict # dictionary with the results generated by the tester | |
| # fmt: on | |
| class SystemInfo: | |
| # fmt: off | |
| gpu: str = '' # Model name of the GPU | |
| device_count: int = 1 # Number of GPUs | |
| cpu: str = '' # Model name of the CPU | |
| runtime: str = '' # Whether CUDA or ROCm | |
| platform: str = '' # Platform string of the machine | |
| torch: str = '' # Torch version | |
| hostname: str = '' # Hostname of the machine | |
| requeues: int = 0 # How many Modal requeues/retries were used before running | |
| # fmt: on | |
| class EvalResult: | |
| # fmt: off | |
| start: datetime.datetime # when did this run start (excluding container setup time) | |
| end: datetime.datetime # and when did it finish | |
| compilation: CompileResult | None # results of compilation | |
| run: RunResult | None # result of actually running the executable/script | |
| profile: ProfileResult | None # result of profiling the executable | |
| # fmt: on | |
| class FullResult: | |
| # fmt: off | |
| success: bool # did the runner (github/modal) execute successfully | |
| error: str # if not success, an error message | |
| system: SystemInfo # specs of the system this was run on | |
| # results of running. There can be multiple runs in one submission, using separate | |
| # 'test' and 'benchmark' keys, for example | |
| runs: dict[str, EvalResult] = dataclasses.field(default_factory=dict) | |
| # fmt: on | |
| def _make_cmd(args: list[str]): | |
| return " ".join(map(shlex.quote, args)) | |
| def _limit_length(text: Union[NoneType, str, bytes], max_len: int = 16384): | |
| if text is None: | |
| return "" | |
| if isinstance(text, bytes): | |
| text = text.decode("utf-8") | |
| lines = text.split("\n") | |
| size = 0 | |
| for i, line in enumerate(lines): | |
| size += len(line) + 1 | |
| if size + 100 > max_len: | |
| lines = lines[:i] + [f"[...] {len(lines) - i} lines omitted"] | |
| return "\n".join(lines) | |
| return text | |
| def _create_files(files: Optional[dict[str, str]]): | |
| """ | |
| Create text files | |
| Args: | |
| files: A dictionary mapping file names to their contents. | |
| Raises: | |
| AssertionError, if the file is not within the current working directory. | |
| """ | |
| if files is None: | |
| return | |
| for name, content in files.items(): | |
| assert Path(name).resolve().is_relative_to(Path.cwd()) | |
| Path(name).write_text(content) | |
| def _directory_to_zip_bytes(directory_path) -> str: | |
| """Create a zip archive and return as base64 encoded bytes.""" | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| archive_path = os.path.join(temp_dir, "archive") | |
| shutil.make_archive(archive_path, "zip", directory_path) | |
| with open(archive_path + ".zip", "rb") as f: | |
| data = f.read() | |
| return base64.b64encode(data).decode("utf-8") | |
| def _filter_ncu_report(report: str, tables: list): # noqa: C901 | |
| """ | |
| Extract the Speed-of-light section from the full ncu terminal report. | |
| For expert users, we just attach the full ncu profile to the result, | |
| and they can view whichever metrics they are interested in. But to | |
| encourage novice users to try out profiling, we want to have a | |
| *simple* set of things to display automatically, short enough to fit | |
| in a *single* discord message. | |
| """ | |
| result = "" | |
| n_kernels = 0 | |
| collect = False | |
| length = 0 | |
| for line in report.splitlines(): | |
| if len(line) >= 3 and line[1] == " " and line[2] != " ": | |
| if n_kernels != 0: | |
| result += "\n" | |
| n_kernels += 1 | |
| if n_kernels == 3: | |
| result += "\nAdditional kernel launches follow. Please check the .ncu-rep file for more details.\n" # noqa: E501 | |
| result += line + "\n" | |
| if n_kernels > 2: | |
| continue | |
| if "Table Name : " in line: | |
| table = line[line.find("Table Name :") + len("Table Name :") :].strip() | |
| if table in tables: | |
| result += "\n" | |
| collect = True | |
| else: | |
| collect = False | |
| if len(line.strip()) == 0: | |
| collect = False | |
| if collect: | |
| result += line + "\n" | |
| length += 1 | |
| # just as a precaution, also limit lines directly | |
| if length > 100: | |
| result += "\n[...]\nReport has been truncated. Please check the .ncu-rep file for more details.\n" # noqa: E501 | |
| break | |
| return result | |
| def compile_cuda_script( # # noqa: C901 | |
| files: list[str], | |
| arch: Optional[int] = None, | |
| include_dirs: Optional[list[str]] = None, | |
| defines: Optional[dict[str, str]] = None, | |
| libraries: Optional[list[str]] = None, | |
| flags: Optional[list[str]] = None, | |
| verbose: bool = False, | |
| ) -> CompileResult: | |
| """ | |
| Compiles a set of cuda files with nvcc. | |
| Args: | |
| files: List of files to compile. | |
| arch: Architecture to compile for. If None, uses `native` | |
| include_dirs: additional include directories to supply to nvcc | |
| defines: Additional defines for the preprocessor | |
| libraries: Additional libraries to link to | |
| flags: Other compiler flags | |
| verbose: whether to print progress or be silent | |
| Returns: | |
| A `CompileResult` that summarizes the compilation process. | |
| """ | |
| if flags is None: | |
| flags = CUDA_FLAGS | |
| if include_dirs is not None: | |
| flags += [f"-I{d}" for d in include_dirs] | |
| # validate include directories | |
| for directory in include_dirs: | |
| if not Path(directory).exists(): | |
| raise FileNotFoundError(f"Directory `{directory}` does not exist") | |
| elif not Path(directory).is_dir(): | |
| raise NotADirectoryError(f"`{directory}` is not a directory") | |
| if libraries is not None: | |
| flags += [f"-l{lib}" for lib in libraries] | |
| if defines is not None: | |
| for name, value in defines.items(): | |
| # restrict macro names to valid identifiers | |
| if not name.isidentifier(): | |
| raise ValueError(f"Define key `{name}` contains invalid character") | |
| if value is not None: | |
| flags.append(f"-D{name}={value}") | |
| else: | |
| flags.append(f"-D{name}") | |
| for flag in flags: | |
| if not flag.startswith("-"): | |
| raise ValueError(f"Flag `{flag}` should start with a dash.") | |
| if verbose: | |
| print_ = print | |
| else: | |
| print_ = lambda *args, **kwargs: None # noqa | |
| # Check CUDA is available and installed correctly | |
| print_("[CUDA Env Check]") | |
| try: | |
| # these check cuda compiler is also available | |
| nvcc = subprocess.check_output(["which", "nvcc"], encoding="utf-8").strip() | |
| nvcc_version = subprocess.check_output(["nvcc", "--version"], encoding="utf-8") | |
| except subprocess.CalledProcessError as e: | |
| return CompileResult( | |
| nvcc_found=False, | |
| success=False, | |
| nvcc_version="", | |
| command=_make_cmd(e.cmd), | |
| stdout=_limit_length(e.stdout), | |
| stderr=_limit_length(e.stderr), | |
| exit_code=e.returncode, | |
| ) | |
| if arch is None: | |
| ARCH = "-arch=native" | |
| else: | |
| ARCH = f"-gencode=arch=compute_{arch},code=sm_{arch}" | |
| command = [nvcc] + flags + files + [ARCH, "-o", "eval.out"] | |
| print_("[Compiling]") | |
| try: | |
| compile_process = subprocess.run( | |
| command, capture_output=True, text=True, check=True, timeout=Timeout.COMPILE | |
| ) | |
| except subprocess.CalledProcessError as e: | |
| return CompileResult( | |
| nvcc_found=True, | |
| success=False, | |
| nvcc_version=nvcc_version, | |
| command=_make_cmd(e.cmd), | |
| stdout=_limit_length(e.stdout), | |
| stderr=_limit_length(e.stderr), | |
| exit_code=e.returncode, | |
| ) | |
| return CompileResult( | |
| nvcc_found=True, | |
| success=True, | |
| nvcc_version=nvcc_version, | |
| command=_make_cmd(compile_process.args), | |
| stdout=_limit_length(compile_process.stdout), | |
| stderr=_limit_length(compile_process.stderr), | |
| exit_code=compile_process.returncode, | |
| ) | |
| def run_program( | |
| args: list[str], | |
| seed: Optional[int], | |
| timeout: int, | |
| multi_gpu: bool = False, | |
| extra_env: Optional[dict[str, str]] = None, | |
| ) -> RunResult: | |
| print("[Running]") | |
| # set up a pipe so the tester can communicate its verdict with us | |
| env = os.environ.copy() | |
| if extra_env is not None: | |
| env.update(extra_env) | |
| pipe_read, pipe_write = os.pipe() | |
| env["POPCORN_FD"] = str(pipe_write) | |
| if seed is not None: | |
| env["POPCORN_SEED"] = str(seed) | |
| if multi_gpu: | |
| import torch | |
| env["POPCORN_GPUS"] = str(torch.cuda.device_count()) | |
| execution_start_time = time.perf_counter() | |
| try: | |
| run_process = subprocess.run( | |
| args, | |
| capture_output=True, | |
| text=True, | |
| check=False, | |
| env=env, | |
| pass_fds=[pipe_write], | |
| timeout=timeout, | |
| ) | |
| except subprocess.TimeoutExpired as e: | |
| return RunResult( | |
| success=False, | |
| passed=False, | |
| command=_make_cmd(e.cmd), | |
| stdout=_limit_length(e.stdout), | |
| stderr=_limit_length(e.stderr), | |
| exit_code=ExitCode.TIMEOUT_EXPIRED, | |
| duration=timeout, | |
| result={}, | |
| ) | |
| execution_end_time = time.perf_counter() | |
| # terminate output writing | |
| os.close(pipe_write) | |
| # and fetch pipe's content | |
| result = os.fdopen(pipe_read, "r").read() | |
| result_dict = {} | |
| for line in result.splitlines(): | |
| key, _, value = line.partition(":") | |
| if key != "" or value != "": | |
| result_dict[key.strip()] = value.strip() | |
| return RunResult( | |
| success=( | |
| run_process.returncode == ExitCode.SUCCESS | |
| or run_process.returncode == ExitCode.VALIDATE_FAIL | |
| ), | |
| passed=result_dict.get("check", None) == "pass", | |
| command=_make_cmd(run_process.args), | |
| stdout=_limit_length(run_process.stdout), | |
| stderr=_limit_length(run_process.stderr), | |
| exit_code=run_process.returncode, | |
| duration=execution_end_time - execution_start_time, | |
| result=result_dict, | |
| ) | |
| def profile_program_roc( | |
| call: list[str], | |
| seed: Optional[int], | |
| timeout: int, | |
| multi_gpu: bool, | |
| output_dir: Path, | |
| ) -> tuple[RunResult, Optional[ProfileResult]]: | |
| # Wrap program in rocprof | |
| call = [ | |
| "rocprofv3", | |
| "--log-level", | |
| "fatal", | |
| "--hip-trace", | |
| "--kernel-trace", | |
| "--rccl-trace", | |
| "--marker-trace", | |
| "--hip-trace", | |
| "--memory-copy-trace", | |
| # New? Doesn't work in the runner | |
| # "--memory-allocation-trace", | |
| "--scratch-memory-trace", | |
| # The HSA trace output is very large, so skip it for now | |
| # "--hsa-trace", | |
| "--output-format", | |
| "pftrace", | |
| "csv", | |
| "-d", | |
| str(output_dir), | |
| # Just store the files as %pid%_tracename.ext instead of putting them in an | |
| # additional directory named after the hostname. | |
| "-o", | |
| # Insert an extra path here so that the resulting zip has all files | |
| # in the profile_data/ directory rather than the root. | |
| "%pid%", | |
| "--", | |
| ] + call | |
| run_result = run_program( | |
| call, | |
| seed=seed, | |
| timeout=timeout, | |
| multi_gpu=multi_gpu, | |
| extra_env={ | |
| "GPU_DUMP_CODE_OBJECT": "1", | |
| }, | |
| ) | |
| profile_result = None | |
| if run_result.success: | |
| # Post-process trace data. | |
| # rocPROF generates one trace for every process, but its more useful to | |
| # have all traces be in the same file. Fortunately we can do that by | |
| # concatenating. | |
| traces = list(output_dir.glob("*.pftrace")) | |
| with (output_dir / "combined.pftrace").open("wb") as combined: | |
| for trace_path in traces: | |
| with trace_path.open("rb") as trace: | |
| shutil.copyfileobj(trace, combined) | |
| # After we've created the combined trace, there is no point in | |
| # keeping the individual traces around. | |
| trace_path.unlink() | |
| # Also move the code objects to the profiling output directory. | |
| for code_obj in list(Path.cwd().glob("_code_object*.o")): | |
| code_obj.rename(output_dir / code_obj.name) | |
| profile_result = ProfileResult( | |
| profiler="rocPROF", | |
| trace=_directory_to_zip_bytes(output_dir), | |
| download_url=None, | |
| ) | |
| return run_result, profile_result | |
| def profile_program_ncu( | |
| call: list[str], | |
| seed: Optional[int], | |
| timeout: int, | |
| multi_gpu: bool, | |
| output_dir: Path, | |
| ) -> tuple[RunResult, Optional[ProfileResult]]: | |
| assert not multi_gpu, "Multi-GPU profiling not supported for ncu." | |
| # Wrap program in ncu | |
| call = [ | |
| "ncu", | |
| "--set", | |
| "full", | |
| "--nvtx", | |
| "--nvtx-include", | |
| "custom_kernel/", | |
| "--import-source", | |
| "1", | |
| "-c", | |
| "10", | |
| "-o", | |
| f"{str(output_dir / 'profile.ncu-rep')}", | |
| "--", | |
| ] + call | |
| run_result = run_program( | |
| call, seed=seed, timeout=timeout, multi_gpu=multi_gpu, extra_env={"POPCORN_NCU": "1"} | |
| ) | |
| profile_result = None | |
| try: | |
| get_tables = [ | |
| "GPU Throughput", | |
| "Pipe Utilization (% of active cycles)", | |
| "Warp State (All Cycles)", | |
| ] | |
| ncu_cmd = [ | |
| "ncu", | |
| "--import", | |
| f"{str(output_dir / 'profile.ncu-rep')}", | |
| "--print-details", | |
| "body", | |
| ] | |
| report = subprocess.check_output(ncu_cmd, text=True) | |
| report = _filter_ncu_report(report, get_tables) | |
| run_result.result["benchmark.0.report"] = base64.b64encode(report.encode("utf-8")).decode( | |
| "utf-8" | |
| ) | |
| except subprocess.CalledProcessError: | |
| pass | |
| if run_result.success: | |
| profile_result = ProfileResult( | |
| profiler="Nsight-Compute", | |
| trace=_directory_to_zip_bytes(output_dir), | |
| download_url=None, | |
| ) | |
| return run_result, profile_result | |
| def profile_program( | |
| system: SystemInfo, | |
| call: list[str], | |
| seed: Optional[int], | |
| timeout: int, | |
| multi_gpu: bool, | |
| ) -> tuple[RunResult, Optional[ProfileResult]]: | |
| # The runner-specific configuration should implement logic | |
| # to fetch the data in this directory and return it as | |
| # ProfileResult.download_url. | |
| # Insert an extra nested path here so that the resulting zip has all files | |
| # in the profile_data/ directory rather than directly in the root. | |
| with tempfile.TemporaryDirectory(dir=".") as tmpdir: | |
| output_dir = Path(tmpdir) / "profile_data" | |
| output_dir.mkdir() | |
| if system.runtime == "ROCm": | |
| return profile_program_roc(call, seed, timeout, multi_gpu, output_dir) | |
| elif system.runtime == "CUDA": | |
| return profile_program_ncu(call, seed, timeout, multi_gpu, output_dir) | |
| else: | |
| raise ValueError(f"Unknown runtime {system.runtime}") | |
| def run_single_evaluation( | |
| call: list[str], | |
| mode: str, | |
| *, | |
| system: SystemInfo, | |
| multi_gpu: bool = False, | |
| tests: Optional[str] = None, | |
| benchmarks: Optional[str] = None, | |
| test_timeout: int = Timeout.TEST, | |
| benchmark_timeout: int = Timeout.BENCHMARK, | |
| ranked_timeout: int = Timeout.RANKED, | |
| ranking_by: str = "last", | |
| seed: Optional[int] = None, | |
| ) -> tuple[RunResult, Optional[ProfileResult]]: | |
| """ | |
| A single runner run, either in the context of test files, or in the | |
| context of benchmark files. | |
| """ | |
| with tempfile.NamedTemporaryFile("w") as cases: | |
| if mode == "test": | |
| timeout = test_timeout | |
| cases.write(tests) | |
| elif mode in ["benchmark", "profile", "leaderboard"]: | |
| timeout = ranked_timeout if mode == "leaderboard" else benchmark_timeout | |
| if ranking_by == "last": | |
| cases.write(benchmarks.splitlines(keepends=True)[-1]) | |
| else: | |
| cases.write(benchmarks) | |
| else: | |
| raise ValueError(f"Invalid mode {mode}") | |
| cases.flush() | |
| call = call + [mode, cases.name] | |
| if mode == "profile": | |
| return profile_program(system, call, seed=seed, timeout=timeout, multi_gpu=multi_gpu) | |
| return run_program(call, seed=seed, timeout=timeout, multi_gpu=multi_gpu), None | |
| def make_system_info() -> SystemInfo: # noqa: C901 | |
| info = SystemInfo() | |
| try: | |
| import torch | |
| info.torch = torch.torch_version.internal_version | |
| # Note: cuda.is_available() also covers HiP | |
| # https://pytorch.org/docs/stable/notes/hip.html | |
| if torch.cuda.is_available(): | |
| info.gpu = torch.cuda.get_device_name() | |
| info.device_count = torch.cuda.device_count() | |
| if torch.version.hip is not None: | |
| info.runtime = "ROCm" | |
| elif torch.version.cuda is not None: | |
| info.runtime = "CUDA" | |
| except ImportError: | |
| # get GPU info manually | |
| try: | |
| info.gpu = subprocess.check_output( | |
| ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], encoding="utf-8" | |
| ) | |
| info.device_count = info.gpu.count("\n") | |
| info.runtime = "CUDA" | |
| except subprocess.CalledProcessError: | |
| # try again for HIP | |
| try: | |
| rocm_info = json.loads( | |
| subprocess.check_output( | |
| ["rocm-smi", "--showproductname", "--json"], encoding="utf-8" | |
| ) | |
| ) | |
| if len(rocm_info) > 0: | |
| info.gpu = next(rocm_info.__iter__())["Card Series"] | |
| info.device_count = len(rocm_info) | |
| info.runtime = "ROCm" | |
| except subprocess.CalledProcessError: | |
| # OK, no GPU info available | |
| pass | |
| try: | |
| cpu_info_str = Path("/proc/cpuinfo").read_text() | |
| cpu_info_dict = {} | |
| for line in cpu_info_str.splitlines(): | |
| key, _, val = line.partition(":") | |
| cpu_info_dict[key.strip()] = val.strip() | |
| info.cpu = cpu_info_dict.get("model name", "") | |
| # on modal, we don't get to know the exact CPU model | |
| # make due with the vendor in that case | |
| if info.cpu == "unknown": | |
| # ¯\_(ツ)_/¯ | |
| info.cpu = cpu_info_dict.get("vendor_id", "") | |
| except PermissionError: | |
| # nothing we can do here; we're not getting CPU info | |
| pass | |
| import platform | |
| info.hostname = socket.gethostname() | |
| info.platform = platform.platform() | |
| return info | |
| def run_cuda_script( # # noqa: C901 | |
| sources: dict[str, str], | |
| headers: Optional[dict[str, str]] = None, | |
| arch: Optional[int] = None, | |
| defines: Optional[dict[str, str]] = None, | |
| include_dirs: Optional[list[str]] = None, | |
| libraries: Optional[list[str]] = None, | |
| flags: Optional[list[str]] = None, | |
| **kwargs, | |
| ) -> EvalResult: | |
| """ | |
| Executes the provided CUDA kernel in an isolated environment | |
| Args: | |
| sources: The source files to compile. Mapping file name to content. | |
| headers: Additional header files to create for the compile run. | |
| Mapping of file name to file contents. These files will _not_ be added to the | |
| compile command. | |
| arch: The arch code for the compute/sm versions. If None, native arch is used. | |
| include_dirs: Additional include directories, e.g., for thunderkittens/cutlass etc | |
| defines: Preprocessor defines | |
| libraries: Additional libraries to link to | |
| flags: Additional flags to give to the compiler | |
| seed: Random seed to initialize the RNG for testing | |
| Returns: | |
| tuple[CompileResult, RunResult]: CUDA compile/eval result information | |
| """ | |
| start = datetime.datetime.now() | |
| try: | |
| # Write submission files to directory | |
| _create_files(sources) | |
| _create_files(headers) | |
| compile_result = compile_cuda_script( | |
| files=list(sources.keys()), | |
| arch=arch, | |
| include_dirs=include_dirs, | |
| defines=defines, | |
| libraries=libraries, | |
| flags=flags, | |
| verbose=True, | |
| ) | |
| if not compile_result.success: | |
| return EvalResult( | |
| start=start, | |
| end=datetime.datetime.now(), | |
| compilation=compile_result, | |
| run=None, | |
| profile=None, | |
| ) | |
| # cleaning up all source files _before_ we let the user code run, just in | |
| # case there's something in there that the user isn't supposed to snoop | |
| finally: | |
| tmp_files = list(sources.keys()) + list((headers or {}).keys()) | |
| for f in tmp_files: | |
| if os.path.exists(f): | |
| os.remove(f) | |
| run_result, profile_result = run_single_evaluation(["./eval.out"], **kwargs) | |
| return EvalResult( | |
| start=start, | |
| end=datetime.datetime.now(), | |
| compilation=compile_result, | |
| run=run_result, | |
| profile=profile_result, | |
| ) | |
| def run_pytorch_script( # noqa: C901 | |
| sources: dict[str, str], | |
| main: str, | |
| **kwargs, | |
| ) -> EvalResult: | |
| """ | |
| Executes the provided PyTorch GPU kernel in an isolated environment | |
| Args: | |
| sources: Files to generate | |
| main: Which file to run. Must be one of the keys in sources. | |
| seed: Random seed to initialize the RNG for testing | |
| Returns: | |
| RunResult | |
| """ | |
| start = datetime.datetime.now() | |
| try: | |
| assert main in sources.keys() | |
| # Write submission files to directory | |
| _create_files(sources) | |
| # "compile" step: execute the script once. Will populate | |
| # `load_inline`'s compile cache, so the actual runs will be faster. | |
| try: | |
| compile_run = run_program(["python3", "submission.py"], seed=1, timeout=Timeout.COMPILE) | |
| if "-DTORCH_EXTENSION_NAME" in compile_run.stdout: | |
| comp = CompileResult( | |
| nvcc_found=True, | |
| nvcc_version="", | |
| success=True, | |
| command=compile_run.command, | |
| stdout=compile_run.stdout, | |
| stderr=compile_run.stderr, | |
| exit_code=compile_run.exit_code, | |
| ) | |
| else: | |
| comp = None | |
| except subprocess.CalledProcessError as e: | |
| # This step is purely optional, so we just go on | |
| # if it fails | |
| comp = CompileResult( | |
| nvcc_found=False, | |
| nvcc_version="", | |
| success=False, | |
| command="python submission.py", | |
| stdout=e.stdout, | |
| stderr=e.stderr, | |
| exit_code=e.returncode, | |
| ) | |
| run, profile = run_single_evaluation(["python3", main], **kwargs) | |
| return EvalResult( | |
| start=start, | |
| end=datetime.datetime.now(), | |
| compilation=comp, | |
| run=run, | |
| profile=profile, | |
| ) | |
| finally: | |
| for f in sources.keys(): | |
| if os.path.exists(f): | |
| os.remove(f) | |
| class _EvalRunner(Protocol): | |
| def __call__(self, mode: str, **kwargs) -> EvalResult: ... | |
| def run_evaluation( | |
| call: _EvalRunner, | |
| mode: str, | |
| common_args: dict, | |
| ) -> dict[str, EvalResult]: | |
| """ | |
| Given a "runner" function `call`, interprets the mode | |
| and calls the runner with the right arguments. | |
| Simple modes (test, benchmark, profile) just | |
| invoke the runner once, but private/leaderboard | |
| require multiple runner calls. | |
| """ | |
| results: dict[str, EvalResult] = {} | |
| if mode == "profile": | |
| benchmarks = copy.deepcopy(common_args["benchmarks"]) | |
| for i, benchmark in enumerate(benchmarks.splitlines()): | |
| common_args["benchmarks"] = benchmark | |
| results[f"{mode}.{i}"] = call(mode=mode, **common_args) | |
| elif mode in ["test", "benchmark"]: | |
| results[mode] = call(mode=mode, **common_args) | |
| elif mode in ["private", "leaderboard"]: | |
| # first, run the tests | |
| results["test"] = call(mode="test", **common_args) | |
| if not results["test"].run or not results["test"].run.passed: | |
| return results | |
| # Unnecessary | |
| # results["benchmark"] = call(mode="benchmark", **common_args) | |
| # if not results["benchmark"].run or not results["benchmark"].run.passed: | |
| # return results | |
| # if they pass, run the leaderboard validation | |
| results["leaderboard"] = call(mode="leaderboard", **common_args) | |
| else: | |
| raise AssertionError("Invalid mode") | |
| return results | |
| def build_test_string(tests: list[dict]): | |
| as_str = "" | |
| for test in tests: | |
| kvs = [] | |
| for k, v in test.items(): | |
| kvs.append(f"{k}: {v}") | |
| as_str += "; ".join(kvs) + "\n" | |
| return as_str | |
| def run_config(config: dict): | |
| system = make_system_info() | |
| common_args = { | |
| "system": system, | |
| "tests": build_test_string(config.get("tests", [])), | |
| "benchmarks": build_test_string(config.get("benchmarks", [])), | |
| "seed": config.get("seed", None), | |
| "ranking_by": config.get("ranking_by", "last"), | |
| "ranked_timeout": config.get("ranked_timeout", Timeout.RANKED), | |
| "benchmark_timeout": config.get("benchmark_timeout", Timeout.BENCHMARK), | |
| "test_timeout": config.get("test_timeout", Timeout.TEST), | |
| "multi_gpu": config.get("multi_gpu", False), | |
| } | |
| if config["lang"] == "py": | |
| runner = functools.partial( | |
| run_pytorch_script, sources=config["sources"], main=config["main"] | |
| ) | |
| elif config["lang"] == "cu": | |
| runner = functools.partial( | |
| run_cuda_script, | |
| sources=config["sources"], | |
| headers=config.get("headers", {}), | |
| arch=config.get("arch", None), | |
| defines=config.get("defines", {}), | |
| include_dirs=config.get("include_dirs", []), | |
| libraries=config.get("libraries", []), | |
| flags=CUDA_FLAGS, | |
| ) | |
| else: | |
| raise ValueError(f"Invalid language {config['lang']}") | |
| results = run_evaluation(runner, config["mode"], common_args) | |
| return FullResult(success=True, error="", runs=results, system=system) | |