diff --git a/.github/actions/install-deps/action.yml b/.github/actions/install-deps/action.yml index 6a6241a40..10fafbf52 100644 --- a/.github/actions/install-deps/action.yml +++ b/.github/actions/install-deps/action.yml @@ -24,6 +24,13 @@ runs: shell: bash run: uv venv + # Installed before the other dependencies, so that it also works on Python versions for which + # no stable torch wheel exists yet. + - name: Install torch nightly + if: ${{ inputs.torch_channel == 'nightly' }} + shell: bash + run: uv pip install --pre torch torchvision -P torch -P torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu + - name: Install dependencies (options=[${{ inputs.options }}], groups=[${{ inputs.groups }}]) shell: bash run: | @@ -50,9 +57,8 @@ runs: echo $cmd eval $cmd - - name: Upgrade to torch nightly + - name: Print torch nightly version if: ${{ inputs.torch_channel == 'nightly' }} shell: bash run: | - uv pip install --pre torch torchvision -P torch -P torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cpu uv run python -c "import torch, torchvision; print('torch:', torch.__version__, '| torchvision:', torchvision.__version__)" diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 88207e136..8f415fd54 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -45,6 +45,11 @@ jobs: extra_groups: 'lower_bounds' # Latest: torch nightly build. - torch_channel: 'nightly' + # Upcoming Python version, not yet supported by stable torch. cvxpy is excluded because one + # of its dependencies (sparsediffpy) has no wheel yet and can't be built without BLAS. + - python-version: '3.15' + torch_channel: 'nightly' + options: 'quadprog_projector fairgrad' steps: - name: Checkout repository diff --git a/.gitignore b/.gitignore index 364fd0259..6458b3764 100644 --- a/.gitignore +++ b/.gitignore @@ -6,8 +6,8 @@ worktrees/ CLAUDE.md .claude/ -# Profiling results -tests/profiling/results/ +# Profiler results +tests/profiler/results/ # Trajectories results tests/trajectories/results/ diff --git a/tests/paths.py b/tests/paths.py index 5300c93c8..cdae39740 100644 --- a/tests/paths.py +++ b/tests/paths.py @@ -3,5 +3,5 @@ TORCHJD_DIR = Path(__file__).parent.parent TESTS_DIR = Path(__file__).parent PLOTS_RESULTS_DIR = TESTS_DIR / "plots" / "results" -PROFILING_RESULTS_DIR = TESTS_DIR / "profiling" / "results" +PROFILER_RESULTS_DIR = TESTS_DIR / "profiler" / "results" TRAJECTORIES_RESULTS_DIR = TESTS_DIR / "trajectories" / "results" diff --git a/tests/profiling/__init__.py b/tests/profiler/__init__.py similarity index 100% rename from tests/profiling/__init__.py rename to tests/profiler/__init__.py diff --git a/tests/profiling/compare_traces.py b/tests/profiler/compare_traces.py similarity index 100% rename from tests/profiling/compare_traces.py rename to tests/profiler/compare_traces.py diff --git a/tests/profiling/plot_memory_timeline.py b/tests/profiler/plot_memory_timeline.py similarity index 96% rename from tests/profiling/plot_memory_timeline.py rename to tests/profiler/plot_memory_timeline.py index 398dbc1ed..8eb161ab8 100644 --- a/tests/profiling/plot_memory_timeline.py +++ b/tests/profiler/plot_memory_timeline.py @@ -13,7 +13,7 @@ matplotlib.use("Agg") # Use non-GUI backend to avoid tkinter dependency import matplotlib.pyplot as plt import numpy as np -from paths import PROFILING_RESULTS_DIR +from paths import PROFILER_RESULTS_DIR @dataclass @@ -63,7 +63,7 @@ def plot_memory_timelines(experiment: str, folders: list[str]) -> None: cpu_timelines = [] cuda_timelines = [] for folder in folders: - path = PROFILING_RESULTS_DIR / folder / f"{experiment}.json" + path = PROFILER_RESULTS_DIR / folder / f"{experiment}.json" cpu_timeline, cuda_timeline = extract_memory_timelines(path) cpu_timelines.append(cpu_timeline) cuda_timelines.append(cuda_timeline) @@ -104,7 +104,7 @@ def plot_memory_timelines(experiment: str, folders: list[str]) -> None: fig.tight_layout() - output_dir = Path(PROFILING_RESULTS_DIR / "memory_timelines") + output_dir = Path(PROFILER_RESULTS_DIR / "memory_timelines") output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"{experiment}.png" print(f"\nSaving plot to: {output_path}") diff --git a/tests/profiling/run_profiler.py b/tests/profiler/run_profiler.py similarity index 97% rename from tests/profiling/run_profiler.py rename to tests/profiler/run_profiler.py index 9ee042b34..b2b133248 100644 --- a/tests/profiling/run_profiler.py +++ b/tests/profiler/run_profiler.py @@ -23,7 +23,7 @@ ) from utils.tensors import make_inputs_and_targets -from tests.paths import PROFILING_RESULTS_DIR +from tests.paths import PROFILER_RESULTS_DIR from torchjd.aggregation import UPGrad, UPGradWeighting from torchjd.autogram import Engine @@ -98,7 +98,7 @@ def _save_and_print_trace( batch_size: int, ) -> None: filename = f"{factory}-bs{batch_size}-{DEVICE.type}.json" - output_dir = PROFILING_RESULTS_DIR / method_name + output_dir = PROFILER_RESULTS_DIR / method_name output_dir.mkdir(parents=True, exist_ok=True) trace_path = output_dir / filename diff --git a/tests/profiling/speed_grad_vs_jac_vs_gram.py b/tests/profiler/speed_grad_vs_jac_vs_gram.py similarity index 100% rename from tests/profiling/speed_grad_vs_jac_vs_gram.py rename to tests/profiler/speed_grad_vs_jac_vs_gram.py