mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 15:37:44 +00:00
TPTBusiness
f3a2e2b4f1
- Add Bandit security scanner to requirements and pre-commit hooks - Fix CWE-22 path traversal in tarfile/zipfile extraction (3 files) * Add _safe_extract() validation in submit.py, env.py, kaggle_crawler.py * Prevents malicious archives from writing outside target directory - Fix MD5 hashlib calls with usedforsecurity=False flag (2 files) * submission_format_test.txt files for checksum validation - Configure .bandit.yml for automated security scanning * Skip known false positives: B602 (subprocess), B701 (Jinja2) - Add security runbook documentation in docs/security/ - Add pre-commit hook scripts for automated Bandit scanning All 106 backtesting and security tests pass. Security issues resolved: B201, B202, B324 (9 total fixes)
433 lines
17 KiB
Python
433 lines
17 KiB
Python
# %%
|
|
import bisect
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tarfile
|
|
import time
|
|
import zipfile
|
|
from itertools import chain
|
|
from pathlib import Path
|
|
|
|
import nbformat
|
|
from rich import print
|
|
from selenium import webdriver
|
|
from selenium.webdriver.chrome.service import Service
|
|
from selenium.webdriver.common.by import By
|
|
from webdriver_manager.chrome import ChromeDriverManager
|
|
|
|
from rdagent.core.conf import ExtendedBaseSettings
|
|
from rdagent.core.exception import KaggleError
|
|
from rdagent.core.utils import cache_with_pickle
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.oai.llm_utils import APIBackend
|
|
from rdagent.scenarios.data_science.debug.data import create_debug_data
|
|
from rdagent.utils.agent.tpl import T
|
|
from rdagent.utils.env import MLEBDockerEnv
|
|
|
|
# %%
|
|
options = webdriver.ChromeOptions()
|
|
options.add_argument("--no-sandbox")
|
|
options.add_argument("--disable-dev-shm-usage")
|
|
options.add_argument("--headless")
|
|
|
|
|
|
def crawl_descriptions(
|
|
competition: str, local_data_path: str, wait: float = 3.0, force: bool = False
|
|
) -> dict[str, str] | str:
|
|
if (fp := Path(f"{local_data_path}/{competition}/description.md")).exists() and not force:
|
|
logger.info(f"Found {competition}/description.md, loading from it.")
|
|
return fp.read_text()
|
|
|
|
if (fp := Path(f"{local_data_path}/{competition}.json")).exists() and not force:
|
|
logger.info(f"Found {competition}.json, loading from local file.")
|
|
with fp.open("r") as f:
|
|
return json.load(f)
|
|
|
|
# Use webdriver-manager to automatically download and manage ChromeDriver version
|
|
driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)
|
|
overview_url = f"https://www.kaggle.com/competitions/{competition}/overview"
|
|
driver.get(overview_url)
|
|
time.sleep(wait)
|
|
site_body = driver.find_element(By.ID, "site-content")
|
|
descriptions = {}
|
|
|
|
# Get the subtitles
|
|
elements = site_body.find_elements(By.CSS_SELECTOR, f"a[href^='/competitions/{competition}/overview/']")
|
|
subtitles = []
|
|
for e in elements:
|
|
inner_text = ""
|
|
for child in e.find_elements(By.XPATH, ".//*"):
|
|
inner_text += child.get_attribute("innerHTML").strip()
|
|
subtitles.append(inner_text)
|
|
|
|
def kaggle_description_css_selectors() -> tuple[str, str]:
|
|
# Get the class name of the main contents
|
|
ab_elm = site_body.find_element(By.ID, "abstract")
|
|
others_elm = ab_elm.find_element(By.XPATH, "../*[2]")
|
|
first_elm = others_elm.find_element(By.XPATH, "./*[1]")
|
|
first_content_elm = first_elm.find_element(By.XPATH, "./*[1]/*[2]")
|
|
selector_elm = first_content_elm.find_element(By.XPATH, "./*[1]/*[1]")
|
|
main_class = selector_elm.get_attribute("class").split()[-1]
|
|
|
|
# Get the class name of the citation
|
|
citation_elm = site_body.find_element(By.ID, "citation")
|
|
citation_content_elm = citation_elm.find_element(By.XPATH, "./*[1]/*[2]/*[1]/*[1]")
|
|
citation_class = citation_content_elm.get_attribute("class").split()[-1]
|
|
|
|
return main_class, citation_class
|
|
|
|
main_class, citation_class = kaggle_description_css_selectors()
|
|
|
|
# Get main contents
|
|
contents = []
|
|
elements = site_body.find_elements(By.CSS_SELECTOR, f".{main_class}")
|
|
for e in elements:
|
|
content = e.get_attribute("innerHTML")
|
|
contents.append(content)
|
|
|
|
assert len(subtitles) == len(contents) + 1 and subtitles[-1] == "Citation"
|
|
for i in range(len(subtitles) - 1):
|
|
descriptions[subtitles[i]] = contents[i]
|
|
|
|
# Get the citation
|
|
element = site_body.find_element(By.CSS_SELECTOR, f".{citation_class}")
|
|
citation = element.get_attribute("innerHTML")
|
|
descriptions[subtitles[-1]] = citation
|
|
|
|
data_url = f"https://www.kaggle.com/competitions/{competition}/data"
|
|
driver.get(data_url)
|
|
time.sleep(wait)
|
|
data_element = driver.find_element(By.CSS_SELECTOR, f".{main_class}")
|
|
descriptions["Data Description"] = data_element.get_attribute("innerHTML")
|
|
|
|
driver.quit()
|
|
with open(f"{local_data_path}/{competition}.json", "w") as f:
|
|
json.dump(descriptions, f)
|
|
return descriptions
|
|
|
|
|
|
def download_data(competition: str, settings: ExtendedBaseSettings, enable_create_debug_data: bool = True) -> None:
|
|
local_path = settings.local_data_path
|
|
if settings.if_using_mle_data:
|
|
zipfile_path = f"{local_path}/zip_files"
|
|
zip_competition_path = Path(zipfile_path) / competition
|
|
competition_local_path = Path(local_path) / competition
|
|
|
|
if not zip_competition_path.exists():
|
|
mleb_env = MLEBDockerEnv()
|
|
mleb_env.prepare()
|
|
(Path(zipfile_path)).mkdir(parents=True, exist_ok=True)
|
|
mleb_env.check_output(
|
|
f"mlebench prepare -c {competition} --data-dir ./zip_files",
|
|
local_path=local_path,
|
|
running_extra_volume={str(Path("~/.kaggle").expanduser().absolute()): "/root/.kaggle"},
|
|
)
|
|
|
|
if not competition_local_path.exists() or list(competition_local_path.iterdir()) == []:
|
|
competition_local_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
mleb_env = MLEBDockerEnv()
|
|
mleb_env.prepare()
|
|
mleb_env.check_output(
|
|
f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path
|
|
)
|
|
|
|
for zip_path in competition_local_path.rglob("*.zip"):
|
|
with zipfile.ZipFile(zip_path, "r") as zip_ref:
|
|
if len(zip_ref.namelist()) == 1:
|
|
mleb_env.check_output(
|
|
f"unzip -o ./{zip_path.relative_to(competition_local_path)} -d {zip_path.parent.relative_to(competition_local_path)}",
|
|
local_path=competition_local_path,
|
|
)
|
|
else:
|
|
mleb_env.check_output(
|
|
f"mkdir -p ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}; unzip -o ./{zip_path.relative_to(competition_local_path)} -d ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}",
|
|
local_path=competition_local_path,
|
|
)
|
|
for tar_path in competition_local_path.rglob("*.tar*"):
|
|
if not tarfile.is_tarfile(tar_path):
|
|
logger.error(f"{tar_path} is not a valid tar file.")
|
|
continue
|
|
is_gzip_file = open(tar_path, "rb").read(2) == b"\x1f\x8b"
|
|
with tarfile.open(tar_path, "r:gz") if is_gzip_file else tarfile.open(tar_path, "r") as tar_ref:
|
|
if len(tar_ref.getmembers()) == 1:
|
|
mleb_env.check_output(
|
|
f"tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C {tar_path.parent.relative_to(competition_local_path)}",
|
|
local_path=competition_local_path,
|
|
)
|
|
else:
|
|
folder_name = tar_path.name.replace(".tar", "").replace(".gz", "")
|
|
mleb_env.check_output(
|
|
f"mkdir -p ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}; tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}",
|
|
local_path=competition_local_path,
|
|
)
|
|
# NOTE:
|
|
# Patching: due to mle has special renaming mechanism for different competition;
|
|
# We have to switch the schema back to a uniform one;
|
|
if competition in {"new-york-city-taxi-fare-prediction"}:
|
|
cpath = Path(local_path) / f"{competition}"
|
|
labels_path = cpath / "labels.csv"
|
|
train_path = cpath / "train.csv"
|
|
if labels_path.exists():
|
|
shutil.copy(labels_path, train_path)
|
|
else:
|
|
logger.error(f"labels.csv not found in {cpath}")
|
|
raise FileNotFoundError(f"{labels_path} does not exist")
|
|
else:
|
|
zipfile_path = f"{local_path}/zip_files"
|
|
if not Path(f"{zipfile_path}/{competition}.zip").exists():
|
|
try:
|
|
subprocess.run(
|
|
["kaggle", "competitions", "download", "-c", competition, "-p", zipfile_path],
|
|
check=True,
|
|
stderr=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
)
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Download failed: {e}, stderr: {e.stderr}, stdout: {e.stdout}")
|
|
raise KaggleError(f"Download failed: {e}, stderr: {e.stderr}, stdout: {e.stdout}")
|
|
|
|
# unzip data
|
|
unzip_path = f"{local_path}/{competition}"
|
|
if not Path(unzip_path).exists():
|
|
unzip_data(unzip_file_path=f"{zipfile_path}/{competition}.zip", unzip_target_path=unzip_path)
|
|
for sub_zip_file in Path(unzip_path).rglob("*.zip"):
|
|
unzip_data(sub_zip_file, unzip_target_path=unzip_path)
|
|
|
|
# sample data
|
|
if enable_create_debug_data and not Path(f"{local_path}/sample/{competition}").exists():
|
|
create_debug_data(competition, dataset_path=local_path)
|
|
|
|
|
|
def _safe_extract_zip(zip_ref: zipfile.ZipFile, path: str) -> None:
|
|
"""Extract zipfile safely, preventing path traversal attacks (CWE-22)."""
|
|
abs_path = os.path.realpath(path)
|
|
for member in zip_ref.infolist():
|
|
member_path = os.path.realpath(os.path.join(abs_path, member.filename))
|
|
if not member_path.startswith(abs_path):
|
|
raise ValueError(f"Attempted path traversal in zip file: {member.filename}")
|
|
zip_ref.extractall(path=path) # nosec B202:tarfile_unsafe_members - validated above
|
|
|
|
|
|
def unzip_data(unzip_file_path: str, unzip_target_path: str) -> None:
|
|
with zipfile.ZipFile(unzip_file_path, "r") as zip_ref:
|
|
_safe_extract_zip(zip_ref, unzip_target_path)
|
|
|
|
|
|
@cache_with_pickle(hash_func=lambda x: x, force=True)
|
|
def leaderboard_scores(competition: str) -> list[float]:
|
|
from kaggle.api.kaggle_api_extended import KaggleApi
|
|
|
|
api = KaggleApi()
|
|
api.authenticate()
|
|
|
|
return [i.score for i in api.competition_leaderboard_view(competition)]
|
|
|
|
|
|
def get_metric_direction(competition: str) -> bool:
|
|
"""
|
|
Return **True** if the metric is *bigger is better*, **False** if *smaller is better*.
|
|
"""
|
|
if competition == "aerial-cactus-identification":
|
|
return True
|
|
if competition == "leaf-classification":
|
|
return False
|
|
leaderboard = leaderboard_scores(competition)
|
|
|
|
return float(leaderboard[0]) > float(leaderboard[-1])
|
|
|
|
|
|
# FIXME: current score_rank is incorrect because kaggle api returns only the first page leaderboard
|
|
def score_rank(competition: str, score: float) -> tuple[int, float]:
|
|
"""
|
|
Return
|
|
------
|
|
rank: int
|
|
rank_percent: float
|
|
"""
|
|
scores = leaderboard_scores(competition)
|
|
if scores[0] < scores[-1]: # Ascending order
|
|
rank = bisect.bisect_right(scores, score)
|
|
else: # Descending order
|
|
scores = scores[::-1] # Reverse the list to use bisect
|
|
rank = len(scores) - bisect.bisect_right(scores, score)
|
|
|
|
rank = rank + 1
|
|
rank_percent = rank / len(scores) * 100
|
|
|
|
return rank, rank_percent
|
|
|
|
|
|
def download_notebooks(competition: str, local_path: str, num: int = 15) -> None:
|
|
data_path = Path(f"{local_path}/{competition}")
|
|
from kaggle.api.kaggle_api_extended import KaggleApi
|
|
|
|
api = KaggleApi()
|
|
api.authenticate()
|
|
|
|
# judge the sort_by
|
|
ll = api.competition_leaderboard_view(competition)
|
|
score_diff = float(ll[0].score) - float(ll[-1].score)
|
|
if score_diff > 0:
|
|
sort_by = "scoreDescending"
|
|
else:
|
|
sort_by = "scoreAscending"
|
|
|
|
# download notebooks
|
|
nl = api.kernels_list(competition=competition, sort_by=sort_by, page=1, page_size=num)
|
|
for nb in nl:
|
|
author = nb.ref.split("/")[0]
|
|
api.kernels_pull(nb.ref, path=data_path / author)
|
|
print(f"Downloaded {len(nl)} notebooks for {competition}. ([red]{sort_by}[/red])")
|
|
|
|
|
|
def notebook_to_knowledge(notebook_text: str) -> str:
|
|
sys_prompt = T(".prompts:gen_knowledge_from_code_mini_case.system").r()
|
|
user_prompt = T(".prompts:gen_knowledge_from_code_mini_case.user").r(notebook=notebook_text)
|
|
|
|
response = APIBackend().build_messages_and_create_chat_completion(
|
|
user_prompt=user_prompt,
|
|
system_prompt=sys_prompt,
|
|
json_mode=False,
|
|
)
|
|
return response
|
|
|
|
|
|
def convert_notebooks_to_text(competition: str, local_path: str) -> None:
|
|
data_path = Path(f"{local_path}/{competition}")
|
|
converted_num = 0
|
|
|
|
# convert ipynb and irnb files
|
|
for nb_path in chain(data_path.glob("**/*.ipynb"), data_path.glob("**/*.irnb")):
|
|
with nb_path.open("r", encoding="utf-8") as f:
|
|
nb = nbformat.read(f, as_version=4)
|
|
text = []
|
|
for cell in nb.cells:
|
|
if cell.cell_type == "markdown":
|
|
text.append(f"```markdown\n{cell.source}```")
|
|
elif cell.cell_type == "code":
|
|
text.append(f"```code\n{cell.source}```")
|
|
text = "\n\n".join(text)
|
|
|
|
text = notebook_to_knowledge(text)
|
|
|
|
text_path = nb_path.with_suffix(".txt")
|
|
text_path.write_text(text, encoding="utf-8")
|
|
converted_num += 1
|
|
|
|
# convert py files
|
|
for py_path in data_path.glob("**/*.py"):
|
|
with py_path.open("r", encoding="utf-8") as f:
|
|
text = f"```code\n{f.read()}```"
|
|
|
|
text = notebook_to_knowledge(text)
|
|
|
|
text_path = py_path.with_suffix(".txt")
|
|
text_path.write_text(text, encoding="utf-8")
|
|
converted_num += 1
|
|
|
|
print(f"Converted {converted_num} notebooks to text files.")
|
|
|
|
|
|
def collect_knowledge_texts(notebooks_path: str | Path) -> dict[str, list[str]]:
|
|
"""
|
|
{
|
|
"competition1": [
|
|
"knowledge_text1",
|
|
"knowledge_text2",
|
|
...
|
|
],
|
|
“competition2”: [
|
|
"knowledge_text1",
|
|
"knowledge_text2",
|
|
...
|
|
],
|
|
...
|
|
}
|
|
"""
|
|
notebooks_dir = Path(notebooks_path)
|
|
|
|
competition_knowledge_texts_dict = {}
|
|
for competition_dir in notebooks_dir.iterdir():
|
|
knowledge_texts = []
|
|
for text_path in competition_dir.glob("**/*.txt"):
|
|
text = text_path.read_text(encoding="utf-8")
|
|
knowledge_texts.append(text)
|
|
|
|
competition_knowledge_texts_dict[competition_dir.name] = knowledge_texts
|
|
|
|
return competition_knowledge_texts_dict
|
|
|
|
|
|
# %%
|
|
if __name__ == "__main__":
|
|
mini_case_cs = [
|
|
"feedback-prize-english-language-learning",
|
|
"playground-series-s3e11",
|
|
"playground-series-s3e14",
|
|
"spaceship-titanic",
|
|
"playground-series-s3e18",
|
|
"playground-series-s3e16",
|
|
"playground-series-s3e9",
|
|
"playground-series-s3e25",
|
|
"playground-series-s3e26",
|
|
"playground-series-s3e24",
|
|
"playground-series-s3e23",
|
|
]
|
|
|
|
other_cs = [
|
|
"amp-parkinsons-disease-progression-prediction",
|
|
"arc-prize-2024",
|
|
"ariel-data-challenge-2024",
|
|
"child-mind-institute-detect-sleep-states",
|
|
"connectx",
|
|
"contradictory-my-dear-watson",
|
|
"digit-recognizer",
|
|
"fathomnet-out-of-sample-detection",
|
|
"forest-cover-type-prediction",
|
|
"gan-getting-started",
|
|
"google-research-identify-contrails-reduce-global-warming",
|
|
"house-prices-advanced-regression-techniques",
|
|
"isic-2024-challenge",
|
|
"leash-BELKA",
|
|
"llm-20-questions",
|
|
"nlp-getting-started",
|
|
"playground-series-s4e1",
|
|
"playground-series-s4e2",
|
|
"playground-series-s4e3",
|
|
"playground-series-s4e4",
|
|
"playground-series-s4e5",
|
|
"playground-series-s4e6",
|
|
"playground-series-s4e7",
|
|
"playground-series-s4e8",
|
|
"rsna-2024-lumbar-spine-degenerative-classification",
|
|
"sf-crime",
|
|
"store-sales-time-series-forecasting",
|
|
"titanic",
|
|
"tpu-getting-started",
|
|
# scenario competition
|
|
"covid19-global-forecasting-week-1",
|
|
"statoil-iceberg-classifier-challenge",
|
|
"optiver-realized-volatility-prediction",
|
|
"facebook-v-predicting-check-ins",
|
|
]
|
|
|
|
# all_cs = mini_case_cs + other_cs
|
|
# for c in all_cs:
|
|
# convert_notebooks_to_text(c)
|
|
# exit()
|
|
# from kaggle.api.kaggle_api_extended import KaggleApi
|
|
|
|
# api = KaggleApi()
|
|
# api.authenticate()
|
|
# cs = api.competitions_list()
|
|
# for c in cs:
|
|
# name = c.ref.split("/")[-1]
|
|
# crawl_descriptions(name)
|
|
res = leaderboard_scores(competition="playground-series-s4e8")
|
|
rank, rank_percent = score_rank(competition="playground-series-s4e8", score=0.9832)
|
|
print(rank, rank_percent)
|
|
# %%
|