diff --git a/src/geant4_python_application/files/2datasets.py b/src/geant4_python_application/files/2datasets.py index 7cf8c9b..a282858 100644 --- a/src/geant4_python_application/files/2datasets.py +++ b/src/geant4_python_application/files/2datasets.py @@ -7,7 +7,6 @@ import tarfile import tempfile from collections import namedtuple -from pathlib import Path import requests from tqdm import tqdm @@ -200,4 +199,3 @@ def reinstall_datasets(): # Uninstall and reinstall datasets uninstall_datasets() install_datasets(force=True) - diff --git a/src/geant4_python_application/files/alt_datasets.py b/src/geant4_python_application/files/alt_datasets.py index c5f9ce6..1c59d5d 100644 --- a/src/geant4_python_application/files/alt_datasets.py +++ b/src/geant4_python_application/files/alt_datasets.py @@ -18,20 +18,81 @@ # List of all Geant4 datasets and their download links datasets = ( - Dataset(name="G4NDL", version="4.7.1", url="https://cern.ch/geant4-data/datasets/G4NDL.4.7.1.tar.gz", env="G4NEUTRONHPDATA"), - Dataset(name="G4EMLOW", version="8.5", url="https://cern.ch/geant4-data/datasets/G4EMLOW.8.5.tar.gz", env="G4LEDATA"), - Dataset(name="PhotonEvaporation", version="5.7", url="https://cern.ch/geant4-data/datasets/G4PhotonEvaporation.5.7.tar.gz", env="G4LEVELGAMMADATA"), - Dataset(name="RadioactiveDecay", version="5.6", url="https://cern.ch/geant4-data/datasets/G4RadioactiveDecay.5.6.tar.gz", env="G4RADIOACTIVEDATA"), - Dataset(name="G4PARTICLEXS", version="4.0", url="https://cern.ch/geant4-data/datasets/G4PARTICLEXS.4.0.tar.gz", env="G4PARTICLEXSDATA"), - Dataset(name="G4PII", version="1.3", url="https://cern.ch/geant4-data/datasets/G4PII.1.3.tar.gz", env="G4PIIDATA"), - Dataset(name="RealSurface", version="2.2", url="https://cern.ch/geant4-data/datasets/G4RealSurface.2.2.tar.gz", env="G4REALSURFACEDATA"), - Dataset(name="G4SAIDDATA", version="2.0", url="https://cern.ch/geant4-data/datasets/G4SAIDDATA.2.0.tar.gz", env="G4SAIDXSDATA"), - Dataset(name="G4ABLA", version="3.3", url="https://cern.ch/geant4-data/datasets/G4ABLA.3.3.tar.gz", env="G4ABLADATA"), - Dataset(name="G4INCL", version="1.2", url="https://cern.ch/geant4-data/datasets/G4INCL.1.2.tar.gz", env="G4INCLDATA"), - Dataset(name="G4ENSDFSTATE", version="2.3", url="https://cern.ch/geant4-data/datasets/G4ENSDFSTATE.2.3.tar.gz", env="G4ENSDFSTATEDATA"), - Dataset(name="G4TENDL", version="1.4", url="https://cern.ch/geant4-data/datasets/G4TENDL.1.4.tar.gz", env="G4TENDLDATA"), + Dataset( + name="G4NDL", + version="4.7.1", + url="https://cern.ch/geant4-data/datasets/G4NDL.4.7.1.tar.gz", + env="G4NEUTRONHPDATA", + ), + Dataset( + name="G4EMLOW", + version="8.5", + url="https://cern.ch/geant4-data/datasets/G4EMLOW.8.5.tar.gz", + env="G4LEDATA", + ), + Dataset( + name="PhotonEvaporation", + version="5.7", + url="https://cern.ch/geant4-data/datasets/G4PhotonEvaporation.5.7.tar.gz", + env="G4LEVELGAMMADATA", + ), + Dataset( + name="RadioactiveDecay", + version="5.6", + url="https://cern.ch/geant4-data/datasets/G4RadioactiveDecay.5.6.tar.gz", + env="G4RADIOACTIVEDATA", + ), + Dataset( + name="G4PARTICLEXS", + version="4.0", + url="https://cern.ch/geant4-data/datasets/G4PARTICLEXS.4.0.tar.gz", + env="G4PARTICLEXSDATA", + ), + Dataset( + name="G4PII", + version="1.3", + url="https://cern.ch/geant4-data/datasets/G4PII.1.3.tar.gz", + env="G4PIIDATA", + ), + Dataset( + name="RealSurface", + version="2.2", + url="https://cern.ch/geant4-data/datasets/G4RealSurface.2.2.tar.gz", + env="G4REALSURFACEDATA", + ), + Dataset( + name="G4SAIDDATA", + version="2.0", + url="https://cern.ch/geant4-data/datasets/G4SAIDDATA.2.0.tar.gz", + env="G4SAIDXSDATA", + ), + Dataset( + name="G4ABLA", + version="3.3", + url="https://cern.ch/geant4-data/datasets/G4ABLA.3.3.tar.gz", + env="G4ABLADATA", + ), + Dataset( + name="G4INCL", + version="1.2", + url="https://cern.ch/geant4-data/datasets/G4INCL.1.2.tar.gz", + env="G4INCLDATA", + ), + Dataset( + name="G4ENSDFSTATE", + version="2.3", + url="https://cern.ch/geant4-data/datasets/G4ENSDFSTATE.2.3.tar.gz", + env="G4ENSDFSTATEDATA", + ), + Dataset( + name="G4TENDL", + version="1.4", + url="https://cern.ch/geant4-data/datasets/G4TENDL.1.4.tar.gz", + env="G4TENDLDATA", + ), ) + # Function to determine the data directory dynamically def data_directory() -> str: return os.path.join( @@ -42,23 +103,32 @@ def data_directory() -> str: "data", ) + # Check if a dataset is already installed def is_dataset_installed(dataset: Dataset) -> bool: dataset_path = os.path.join(data_directory(), f"{dataset.name}.{dataset.version}") return os.path.exists(dataset_path) and os.path.isdir(dataset_path) + # Get the size of the dataset for progress tracking def _get_dataset_download_size(dataset: Dataset) -> int: r = requests.head(dataset.url) r.raise_for_status() return int(r.headers.get("content-length", 0)) + # Get the total download size for all selected datasets def _get_total_download_size(datasets_to_download: list[Dataset] = datasets) -> int: - with concurrent.futures.ThreadPoolExecutor(max_workers=len(datasets_to_download)) as executor: - futures = [executor.submit(_get_dataset_download_size, dataset) for dataset in datasets_to_download] + with concurrent.futures.ThreadPoolExecutor( + max_workers=len(datasets_to_download) + ) as executor: + futures = [ + executor.submit(_get_dataset_download_size, dataset) + for dataset in datasets_to_download + ] return sum(f.result() for f in concurrent.futures.as_completed(futures)) + # Download and extract the dataset with progress tracking def _download_extract_dataset(dataset: Dataset, pbar: tqdm): r = requests.get(dataset.url, stream=True) @@ -74,14 +144,17 @@ def _download_extract_dataset(dataset: Dataset, pbar: tqdm): with tarfile.open(fileobj=f, mode="r:gz") as tar: tar.extractall(data_directory()) + # Main function to handle the installation of datasets def install_datasets(force: bool = False, show_progress: bool = True): os.environ["GEANT4_DATA_DIR"] = data_directory() datasets_to_download = [] - + # Check if datasets are already downloaded, unless force is True for dataset in datasets: - dataset_path = os.path.join(data_directory(), f"{dataset.name}.{dataset.version}") + dataset_path = os.path.join( + data_directory(), f"{dataset.name}.{dataset.version}" + ) os.environ[dataset.env] = dataset_path if not is_dataset_installed(dataset) or force: datasets_to_download.append(dataset) @@ -91,7 +164,7 @@ def install_datasets(force: bool = False, show_progress: bool = True): return # All datasets are already installed os.makedirs(data_directory(), exist_ok=True) - + # Show progress for the dataset download and extraction if show_progress: print( @@ -124,6 +197,7 @@ def install_datasets(force: bool = False, show_progress: bool = True): ) / (1024**3) print(f"Geant4 datasets size on disk after extraction: {total_size_gb:.2f}GB") + # Uninstall datasets (remove the dataset directory) def uninstall_datasets(): dir_to_remove = os.path.dirname(data_directory()) @@ -135,6 +209,7 @@ def uninstall_datasets(): ) shutil.rmtree(dir_to_remove, ignore_errors=True) + # Reinstall datasets (uninstall first, then reinstall) def reinstall_datasets(): uninstall_datasets() diff --git a/src/geant4_python_application/files/dataset1.py b/src/geant4_python_application/files/dataset1.py index 9062e2f..fef25b6 100644 --- a/src/geant4_python_application/files/dataset1.py +++ b/src/geant4_python_application/files/dataset1.py @@ -15,6 +15,7 @@ url = "https://cern.ch/geant4-data/datasets" + def data_directory() -> str: return os.path.join( geant4_python_application.application_directory(), @@ -24,35 +25,72 @@ def data_directory() -> str: "data", ) + Dataset = namedtuple("Dataset", ["name", "version", "filename", "env"]) datasets = ( Dataset(name="G4NDL", version="4.7.1", filename="G4NDL", env="G4NEUTRONHPDATA"), Dataset(name="G4EMLOW", version="8.5", filename="G4EMLOW", env="G4LEDATA"), - Dataset(name="PhotonEvaporation", version="5.7", filename="G4PhotonEvaporation", env="G4LEVELGAMMADATA"), - Dataset(name="RadioactiveDecay", version="5.6", filename="G4RadioactiveDecay", env="G4RADIOACTIVEDATA"), - Dataset(name="G4PARTICLEXS", version="4.0", filename="G4PARTICLEXS", env="G4PARTICLEXSDATA"), + Dataset( + name="PhotonEvaporation", + version="5.7", + filename="G4PhotonEvaporation", + env="G4LEVELGAMMADATA", + ), + Dataset( + name="RadioactiveDecay", + version="5.6", + filename="G4RadioactiveDecay", + env="G4RADIOACTIVEDATA", + ), + Dataset( + name="G4PARTICLEXS", + version="4.0", + filename="G4PARTICLEXS", + env="G4PARTICLEXSDATA", + ), Dataset(name="G4PII", version="1.3", filename="G4PII", env="G4PIIDATA"), - Dataset(name="RealSurface", version="2.2", filename="G4RealSurface", env="G4REALSURFACEDATA"), - Dataset(name="G4SAIDDATA", version="2.0", filename="G4SAIDDATA", env="G4SAIDXSDATA"), + Dataset( + name="RealSurface", + version="2.2", + filename="G4RealSurface", + env="G4REALSURFACEDATA", + ), + Dataset( + name="G4SAIDDATA", version="2.0", filename="G4SAIDDATA", env="G4SAIDXSDATA" + ), Dataset(name="G4ABLA", version="3.3", filename="G4ABLA", env="G4ABLADATA"), Dataset(name="G4INCL", version="1.2", filename="G4INCL", env="G4INCLDATA"), - Dataset(name="G4ENSDFSTATE", version="2.3", filename="G4ENSDFSTATE", env="G4ENSDFSTATEDATA"), + Dataset( + name="G4ENSDFSTATE", + version="2.3", + filename="G4ENSDFSTATE", + env="G4ENSDFSTATEDATA", + ), ) + def _dataset_url(dataset: Dataset) -> str: return f"{url}/{dataset.filename}.{dataset.version}.tar.gz" + def _get_dataset_download_size(dataset: Dataset) -> int: r = requests.head(_dataset_url(dataset)) r.raise_for_status() return int(r.headers.get("content-length", 0)) + def _get_total_download_size(datasets_to_download: list[Dataset] = datasets) -> int: - with concurrent.futures.ThreadPoolExecutor(max_workers=len(datasets_to_download)) as executor: - futures = [executor.submit(_get_dataset_download_size, dataset) for dataset in datasets_to_download] + with concurrent.futures.ThreadPoolExecutor( + max_workers=len(datasets_to_download) + ) as executor: + futures = [ + executor.submit(_get_dataset_download_size, dataset) + for dataset in datasets_to_download + ] return sum(f.result() for f in concurrent.futures.as_completed(futures)) + def _download_extract_dataset(dataset: Dataset, pbar: tqdm): filename = dataset.filename urlpath = f"{url}/{filename}.{dataset.version}.tar.gz" @@ -69,10 +107,11 @@ def _download_extract_dataset(dataset: Dataset, pbar: tqdm): with tarfile.open(fileobj=f, mode="r:gz") as tar: tar.extractall(data_directory()) + def install_datasets(force: bool = False, show_progress: bool = True): os.environ["GEANT4_DATA_DIR"] = data_directory() datasets_to_download = [] - + # Check if datasets are already downloaded for dataset in datasets: path = os.path.join(data_directory(), dataset.name + dataset.version) @@ -115,6 +154,7 @@ def install_datasets(force: bool = False, show_progress: bool = True): ) / (1024**3) print(f"Geant4 datasets size on disk after extraction: {total_size_gb:.2f}GB") + def uninstall_datasets(): dir_to_remove = os.path.dirname(data_directory()) package_dir = os.path.dirname(__file__) @@ -125,6 +165,7 @@ def uninstall_datasets(): ) shutil.rmtree(dir_to_remove, ignore_errors=True) + def reinstall_datasets(): uninstall_datasets() install_datasets(force=True) diff --git a/src/geant4_python_application/files/datasets.py b/src/geant4_python_application/files/datasets.py index 9062e2f..0d3293e 100644 --- a/src/geant4_python_application/files/datasets.py +++ b/src/geant4_python_application/files/datasets.py @@ -5,6 +5,7 @@ import shutil import tarfile import tempfile +import time from collections import namedtuple from pathlib import Path @@ -15,6 +16,7 @@ url = "https://cern.ch/geant4-data/datasets" + def data_directory() -> str: return os.path.join( geant4_python_application.application_directory(), @@ -24,55 +26,112 @@ def data_directory() -> str: "data", ) + Dataset = namedtuple("Dataset", ["name", "version", "filename", "env"]) datasets = ( Dataset(name="G4NDL", version="4.7.1", filename="G4NDL", env="G4NEUTRONHPDATA"), Dataset(name="G4EMLOW", version="8.5", filename="G4EMLOW", env="G4LEDATA"), - Dataset(name="PhotonEvaporation", version="5.7", filename="G4PhotonEvaporation", env="G4LEVELGAMMADATA"), - Dataset(name="RadioactiveDecay", version="5.6", filename="G4RadioactiveDecay", env="G4RADIOACTIVEDATA"), - Dataset(name="G4PARTICLEXS", version="4.0", filename="G4PARTICLEXS", env="G4PARTICLEXSDATA"), + Dataset( + name="PhotonEvaporation", + version="5.7", + filename="G4PhotonEvaporation", + env="G4LEVELGAMMADATA", + ), + Dataset( + name="RadioactiveDecay", + version="5.6", + filename="G4RadioactiveDecay", + env="G4RADIOACTIVEDATA", + ), + Dataset( + name="G4PARTICLEXS", + version="4.0", + filename="G4PARTICLEXS", + env="G4PARTICLEXSDATA", + ), Dataset(name="G4PII", version="1.3", filename="G4PII", env="G4PIIDATA"), - Dataset(name="RealSurface", version="2.2", filename="G4RealSurface", env="G4REALSURFACEDATA"), - Dataset(name="G4SAIDDATA", version="2.0", filename="G4SAIDDATA", env="G4SAIDXSDATA"), + Dataset( + name="RealSurface", + version="2.2", + filename="G4RealSurface", + env="G4REALSURFACEDATA", + ), + Dataset( + name="G4SAIDDATA", version="2.0", filename="G4SAIDDATA", env="G4SAIDXSDATA" + ), Dataset(name="G4ABLA", version="3.3", filename="G4ABLA", env="G4ABLADATA"), Dataset(name="G4INCL", version="1.2", filename="G4INCL", env="G4INCLDATA"), - Dataset(name="G4ENSDFSTATE", version="2.3", filename="G4ENSDFSTATE", env="G4ENSDFSTATEDATA"), + Dataset( + name="G4ENSDFSTATE", + version="2.3", + filename="G4ENSDFSTATE", + env="G4ENSDFSTATEDATA", + ), ) + def _dataset_url(dataset: Dataset) -> str: return f"{url}/{dataset.filename}.{dataset.version}.tar.gz" + def _get_dataset_download_size(dataset: Dataset) -> int: r = requests.head(_dataset_url(dataset)) r.raise_for_status() return int(r.headers.get("content-length", 0)) + def _get_total_download_size(datasets_to_download: list[Dataset] = datasets) -> int: - with concurrent.futures.ThreadPoolExecutor(max_workers=len(datasets_to_download)) as executor: - futures = [executor.submit(_get_dataset_download_size, dataset) for dataset in datasets_to_download] + with concurrent.futures.ThreadPoolExecutor( + max_workers=len(datasets_to_download) + ) as executor: + futures = [ + executor.submit(_get_dataset_download_size, dataset) + for dataset in datasets_to_download + ] return sum(f.result() for f in concurrent.futures.as_completed(futures)) + +chunk_size = 1024 * 1024 +max_attempts = 5 + + def _download_extract_dataset(dataset: Dataset, pbar: tqdm): - filename = dataset.filename - urlpath = f"{url}/{filename}.{dataset.version}.tar.gz" - r = requests.get(urlpath, stream=True) - r.raise_for_status() + urlpath = _dataset_url(dataset) - chunk_size = 1024 with tempfile.TemporaryFile() as f: - for chunk in r.iter_content(chunk_size=chunk_size): - f.write(chunk) - pbar.update(chunk_size) + downloaded = 0 + for attempt in range(max_attempts): + # resume where the previous attempt died instead of restarting a multi-GB download + headers = {"Range": f"bytes={downloaded}-"} if downloaded else {} + try: + r = requests.get(urlpath, stream=True, headers=headers, timeout=60) + r.raise_for_status() + if downloaded and r.status_code != 206: + # server ignored the range request, so start over + f.seek(0) + f.truncate() + pbar.update(-downloaded) + downloaded = 0 + for chunk in r.iter_content(chunk_size=chunk_size): + f.write(chunk) + downloaded += len(chunk) + pbar.update(len(chunk)) + break + except requests.RequestException: + if attempt == max_attempts - 1: + raise + time.sleep(2**attempt) f.seek(0) with tarfile.open(fileobj=f, mode="r:gz") as tar: tar.extractall(data_directory()) + def install_datasets(force: bool = False, show_progress: bool = True): os.environ["GEANT4_DATA_DIR"] = data_directory() datasets_to_download = [] - + # Check if datasets are already downloaded for dataset in datasets: path = os.path.join(data_directory(), dataset.name + dataset.version) @@ -107,7 +166,9 @@ def install_datasets(force: bool = False, show_progress: bool = True): executor.submit(_download_extract_dataset, dataset, pbar) for dataset in datasets_to_download ] - concurrent.futures.wait(futures) + for future in concurrent.futures.as_completed(futures): + # surface download/extraction errors instead of reporting a partial install as success + future.result() if show_progress: total_size_gb = sum( @@ -115,6 +176,7 @@ def install_datasets(force: bool = False, show_progress: bool = True): ) / (1024**3) print(f"Geant4 datasets size on disk after extraction: {total_size_gb:.2f}GB") + def uninstall_datasets(): dir_to_remove = os.path.dirname(data_directory()) package_dir = os.path.dirname(__file__) @@ -125,6 +187,7 @@ def uninstall_datasets(): ) shutil.rmtree(dir_to_remove, ignore_errors=True) + def reinstall_datasets(): uninstall_datasets() install_datasets(force=True)