diff --git a/.github/actions/deploy/action.yaml b/.github/actions/deploy/action.yaml new file mode 100644 index 0000000..06200ae --- /dev/null +++ b/.github/actions/deploy/action.yaml @@ -0,0 +1,64 @@ +name: deploy-suite +description: | + A Github action to build and deploy a suite deployed on an ecflow server. +inputs: + troika_user: + description: User used to submit troika job. + required: true + sbatch_options: + description: List of SBATCH directives + required: false + site: + description: HPC site name. + required: false + default: hpc-batch + github_token: + description: Github token. + required: true + ecflow_host: + description: ecflow server hostname. + required: true + ecflow_port: + description: ecflow server port. + required: true + suite_name: + description: Name of the suite. + required: true + build_options: + description: List of options to pass to the build script. + required: false + default: "limit_workers=1 " +runs: + using: composite + steps: + - name: Build and deploy suite to ecflow server + uses: ecmwf-actions/reusable-workflows/ci-hpc-generic@v2 + with: + template: | + set -eux + echo "Job is running on ${{ runner.os }}" + + # Load modules of interest + module load ecflow + module load wellies + + # Clone the repository + PACKAGE_NAME=anemoi-docs + PACKAGE_BRANCH=${{ github.head_ref || github.ref_name }} + git clone -b $PACKAGE_BRANCH https://${{ inputs.github_token }}@github.com/ecmwf/${PACKAGE_NAME}.git + cd $PACKAGE_NAME + cd tests/system-level + + # Set ecflow server variables + export ECF_HOST=${{ inputs.ecflow_host }} + export ECF_PORT=${{ inputs.ecflow_port }} + + # Build the suite definition file + BUILD_DIR=${SCRATCHDIR}/pyflow/${{ inputs.suite_name }}/build + ./build.sh -s name=${{ inputs.suite_name }} ${{ inputs.build_options }} -y + + # Deploy the suite. If a suite is already running it will fail. User should deal with its zombies + ecflow_client --replace=/anemoi_tests/${{ inputs.suite_name }} $HOME/pyflow/anemoi_tests/${{ inputs.suite_name }}/${{ inputs.suite_name }}.def + sbatch_options: ${{ inputs.sbatch_options }} + troika_user: ${{ inputs.troika_user }} + site: ${{ inputs.site }} diff --git a/.github/workflows/nightly_ci_hpc.yaml b/.github/workflows/nightly_ci_hpc.yaml new file mode 100644 index 0000000..ed67a68 --- /dev/null +++ b/.github/workflows/nightly_ci_hpc.yaml @@ -0,0 +1,51 @@ +name: nightly-system-level-test + +on: + schedule: + - cron: '0 23 * * 0' # Nightly at 23pm on all main branches + # pull_request: + # types: [opened, synchronize, reopened] + # branches: + # - main + # - develop + + +jobs: + check-commits: + runs-on: ubuntu-latest + name: Check latest commits + outputs: + should_run: ${{ steps.check-commits.outputs.should_run }} + steps: + - name: Check for recent commits on main in anemoi repos + id: check-commits + run: | + REPOS=("anemoi-core" "anemoi-datasets" "anemoi-docs") + SHOULD_RUN=false + + for REPO in "${REPOS[@]}"; do + echo "Checking repository: $REPO" + + # Fetch the latest commit date on the main branch from GitHub API + LATEST_COMMIT_DATE=$(curl -s "https://api.github.com/repos/ecmwf/$REPO/commits/main" | jq -r '.commit.committer.date') + + if [ "$(date -d "$LATEST_COMMIT_DATE" +%s)" -gt "$(date -d "24 hours ago" +%s)" ]; then + echo "Recent commit found in $REPO: $LATEST_COMMIT_DATE" + SHOULD_RUN=true + fi + done + + echo "should_run=$SHOULD_RUN" >> "$GITHUB_OUTPUT" + + system_level_test: + needs: check-commits + if: needs.check-commits.outputs.should_run == 'true' + uses: ./.github/workflows/system_level_test.yaml + with: + suite_name: "nightly" + build_options: "" + secrets: + gh_token: ${{ secrets.GH_REPO_READ_TOKEN }} + troika_user: ${{ secrets.HPC_SYSTEM_TEST_USER }} + ecflow_host: ${{ secrets.HPC_SYSTEM_TEST_HOST }} + ecflow_port: ${{ secrets.HPC_SYSTEM_TEST_PORT }} diff --git a/.github/workflows/on_demand_ci_hpc.yaml b/.github/workflows/on_demand_ci_hpc.yaml new file mode 100644 index 0000000..7ce3d9f --- /dev/null +++ b/.github/workflows/on_demand_ci_hpc.yaml @@ -0,0 +1,61 @@ +name: on-demand-system-level-test + +on: + workflow_dispatch: + inputs: + anemoi_docs_branch: + description: Branch for anemoi-docs to be used to deploy the anemoi_tests suite. + required: false + anemoi_datasets_branch: + description: Branch for anemoi-datasets to be used to build the datasets environment for testing. + required: false + default: "main" + anemoi_training_branch: + description: Branch for anemoi-training to be used to build the training environment for testing. + required: false + default: "main" + + +jobs: + prepare: + runs-on: ubuntu-latest + outputs: + suite_name: ${{ steps.normalize.outputs.suite_name }} + build_options: ${{ steps.set_build_options.outputs.build_options }} + steps: + - name: Normalize suite_name + id: normalize + run: echo "suite_name=${GITHUB_TRIGGERING_ACTOR//-/_}" >> "$GITHUB_OUTPUT" + env: + GITHUB_TRIGGERING_ACTOR: ${{ github.triggering_actor }} + - name: Set build options + id: set_build_options + run: | + build_options="" + + if [[ "${{ github.event_name }}" == "workflow_dispatch" ]]; then + docs_branch="${{ github.event.inputs.anemoi_docs_branch }}" + datasets_branch="${{ github.event.inputs.anemoi_datasets_branch }}" + training_branch="${{ github.event.inputs.anemoi_training_branch }}" + + # Default to ref name if docs branch not specified + if [[ -z "$docs_branch" ]]; then + docs_branch="${{ github.ref_name }}" + fi + + build_options="anemoi_docs_branch=$docs_branch anemoi_datasets_branch=$datasets_branch anemoi_training_branch=$training_branch" + fi + + echo "build_options=$build_options" >> "$GITHUB_OUTPUT" + + system_level_test: + needs: prepare + uses: ./.github/workflows/system_level_test.yaml + with: + suite_name: ${{ needs.prepare.outputs.suite_name }} + build_options: ${{ needs.prepare.outputs.build_options }} + secrets: + gh_token: ${{ secrets.GH_REPO_READ_TOKEN }} + troika_user: ${{ secrets.HPC_SYSTEM_TEST_USER }} + ecflow_host: ${{ secrets.HPC_SYSTEM_TEST_HOST }} + ecflow_port: ${{ secrets.HPC_SYSTEM_TEST_PORT }} diff --git a/.github/workflows/system_level_test.yaml b/.github/workflows/system_level_test.yaml new file mode 100644 index 0000000..4b83977 --- /dev/null +++ b/.github/workflows/system_level_test.yaml @@ -0,0 +1,146 @@ +name: system-level-test +description: | + A Github workflow to build, deploy, run and check that a suite + finished well on an ecflow server. + +on: + workflow_call: + inputs: + build_options: + required: true + type: string + suite_name: + required: true + type: string + secrets: + gh_token: + required: true + troika_user: + required: true + ecflow_host: + required: true + ecflow_port: + required: true + +jobs: + deploy: + runs-on: hpc + name: Deploy + steps: + - uses: actions/checkout@v4 + - name: Build and deploy the suite + uses: ./.github/actions/deploy + with: + sbatch_options: | + #SBATCH --job-name=anemoi_test_build_deploy + #SBATCH --time=00:10:00 + #SBATCH --qos=nf + site: hpc-batch + troika_user: ${{ secrets.troika_user }} + github_token: ${{ secrets.gh_token }} + ecflow_host: ${{ secrets.ecflow_host }} + ecflow_port: ${{ secrets.ecflow_port }} + suite_name: ${{ inputs.suite_name }} + build_options: ${{ inputs.build_options }} + run: + runs-on: hpc + name: Run + needs: deploy + steps: + - uses: ecmwf/reusable-workflows/ci-hpc-generic@v2 + name: Run suite from ecflow server + with: + template: | + set -eux + echo "Job is running on ${{ runner.os }}" + + # Load modules of interest + module load ecflow + + # Set ecflow server variables + export ECF_HOST=${{ secrets.ecflow_host }} + export ECF_PORT=${{ secrets.ecflow_port }} + + # Run the suite + ecflow_client --begin=/anemoi_tests + ecflow_client --resume=/anemoi_tests/${{ inputs.suite_name }} + sbatch_options: | + #SBATCH --job-name=anemoi_test_build + #SBATCH --time=00:10:00 + #SBATCH --qos=nf + troika_user: ${{ secrets.troika_user }} + site: hpc-batch + monitor: + runs-on: hpc + name: Monitor + needs: run + steps: + - name: Check that the suite is complete + uses: ecmwf/reusable-workflows/hpc/ecflow/wait-for-ecflow-suite-to-complete@v2 + with: + sbatch_options: | + #SBATCH --job-name=anemoi_test_monitor + #SBATCH --time=01:00:00 + #SBATCH --qos=nf + site: hpc-batch + troika_user: ${{ secrets.troika_user }} + suite_name: anemoi_tests/${{ inputs.suite_name }} + ecflow_host: ${{ secrets.ecflow_host }} + ecflow_port: ${{ secrets.ecflow_port }} + delay: 60 + timeout: 3600 + print: + runs-on: hpc + name: Print + needs: + - deploy + - monitor + if: always() && needs.deploy.result == 'success' + steps: + - name: Print results + uses: ecmwf/reusable-workflows/ci-hpc-generic@v2 + with: + template: | + set -eux + echo "Job is running on ${{ runner.os }}" + + # Load modules of interest + module load wellies/1.2.0 + + mkdir -p /scratch/${{ secrets.troika_user }}/anemoi_tests/${{ inputs.suite_name }} + /home/mlx/bin/tracksuite-print /anemoi_tests/${{ inputs.suite_name }} --host ${{ secrets.ecflow_host }} -f html &> /scratch/${{ secrets.troika_user }}/anemoi_tests/${{ inputs.suite_name }}/summary.md + sbatch_options: | + #SBATCH --job-name=anemoi_test_print_results + #SBATCH --time=00:10:00 + #SBATCH --qos=nf + site: hpc-batch + troika_user: ${{ secrets.troika_user }} + summary: + runs-on: hpc + name: Summary + needs: print + if: always() && needs.print.result == 'success' + steps: + - name: Print summary + run: | + scp ${{ secrets.troika_user }}@hpc-batch:/scratch/${{ secrets.troika_user }}/anemoi_tests/${{ inputs.suite_name }}/summary.md $GITHUB_STEP_SUMMARY + clean: + runs-on: hpc + name: Clean + needs: + - monitor + - summary + if: always() && needs.monitor.result == 'success' + steps: + - name: Clean the suite + uses: ecmwf/reusable-workflows/hpc/ecflow/remove-suite@v2 + with: + sbatch_options: | + #SBATCH --job-name=anemoi_test_clean + #SBATCH --time=00:10:00 + #SBATCH --qos=nf + site: hpc-batch + troika_user: ${{ secrets.troika_user }} + suite_name: anemoi_tests/${{ inputs.suite_name }} + ecflow_host: ${{ secrets.ecflow_host }} + ecflow_port: ${{ secrets.ecflow_port }} diff --git a/.gitignore b/.gitignore index 8e2733b..fe5bfb8 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ _build/ ?.* ~* *.egg-info/ +__pycache__/ diff --git a/docs/contributing/testing.rst b/docs/contributing/testing.rst index 407f670..6249886 100644 --- a/docs/contributing/testing.rst +++ b/docs/contributing/testing.rst @@ -32,6 +32,15 @@ Integration Tests - Important for data processing pipelines and model training workflows. - Integration tests reside in `tests/integration`. +System-level Tests +================== + +- Test the anemoi packages as a whole, including end-to-end workflows, + from dataset creation to model training and inference. +- These tests ensure that the packages work together as expected. +- The system-level test suite is located in the `anemoi-docs` + repository in the `tests/system-level` directory. + *************** Running Tests *************** @@ -45,7 +54,7 @@ by running: We use the ``pytest-skip-slow`` plugin to skip slow tests by default. -To run all unit tests: +To run all **unit tests**: .. code:: bash @@ -57,7 +66,7 @@ To run tests in a specific file: pytest tests/unit/test_specific_feature.py -To run all integration tests, including slow-running tests, use the +To run all **integration tests**, including slow-running tests, use the `--slow` flag. Follow the package-specific instructions. For integration tests in anemoi-training, for instance, ensure that you have GPU available and run: @@ -91,6 +100,19 @@ docs for more information.

+To run **system-level tests**, navigate to the `anemoi-docs` repository +on github and trigger the workflow `on-demand-system-level-test` via the +GitHub Actions tab. + +.. note:: + + Do not trigger system-level tests if another system-level test + workflow is already running under your user account. This is to avoid + ecflow zombies when replacing a suite that is already running. + + If you need to replace a running suite, please wait for it to finish + or cancel it first. + *************** Writing Tests *************** @@ -186,6 +208,18 @@ or `pytest-mock `_. result = my_api_function() assert result == "mocked" +Test Coverage +============= + +We use pytest-cov to measure test coverage. To check coverage: + +.. code:: bash + + pytest --cov=anemoi_training + +Aim for at least 80% coverage for new features, and strive to maintain +or improve overall project coverage. + *************************** Writing Integration Tests *************************** @@ -233,15 +267,58 @@ approach includes: For more details and package-specific examples, please refer to the package-level documentation. -*************** - Test Coverage -*************** - -We use pytest-cov to measure test coverage. To check coverage: - -.. code:: bash - - pytest --cov=anemoi_training - -Aim for at least 80% coverage for new features, and strive to maintain -or improve overall project coverage. +*************************************************** + Adding a Test Case in the System-level Test Suite +*************************************************** + +To add a test case in the system-level test suite, you need to add +config files in the relevant directory in the `anemoi-docs` repository. +The config files should be placed in the `tests/system-level/configs` +directory as explained below. They will constitute tasks in the +system-level test suite. No pyflow knowledge is required to add new test +cases. + +Dataset Creation Test cases +=========================== + +To add a new test case for dataset creation create a new folder in the +`tests/system-level/anemoi_test/configs/datasets` directory. The name of +the folder will be the name of the test case and of the dataset created. +In the folder add + +#. a `dataset_config.yaml` file with the configuration for the dataset + creation. Currently, the only source supported in the test suite is + `mars`. + +#. a `task_config.yaml` file that specifies additional information + required to configure the task in the suite. The task config should + specify the `anemoi_command` to be used to create the dataset -- + typically, "anemoi-datasets create". Hardware overrides (e.g. path to + the `dataset_config.yaml` file) will be set in the suite. + +If you need additional flexibility in configuring your test case, please +open an issue in the `anemoi-docs` repository. + +Model Training Test cases +========================= + +To add a new test case for model training, create a new folder in the +`tests/system-level/anemoi_test/configs/training` directory. The name of +the folder will be the name of the test case. In the folder add + +#. a `training_config.yaml` file with the configuration for the model + training. The configuration should be a full config file, i.e. not + require hydra to build a config based on defaults. The dataset names + should match the names of datasets created in the previous part of + the suite. (If you want to test training based on an existing + anemoi-dataset, consider adding an integration test in the + `anemoi-training` package instead.) + +#. a `task_config.yaml` that specifies additional information required + to configure the task in the suite. The `task_config.yaml` should + contain a list of dataset names that are required for the training + task. The names should match the names of the datasets specified in + the `training_config.yaml`. The `task_config.yaml` should also + specify the `anemoi_command` to be used to run the training -- + typically, "anemoi-training train". Hardware overrides (e.g. path to + the `training_config.yaml` file) will be set in the suite. diff --git a/tests/system-level/README.md b/tests/system-level/README.md new file mode 100644 index 0000000..6537f83 --- /dev/null +++ b/tests/system-level/README.md @@ -0,0 +1,55 @@ +# System-level Tests + +The `anemoi_test` suite provides system-level testing for the anemoi packages, covering the entire workflow from dataset creation to model training and inference. + +For details on adding a test case or triggering suite deployment via GitHub, see the main +[Testing Documentation](https://anemoi.readthedocs.io/en/latest/contributing/testing.html). + + +## Running Tests Locally During Development + +To build and deploy the anemoi_test suite locally, you will need: + - An ecflow server + - [pyflow-wellies](https://pyflow-wellies.readthedocs.io/latest/) version ≥ 1.2.0 + +### Building the Suite + +Run the build script with your desired output directory: + +./build.sh -s OUTPUT_ROOT_SUITE=desired-output-dir + +Set `OUTPUT_ROOT_SUITE` to the path you want to use as the root output directory of the suite, e.g. `$SCRATCH/workdir`. + +This directory is used to store the suite's outputs and to build the virtual environments using uv. Therefore, the suite will only use the UV cache if it is located in the same file system as `OUTPUT_ROOT_SUITE`. + +### Testing Configuration Changes + +If you modify configuration files in `anemoi_test/configs`, you need to point to a committed branch of `anemoi-docs` so that the deployed suite can pull your config files from that branch. + +- Commit your changes to a branch. +- Push the branch to anemoi-docs. +- Run the build script pointing to that branch: + +``` +./build.sh -s OUTPUT_ROOT_SUITE=$SCRATCH/workdir anemoi_docs_branch=name-of-your-branch +``` + +### Additional Build Options + +You can configure additional build options defined in `configs/user.yaml`, e.g. the `anemoi-datasets` branch used to run the tests, by passing additional overrides: + +``` +./build.sh -s OUTPUT_ROOT_SUITE=$SCRATCH/workdir anemoi_datasets_branch=name-of-your-branch +``` + +### Deploying the Suite to the Ecflow Server + +By default, the suite will be built in your home directory `$HOME/pyflow/anemoi_tests/{USER}`. To deploy the suite to the ecflow server, +- Navigate to this directory -- it should contain a definition file `{USER}.def`. +- Follow the [ecflow documentation](https://ecflow.readthedocs.io/en/5.14.1/quickstart.html) to start the server and configure the host. +- Load your suite definition to the server + + ``` + ecflow_client --load={USER}.def + ``` +- Follow the ecflow documentation to start and monitor the suite. diff --git a/tests/system-level/anemoi_test/__init__.py b/tests/system-level/anemoi_test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/system-level/anemoi_test/config.py b/tests/system-level/anemoi_test/config.py new file mode 100644 index 0000000..32e3fe7 --- /dev/null +++ b/tests/system-level/anemoi_test/config.py @@ -0,0 +1,49 @@ +import os + +import wellies as wl + +ROOT_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__))) + + +class Config: + # Configuration templates global variables + global_vars = { + "ROOT": ROOT_DIR, + "NAME": os.path.splitext(os.path.basename(__file__))[0], + } + + def __init__(self, args): + + # Parse the configuration files + options = wl.parse_profiles( + args.profiles, + config_name=args.name, + set_variables=args.set, + ) + + # put everything from the yaml into class variables + self.__dict__.update(options) + + # Ecflow server options + self.ecflow_server = wl.EcflowServer(**options["ecflow_server"]) + + # HPC server options + self.host, submit_variables = wl.get_host(**options["host"]) + + # Optional suite deployment + self.backup_deploy = options.get("backup_deploy", None) + + # Tools + tools = options.get("tools", {}) + self.tools = wl.ToolStore("$LIB_DIR", tools) + + # Static data + static_data = options.get("static_data", {}) + self.static_data = wl.StaticDataStore("$DATA_DIR", static_data) + + # Suite variables + self.suite_variables = { + "SUITE": self.name, + **options.get("ecflow_variables", {}), + **submit_variables, + } diff --git a/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/dataset_config.yaml b/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/dataset_config.yaml new file mode 100644 index 0000000..609ae17 --- /dev/null +++ b/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/dataset_config.yaml @@ -0,0 +1,59 @@ +name: aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing + +description: 'Subset of aifs-ea-an-oper-0001-mars-o96-1979-2023-6h-v8, + for purposes of system-level testing, with reduced time range and parameters. + Data from 2017, resolution O96. From ERA5 data, deterministic. + ' + +attribution: ECMWF + +licence: CC-BY-4.0 + +dates: + start: 2017-01-01 6:00:00 + end: 2017-01-10 18:00:00 + frequency: 6h + group_by: weekly + +input: + join: + - mars: + class: ea + expver: '0001' + grid: o96 + levtype: sfc + param: + - 10u + - 10v + - 2d + - 2t + - mars: + class: ea + expver: '0001' + grid: o96 + level: + - 50 + - 100 + levtype: pl + param: + - q + - t + - accumulations: + class: ea + expver: '0001' + grid: o96 + param: + - cp + - tp + - constants: + param: + - cos_latitude + - cos_longitude + - sin_latitude + - sin_longitude + - cos_julian_day + - cos_local_time + - sin_julian_day + - sin_local_time + - insolation + template: ${input.join.0.mars} diff --git a/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/task_config.yaml b/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/task_config.yaml new file mode 100644 index 0000000..9d5b2ca --- /dev/null +++ b/tests/system-level/anemoi_test/configs/datasets/aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing/task_config.yaml @@ -0,0 +1 @@ +anemoi_command: "anemoi-datasets create" diff --git a/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/dataset_config.yaml b/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/dataset_config.yaml new file mode 100644 index 0000000..1af637d --- /dev/null +++ b/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/dataset_config.yaml @@ -0,0 +1,91 @@ +name: cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing + +description: 'Subset of cerra-rr-an-oper-0001-mars-5p5km-1984-2020-6h-v2-hmsi, + + comprising only 10 days and a reduced set of parameters for the purpose of + + integration tests, + + HARMONIE-ALADIN reanalysis by SMHI on EURO-CORDEX domain (CERRA project) + + Note that the data includes relative humidity (r), not specific humidity (q). + + ' + +attribution: ECMWF + +licence: CC-BY-4.0 + +dates: + end: '2017-01-10T18:00:00' + frequency: 6h + group_by: monthly + start: '2017-01-01T06:00:00' + +input: + join: + - mars: + class: rr + level: + - 50 + - 100 + levtype: pl + origin: se-al-ec + param: + - r + - t + stream: oper + type: an + - mars: + class: rr + levtype: sfc + origin: se-al-ec + param: + - sp + - msl + stream: oper + type: an + - accumulations: + class: rr + levtype: sfc + origin: se-al-ec + param: + - tp + - sf + stream: oper + - constants: + param: + - cos_latitude + - cos_longitude + - sin_latitude + - sin_longitude + - cos_julian_day + - cos_local_time + - sin_julian_day + - sin_local_time + - cos_solar_zenith_angle + template: ${input.join.0.mars} + +output: + chunking: + dates: 1 + ensembles: 1 + dtype: float32 + ensemble_dimension: 2 + flatten_grid: true + order_by: + - valid_datetime: ascending + - param_level: ascending + - number: ascending + remapping: + param_level: '{param}_{levelist}' + statistics: param_level + +statistics: + allow_nans: [] + +build: + additions: true + group_by: monthly + use_grib_paramid: false + variable_naming: default diff --git a/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/task_config.yaml b/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/task_config.yaml new file mode 100644 index 0000000..9d5b2ca --- /dev/null +++ b/tests/system-level/anemoi_test/configs/datasets/cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing/task_config.yaml @@ -0,0 +1 @@ +anemoi_command: "anemoi-datasets create" diff --git a/tests/system-level/anemoi_test/configs/training/basic_check.sh b/tests/system-level/anemoi_test/configs/training/basic_check.sh new file mode 100644 index 0000000..07ca229 --- /dev/null +++ b/tests/system-level/anemoi_test/configs/training/basic_check.sh @@ -0,0 +1,24 @@ +if [[ ! -d "$CHECKPOINT_DIR" ]]; then + echo "❌ Checkpoint directory not found: $CHECKPOINT_DIR" + exit 1 +fi +cd "$CHECKPOINT_DIR" + +run_dirs=(*/) +if [[ ${#run_dirs[@]} -ne 1 ]]; then + echo "❌ Expected exactly 1 run directory, found ${#run_dirs[@]}" + exit 1 +fi +run_dir="${run_dirs[0]}" + +ckpt_files=( "$run_dir"/anemoi-by_epoch-*.ckpt ) +if [[ ${#ckpt_files[@]} -ne 2 ]]; then + echo "❌ Expected 2 training checkpoints anemoi-by_epoch-*.ckpt, found ${#ckpt_files[@]}" + exit 1 +fi + +inf_ckpt_files=( "$run_dir"/inference-anemoi-by_epoch-*.ckpt ) +if [[ ${#inf_ckpt_files[@]} -ne 2 ]]; then + echo "❌ Expected 2 inference checkpoints inference-anemoi-by_epoch-*.ckpt, found ${#inf_ckpt_files[@]}" + exit 1 +fi diff --git a/tests/system-level/anemoi_test/configs/training/global/task_config.yaml b/tests/system-level/anemoi_test/configs/training/global/task_config.yaml new file mode 100644 index 0000000..fdff4cf --- /dev/null +++ b/tests/system-level/anemoi_test/configs/training/global/task_config.yaml @@ -0,0 +1,3 @@ +datasets: + - aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing +anemoi_command: "anemoi-training train" diff --git a/tests/system-level/anemoi_test/configs/training/global/training_config.yaml b/tests/system-level/anemoi_test/configs/training/global/training_config.yaml new file mode 100644 index 0000000..51975fb --- /dev/null +++ b/tests/system-level/anemoi_test/configs/training/global/training_config.yaml @@ -0,0 +1,505 @@ +config_validation: True + +data: + format: zarr + frequency: 6h + timestep: 6h + + forcing: + - "cos_latitude" + - "cos_longitude" + - "sin_latitude" + - "sin_longitude" + - "cos_julian_day" + - "cos_local_time" + - "sin_julian_day" + - "sin_local_time" + - "insolation" + diagnostic: + - tp + - cp + + normalizer: + default: "mean-std" + + std: + - "tp" + + min-max: + max: + none: + - "cos_latitude" + - "cos_longitude" + - "sin_latitude" + - "sin_longitude" + - "cos_julian_day" + - "cos_local_time" + - "sin_julian_day" + - "sin_local_time" + - "insolation" + + imputer: + default: "none" + + processors: + normalizer: + _target_: anemoi.models.preprocessing.normalizer.InputNormalizer + config: ${data.normalizer} + + num_features: null # number of features in the forecast state + +dataloader: + prefetch_factor: 2 + pin_memory: True + + read_group_size: ${hardware.num_gpus_per_model} + + num_workers: + training: 2 + validation: 2 + test: 2 + batch_size: + training: 2 + validation: 2 + test: 2 + + limit_batches: + training: 2 + validation: 2 + test: 20 + + grid_indices: + _target_: anemoi.training.data.grid_indices.FullGrid + nodes_name: ${graph.data} + + dataset: ${hardware.paths.data}/${hardware.files.dataset} + + training: + dataset: ${dataloader.dataset} + start: null + end: 2017-01-07 + frequency: ${data.frequency} + drop: [] + + validation_rollout: 1 # number of rollouts to use for validation, must be equal or greater than rollout expected by callbacks + + validation: + dataset: ${dataloader.dataset} + start: 2017-01-08 + end: null + frequency: ${data.frequency} + drop: [] + + test: + dataset: ${dataloader.dataset} + start: 2022 + end: null + frequency: ${data.frequency} + drop: [] + +diagnostics: + plot: + callbacks: [] + asynchronous: True # Whether to plot asynchronously + datashader: True # Choose which technique to use for plotting + frequency: # Frequency of the plotting + batch: 750 + epoch: 10 + + parameters: + - z_500 + - t_850 + - u_850 + - v_850 + - 2t + - 10u + - 10v + - sp + - tp + - cp + + sample_idx: 0 + + precip_and_related_fields: [tp, cp] + + colormaps: + precip: + _target_: anemoi.training.utils.custom_colormaps.MatplotlibColormapClevels + clevels: ["#ffffff", "#04e9e7", "#019ff4", "#0300f4", "#02fd02", "#01c501", "#008e00", "#fdf802", "#e5bc00", "#fd9500", "#fd0000", "#d40000", "#bc0000", "#f800fd"] + variables: ${diagnostics.plot.precip_and_related_fields} + + debug: + anomaly_detection: False + + profiler: False + + enable_checkpointing: True + checkpoint: + every_n_minutes: + save_frequency: 30 # Approximate, as this is checked at the end of training steps + num_models_saved: 3 # If set to k, saves the 'last' k model weights in the training. + + every_n_epochs: + save_frequency: 1 + num_models_saved: -1 # If set to -1, all checkpoints are kept ensuring runs can be continued/forked at any point in the training process + + every_n_train_steps: + save_frequency: null # Does not scale with rollout + num_models_saved: 0 + + log: + wandb: + enabled: False + offline: False + log_model: False + project: 'Anemoi' + entity: null + # logger options (these probably come with some overhead) + gradients: False + parameters: False + tensorboard: + enabled: False + mlflow: + enabled: False + offline: False + authentication: False + log_model: False + tracking_uri: null + experiment_name: 'anemoi-debug' + project_name: 'Anemoi' + system: False + terminal: True + run_name: null # If set to null, the run name will be the a random UUID + on_resume_create_child: True + expand_hyperparams: # Which keys in hyperparams to expand + - config + http_max_retries: 35 + interval: 100 # passed to trainer.log_every_n_steps + + enable_progress_bar: True + print_memory_summary: False + + benchmark_profiler: + memory: + enabled: False + steps: 5 # wait warmup steps and then do steps (too many steps would lead to a big file) + warmup: 2 + extra_plots: False + trace_rank0_only: False #set to true and it will profile rank 0 only. Reads SLURM_PROC_ID so won't work when not running via Slurm + time: + enabled: True + verbose: False #If true, output every action the profiler caputres, otherwise output a subset defined in PROFILER_ACTIONS at the top of aifs/diagnostics/profiler.py + speed: + enabled: True + system: + enabled: False + model_summary: + enabled: False + snapshot: + enabled: False + steps: 4 # wait warmup steps and then do steps + warmup: 0 + + +datamodule: + _target_: anemoi.training.data.datamodule.AnemoiDatasetsDataModule + +hardware: + + # number of GPUs per node and number of nodes (for DDP) + accelerator: auto + num_gpus_per_node: 1 + num_nodes: 1 + num_gpus_per_model: 1 + paths: + data: dummy_path/ # will be replaced in the suite + output: ${oc.env:PWD}/tmp_output/ # will be replaced in the suite + truncation: null + logs: + base: ${hardware.paths.output}logs/ + wandb: ${hardware.paths.logs.base} + mlflow: ${hardware.paths.logs.base}mlflow/ + tensorboard: ${hardware.paths.logs.base}tensorboard/ + checkpoints: ${hardware.paths.output}checkpoint/ + plots: ${hardware.paths.output}plots/ + profiler: ${hardware.paths.output}profiler/ + graph: ${hardware.paths.output}graphs/ + + files: + dataset: aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing.zarr + graph: dummy.pt + truncation: null + truncation_inv: null + checkpoint: + every_n_epochs: anemoi-by_epoch-epoch_{epoch:03d}-step_{step:06d} + every_n_train_steps: anemoi-by_step-epoch_{epoch:03d}-step_{step:06d} + every_n_minutes: anemoi-by_time-epoch_{epoch:03d}-step_{step:06d} + warm_start: null + + +graph: + overwrite: True + + data: "data" + hidden: "hidden" + + nodes: + # Data nodes + data: + node_builder: + _target_: anemoi.graphs.nodes.AnemoiDatasetNodes # options: AnemoiDatasetNodes, NPZFileNodes + dataset: ${dataloader.dataset} + attributes: ${graph.attributes.nodes} + # Hidden nodes + hidden: + node_builder: + _target_: anemoi.graphs.nodes.TriNodes # options: AnemoiDatasetNodes, NPZFileNodes, TriNodes + resolution: 5 # grid resolution for npz (o32, o48, ...) + attributes: ${graph.attributes.nodes} + + edges: + # Encoder configuration + - source_name: ${graph.data} + target_name: ${graph.hidden} + edge_builders: + - _target_: anemoi.graphs.edges.CutOffEdges # options: KNNEdges, CutOffEdges + cutoff_factor: 0.6 # only for cutoff method + source_mask_attr_name: null + target_mask_attr_name: null + attributes: ${graph.attributes.edges} + # Processor configuration + - source_name: ${graph.hidden} + target_name: ${graph.hidden} + edge_builders: + - _target_: anemoi.graphs.edges.MultiScaleEdges + x_hops: 1 + scale_resolutions: ${graph.nodes.hidden.node_builder.resolution} + source_mask_attr_name: null + target_mask_attr_name: null + attributes: ${graph.attributes.edges} + # Decoder configuration + - source_name: ${graph.hidden} + target_name: ${graph.data} + edge_builders: + - _target_: anemoi.graphs.edges.KNNEdges # options: KNNEdges, CutOffEdges + num_nearest_neighbours: 3 # only for knn method + source_mask_attr_name: null + target_mask_attr_name: null + attributes: ${graph.attributes.edges} + + + attributes: + nodes: + area_weight: + _target_: anemoi.graphs.nodes.attributes.SphericalAreaWeights # options: Area, Uniform + norm: unit-max # options: l1, l2, unit-max, unit-sum, unit-std + fill_value: 0 + edges: + edge_length: + _target_: anemoi.graphs.edges.attributes.EdgeLength + norm: unit-std + edge_dirs: + _target_: anemoi.graphs.edges.attributes.EdgeDirection + norm: unit-std + + post_processors: [] + +model: + num_channels: 16 + cpu_offload: False + + keep_batch_sharded: True + + model: + _target_: anemoi.models.models.AnemoiModelEncProcDec + + layer_kernels: + LayerNorm: + _target_: anemoi.models.layers.normalization.AutocastLayerNorm + Linear: + _target_: torch.nn.Linear + Activation: + _target_: torch.nn.GELU + + processor: + _target_: anemoi.models.layers.processor.GNNProcessor + trainable_size: ${model.trainable_parameters.hidden2hidden} + sub_graph_edge_attributes: ${model.attributes.edges} + num_layers: 16 + num_chunks: 2 + mlp_extra_layers: 0 + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + + encoder: + _target_: anemoi.models.layers.mapper.GNNForwardMapper + trainable_size: ${model.trainable_parameters.data2hidden} + sub_graph_edge_attributes: ${model.attributes.edges} + num_chunks: 1 + mlp_extra_layers: 0 + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + + decoder: + _target_: anemoi.models.layers.mapper.GNNBackwardMapper + trainable_size: ${model.trainable_parameters.hidden2data} + sub_graph_edge_attributes: ${model.attributes.edges} + num_chunks: 1 + mlp_extra_layers: 0 + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + + output_mask: + _target_: anemoi.training.utils.masks.NoOutputMask + + + trainable_parameters: + data: 8 + hidden: 8 + data2hidden: 8 + hidden2data: 8 + hidden2hidden: 8 + + attributes: + edges: + - edge_length + - edge_dirs + nodes: [] + + # Bounding configuration + bounding: #These are applied in order + - _target_: anemoi.models.layers.bounding.ReluBounding #[0, infinity) + variables: + - tp + +training: + + run_id: null + fork_run_id: null + transfer_learning: False # activate to perform transfer learning + load_weights_only: False # only load model weights, do not restore optimiser states etc. + + deterministic: False + + precision: 16-mixed + + multistep_input: 2 + + accum_grad_batches: 1 + + num_sanity_val_steps: 6 + + gradient_clip: + val: 32. + algorithm: value + + swa: + enabled: False + lr: 1.e-4 + + # Optimizer settings + optimizer: + zero: False # use ZeroRedundancyOptimizer ; saves memory for larger models + kwargs: + betas: [0.9, 0.95] + + # select model + model_task: anemoi.training.train.tasks.GraphForecaster + + # select strategy + strategy: + _target_: anemoi.training.distributed.strategy.DDPGroupStrategy + num_gpus_per_model: ${hardware.num_gpus_per_model} + read_group_size: ${dataloader.read_group_size} + + # loss functions + loss_gradient_scaling: False + + # loss function for the model + training_loss: + # loss class to initialise + _target_: anemoi.training.losses.MSELoss + # Scalers to include in loss calculation + scalers: ['pressure_level', 'general_variable', 'nan_mask_weights', 'node_weights'] + ignore_nans: False + + validation_metrics: + # loss class to initialise + mse: + _target_: anemoi.training.losses.MSELoss + scalers: [] + # other kwargs + ignore_nans: True + + variable_groups: + default: sfc + pl: + param: [q, t, u, v, w] + + metrics: + - z_500 + - t_850 + - u_850 + - v_850 + + # length of the "rollout" window (see Keisler's paper) + rollout: + start: 1 + # increase rollout every n epochs + epoch_increment: 0 + # maximum rollout to use + max: 1 + + # Set max_epochs or max_steps. Training stops at the first limit reached. + max_epochs: 2 + max_steps: 150000 + + lr: + warmup: 1000 # number of warmup iterations + rate: 0.625e-4 #local_lr + iterations: ${training.max_steps} # NOTE: When max_epochs < max_steps, scheduler will run for max_steps + min: 3e-7 #Not scaled by #GPU + + + submodules_to_freeze: [] + + scalers: + general_variable: + _target_: anemoi.training.losses.scalers.GeneralVariableLossScaler + weights: + default: 1 + q: 0.6 #1 + t: 6 #1 + u: 0.8 #0.5 + v: 0.5 #0.33 + w: 0.001 + z: 12 #1 + sp: 10 + 10u: 0.1 + 10v: 0.1 + 2d: 0.5 + tp: 0.025 + cp: 0.0025 + + pressure_level: + _target_: anemoi.training.losses.scalers.ReluVariableLevelScaler + group: pl + y_intercept: 0.2 + slope: 0.001 + + # mask NaNs with zeros in the loss function + nan_mask_weights: + _target_: anemoi.training.losses.scalers.NaNMaskScaler + + stdev_tendency: + _target_: anemoi.training.losses.scalers.StdevTendencyScaler + + var_tendency: + _target_: anemoi.training.losses.scalers.VarTendencyScaler + + # Scalers from node attributes + node_weights: + _target_: anemoi.training.losses.scalers.GraphNodeAttributeScaler + nodes_name: ${graph.data} + nodes_attribute_name: area_weight + norm: unit-sum diff --git a/tests/system-level/anemoi_test/configs/training/lam/task_config.yaml b/tests/system-level/anemoi_test/configs/training/lam/task_config.yaml new file mode 100644 index 0000000..e0bb98d --- /dev/null +++ b/tests/system-level/anemoi_test/configs/training/lam/task_config.yaml @@ -0,0 +1,4 @@ +datasets: + - aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing + - cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing +anemoi_command: "anemoi-training train" diff --git a/tests/system-level/anemoi_test/configs/training/lam/training_config.yaml b/tests/system-level/anemoi_test/configs/training/lam/training_config.yaml new file mode 100644 index 0000000..babc343 --- /dev/null +++ b/tests/system-level/anemoi_test/configs/training/lam/training_config.yaml @@ -0,0 +1,557 @@ +config_validation: True + +data: + format: zarr + # Time frequency requested from dataset + frequency: 6h + # Time step of model (must be multiple of frequency) + timestep: 6h + + forcing: + - "cos_latitude" + - "cos_longitude" + - "sin_latitude" + - "sin_longitude" + - "cos_julian_day" + - "cos_local_time" + - "sin_julian_day" + - "sin_local_time" + + diagnostic: + - tp + + normalizer: + default: "mean-std" + + std: + - "tp" + + min-max: + max: + none: + - "cos_latitude" + - "cos_longitude" + - "sin_latitude" + - "sin_longitude" + - "cos_julian_day" + - "cos_local_time" + - "sin_julian_day" + - "sin_local_time" + + imputer: + default: "none" + + # processors including imputers and normalizers are applied in order of definition + processors: + normalizer: + _target_: anemoi.models.preprocessing.normalizer.InputNormalizer + config: ${data.normalizer} + + # Values set in the code + num_features: null # number of features in the forecast state + +dataloader: + prefetch_factor: 2 + pin_memory: True + + read_group_size: ${hardware.num_gpus_per_model} + + num_workers: + training: 2 + validation: 2 + test: 2 + batch_size: + training: 2 + validation: 2 + test: 4 + limit_batches: + training: 2 + validation: 2 + test: 20 + + dataset: + cutout: + - dataset: ${hardware.paths.data}/${hardware.files.dataset} + thinning: 25 + - dataset: ${hardware.paths.data}/${hardware.files.forcing_dataset} + adjust: all + min_distance_km: 0 + + grid_indices: + _target_: anemoi.training.data.grid_indices.MaskedGrid + nodes_name: data + node_attribute_name: indices_connected_nodes + + training: + dataset: ${dataloader.dataset} + start: null + end: 2017-01-07 + frequency: ${data.frequency} + drop: [] + + validation_rollout: 1 # number of rollouts to use for validation, must be equal or greater than rollout expected by callbacks + + validation: + dataset: ${dataloader.dataset} + start: 2017-01-08 + end: 2021 + frequency: ${data.frequency} + drop: [] + + test: + dataset: ${dataloader.dataset} + start: 2022 + end: null + frequency: ${data.frequency} + drop: [] + +diagnostics: + plot: + callbacks: [] + asynchronous: True # Whether to plot asynchronously + datashader: True # Choose which technique to use for plotting + frequency: # Frequency of the plotting + batch: 750 + epoch: 10 + + # Parameters to plot + parameters: + - z_500 + - t_850 + - u_850 + - v_850 + - 2t + - 10u + - 10v + - sp + - tp + - cp + + # Sample index + sample_idx: 0 + + # Precipitation and related fields + precip_and_related_fields: [tp, cp] + + # select special colormap for precip fields + colormaps: + precip: + _target_: anemoi.training.utils.custom_colormaps.MatplotlibColormapClevels + clevels: ["#ffffff", "#04e9e7", "#019ff4", "#0300f4", "#02fd02", "#01c501", "#008e00", "#fdf802", "#e5bc00", "#fd9500", "#fd0000", "#d40000", "#bc0000", "#f800fd"] + variables: ${diagnostics.plot.precip_and_related_fields} + + debug: + # this will detect and trace back NaNs / Infs etc. but will slow down training + anomaly_detection: False + + profiler: False + + enable_checkpointing: True + checkpoint: + every_n_minutes: + save_frequency: 30 # Approximate, as this is checked at the end of training steps + num_models_saved: 3 # If set to k, saves the 'last' k model weights in the training. + + every_n_epochs: + save_frequency: 1 + num_models_saved: -1 # If set to -1, all checkpoints are kept ensuring runs can be continued/forked at any point in the training process + + every_n_train_steps: + save_frequency: null # Does not scale with rollout + num_models_saved: 0 + + log: + wandb: + enabled: False + offline: False + log_model: False + project: 'Anemoi' + entity: null + # logger options (these probably come with some overhead) + gradients: False + parameters: False + tensorboard: + enabled: False + mlflow: + enabled: False + offline: False + authentication: False + log_model: False + tracking_uri: null + experiment_name: 'anemoi-debug' + project_name: 'Anemoi' + system: False + terminal: True + run_name: null # If set to null, the run name will be the a random UUID + on_resume_create_child: True + expand_hyperparams: # Which keys in hyperparams to expand + - config + http_max_retries: 35 + interval: 100 # passed to trainer.log_every_n_steps + + enable_progress_bar: True + print_memory_summary: False + + benchmark_profiler: + memory: + enabled: False + steps: 5 # wait warmup steps and then do steps (too many steps would lead to a big file) + warmup: 2 + extra_plots: False + trace_rank0_only: False #set to true and it will profile rank 0 only. Reads SLURM_PROC_ID so won't work when not running via Slurm + time: + enabled: True + verbose: False #If true, output every action the profiler caputres, otherwise output a subset defined in PROFILER_ACTIONS at the top of aifs/diagnostics/profiler.py + speed: + enabled: True + system: + enabled: False + model_summary: + enabled: False + snapshot: + enabled: False + steps: 4 # wait warmup steps and then do steps + warmup: 0 + + +datamodule: + _target_: anemoi.training.data.datamodule.AnemoiDatasetsDataModule + +hardware: + + # number of GPUs per node and number of nodes (for DDP) + accelerator: auto + num_gpus_per_node: 1 + num_nodes: 1 + num_gpus_per_model: 1 + paths: + data: dummy_path/ # will be replaced in the suite + output: ${oc.env:PWD}/tmp_output/ # will be replaced in the suite + truncation: null + logs: + base: ${hardware.paths.output}logs/ + wandb: ${hardware.paths.logs.base} + mlflow: ${hardware.paths.logs.base}mlflow/ + tensorboard: ${hardware.paths.logs.base}tensorboard/ + checkpoints: ${hardware.paths.output}checkpoint/ + plots: ${hardware.paths.output}plots/ + profiler: ${hardware.paths.output}profiler/ + graph: ${hardware.paths.output}graphs/ + + files: + dataset: cerra-rr-an-oper-0001-mars-5p5km-2017-2017-6h-v3-testing.zarr + forcing_dataset: aifs-ea-an-oper-0001-mars-o96-2017-2017-6h-v8-testing.zarr + graph: dummy.pt + truncation: null + truncation_inv: null + checkpoint: + every_n_epochs: anemoi-by_epoch-epoch_{epoch:03d}-step_{step:06d} + every_n_train_steps: anemoi-by_step-epoch_{epoch:03d}-step_{step:06d} + every_n_minutes: anemoi-by_time-epoch_{epoch:03d}-step_{step:06d} + warm_start: null + +graph: + overwrite: True + + data: "data" + hidden: "hidden" + + nodes: + # Data nodes + data: + node_builder: + _target_: anemoi.graphs.nodes.AnemoiDatasetNodes + dataset: ${dataloader.training.dataset} + attributes: ${graph.attributes.nodes} + # Hidden nodes + hidden: + node_builder: + _target_: anemoi.graphs.nodes.LimitedAreaTriNodes # options: AnemoiDatasetNodes, NPZFileNodes, TriNodes + resolution: 6 # grid resolution for npz (o32, o48, ...) + reference_node_name: ${graph.data} + mask_attr_name: cutout_mask + margin_radius_km: 10 + + edges: + # Encoder configuration + - source_name: ${graph.data} + target_name: ${graph.hidden} + edge_builders: + - _target_: anemoi.graphs.edges.CutOffEdges # options: KNNEdges, CutOffEdges + cutoff_factor: 0.6 # only for cutoff method + source_mask_attr_name: null + target_mask_attr_name: null + - _target_: anemoi.graphs.edges.CutOffEdges # connects only boundary nodes + cutoff_factor: 1.5 # only for cutoff method + source_mask_attr_name: boundary_mask + target_mask_attr_name: null + attributes: ${graph.attributes.edges} + # Processor configuration + - source_name: ${graph.hidden} + target_name: ${graph.hidden} + edge_builders: + - _target_: anemoi.graphs.edges.MultiScaleEdges + x_hops: 1 + scale_resolutions: ${graph.nodes.hidden.node_builder.resolution} + source_mask_attr_name: null + target_mask_attr_name: null + attributes: ${graph.attributes.edges} + # Decoder configuration + - source_name: ${graph.hidden} + target_name: ${graph.data} + edge_builders: + - _target_: anemoi.graphs.edges.KNNEdges # options: KNNEdges, CutOffEdges + num_nearest_neighbours: 3 # only for knn method + source_mask_attr_name: null + target_mask_attr_name: cutout_mask + attributes: ${graph.attributes.edges} + + post_processors: + - _target_: anemoi.graphs.processors.RemoveUnconnectedNodes + nodes_name: data + ignore: cutout_mask # optional + save_mask_indices_to_attr: indices_connected_nodes # optional + + attributes: + nodes: + # Attributes for data nodes + cutout_mask: + _target_: anemoi.graphs.nodes.attributes.CutOutMask + boundary_mask: + _target_: anemoi.graphs.nodes.attributes.BooleanNot + masks: + _target_: anemoi.graphs.nodes.attributes.CutOutMask + area_weight: + _target_: anemoi.graphs.nodes.attributes.MaskedPlanarAreaWeights # options: Uniform + mask_node_attr_name: cutout_mask + norm: unit-max + edges: + edge_length: + _target_: anemoi.graphs.edges.attributes.EdgeLength + norm: unit-std + edge_dirs: + _target_: anemoi.graphs.edges.attributes.EdgeDirection + norm: unit-std + +model: + num_channels: 16 + cpu_offload: False + + keep_batch_sharded: True + + model: + _target_: anemoi.models.models.AnemoiModelEncProcDec + + layer_kernels: + # The layer_kernels can be adjusted per model component, but are defined here for convenience. + LayerNorm: + _target_: torch.nn.LayerNorm + Linear: + _target_: torch.nn.Linear + Activation: + _target_: torch.nn.GELU + QueryNorm: + _target_: anemoi.models.layers.normalization.AutocastLayerNorm + bias: False + KeyNorm: + _target_: anemoi.models.layers.normalization.AutocastLayerNorm + bias: False + + processor: + _target_: anemoi.models.layers.processor.GraphTransformerProcessor + trainable_size: ${model.trainable_parameters.hidden2hidden} + sub_graph_edge_attributes: ${model.attributes.edges} + num_layers: 16 + num_chunks: 2 + mlp_hidden_ratio: 4 # GraphTransformer or Transformer only + num_heads: 16 # GraphTransformer or Transformer only + qk_norm: False + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + + encoder: + _target_: anemoi.models.layers.mapper.GraphTransformerForwardMapper + trainable_size: ${model.trainable_parameters.data2hidden} + sub_graph_edge_attributes: ${model.attributes.edges} + num_chunks: 4 + mlp_hidden_ratio: 4 # GraphTransformer or Transformer only + num_heads: 16 # GraphTransformer or Transformer only + qk_norm: False + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + shard_strategy: "edges" + + decoder: + _target_: anemoi.models.layers.mapper.GraphTransformerBackwardMapper + trainable_size: ${model.trainable_parameters.hidden2data} + sub_graph_edge_attributes: ${model.attributes.edges} + num_chunks: 4 + mlp_hidden_ratio: 4 # GraphTransformer or Transformer only + num_heads: 16 # GraphTransformer or Transformer only + initialise_data_extractor_zero: False + qk_norm: False + cpu_offload: ${model.cpu_offload} + layer_kernels: ${model.layer_kernels} + + output_mask: + _target_: anemoi.training.utils.masks.Boolean1DMask + nodes_name: ${graph.data} + attribute_name: cutout_mask + + trainable_parameters: + data: 8 + hidden: 8 + data2hidden: 8 + hidden2data: 8 + hidden2hidden: 8 # GNN and GraphTransformer Processor only + + attributes: + edges: + - edge_length + - edge_dirs + nodes: [] + + # Bounding configuration + bounding: #These are applied in order + - _target_: anemoi.models.layers.bounding.ReluBounding #[0, infinity) + variables: + - tp + +training: + # resume or fork a training from a checkpoint last.ckpt or specified in hardware.files.warm_start + run_id: null + fork_run_id: null + transfer_learning: False # activate to perform transfer learning + load_weights_only: False # only load model weights, do not restore optimiser states etc. + + # run in deterministic mode ; slows down + deterministic: False + + # miscellaneous + precision: 16-mixed + + multistep_input: 2 + + accum_grad_batches: 1 + + num_sanity_val_steps: 6 + + # clipp gradients, 0 : don't clip, default algorithm: norm, alternative: value + gradient_clip: + val: 32. + algorithm: value + + swa: + enabled: False + lr: 1.e-4 + + # Optimizer settings + optimizer: + zero: False # use ZeroRedundancyOptimizer ; saves memory for larger models + kwargs: + betas: [0.9, 0.95] + + # select model + model_task: anemoi.training.train.tasks.GraphForecaster + + # select strategy + strategy: + _target_: anemoi.training.distributed.strategy.DDPGroupStrategy + num_gpus_per_model: ${hardware.num_gpus_per_model} + read_group_size: ${dataloader.read_group_size} + + loss_gradient_scaling: False + + # loss function for the model + training_loss: + # loss class to initialise + _target_: anemoi.training.losses.MSELoss + # Scalers to include in loss calculation + scalers: ['pressure_level', 'general_variable', 'nan_mask_weights', 'node_weights'] + ignore_nans: False + + validation_metrics: + # loss class to initialise + mse_inside_lam: + _target_: anemoi.training.losses.MSELoss + scalers: ["node_weights"] + # other kwargs + ignore_nans: True + + variable_groups: + default: sfc + pl: [q, t, u, v, w, z] + + metrics: + - z_500 + - t_850 + - u_850 + - v_850 + + # length of the "rollout" window (see Keisler's paper) + rollout: + start: 1 + # increase rollout every n epochs + epoch_increment: 0 + # maximum rollout to use + max: 1 + + # Set max_epochs or max_steps. Training stops at the first limit reached. + max_epochs: 2 + max_steps: 150000 + + lr: + warmup: 1000 # number of warmup iterations + rate: 0.625e-4 #local_lr + iterations: ${training.max_steps} # NOTE: When max_epochs < max_steps, scheduler will run for max_steps + min: 3e-7 #Not scaled by #GPU + + + submodules_to_freeze: [] + + scalers: + general_variable: + _target_: anemoi.training.losses.scalers.GeneralVariableLossScaler + weights: + default: 1 + q: 0.6 #1 + t: 6 #1 + u: 0.8 #0.5 + v: 0.5 #0.33 + w: 0.001 + z: 12 #1 + sp: 10 + 10u: 0.1 + 10v: 0.1 + 2d: 0.5 + tp: 0.025 + cp: 0.0025 + + pressure_level: + _target_: anemoi.training.losses.scalers.ReluVariableLevelScaler + group: pl + y_intercept: 0.2 + slope: 0.001 + + # mask NaNs with zeros in the loss function + nan_mask_weights: + _target_: anemoi.training.losses.scalers.NaNMaskScaler + + stdev_tendency: + _target_: anemoi.training.losses.scalers.StdevTendencyScaler + + var_tendency: + _target_: anemoi.training.losses.scalers.VarTendencyScaler + + # Scalers from node attributes + node_weights: + _target_: anemoi.training.losses.scalers.GraphNodeAttributeScaler + nodes_name: ${graph.data} + nodes_attribute_name: area_weight + norm: "unit-sum" + + limited_area_mask: + _target_: anemoi.training.losses.scalers.GraphNodeAttributeScaler + nodes_name: ${graph.data} + nodes_attribute_name: cutout_mask + norm: null diff --git a/tests/system-level/anemoi_test/nodes.py b/tests/system-level/anemoi_test/nodes.py new file mode 100644 index 0000000..5844f31 --- /dev/null +++ b/tests/system-level/anemoi_test/nodes.py @@ -0,0 +1,174 @@ +import os +from pathlib import Path +from typing import Optional + +import pyflow as pf +import wellies as wl +import yaml + +SUITE_DIR = Path(__file__).resolve().parent + +STATIC_DATA_DIR = Path("$DATA_DIR/anemoi_test_configs") +RESULTS_DIR_TRAINING = Path("$RESULTS_DIR/training") +RESULTS_DIR_DATASETS = Path("$RESULTS_DIR/datasets") + + +def load_yaml(config_file: Path) -> dict: + with open(config_file, "r") as file: + config = yaml.load(file, Loader=yaml.FullLoader) + return config + + +def dict_to_overrides_string(overrides: dict) -> str: + return " ".join(f"{key}={value}" for key, value in overrides.items()) + + +class CreateDatasetFamily(pf.AnchorFamily): + def __init__(self, config, **kwargs): + super().__init__(name="datasets", **kwargs) + dataset_config_dir = SUITE_DIR / "configs/datasets" + + with self: + for folder in os.listdir(dataset_config_dir): + local_config_folder = dataset_config_dir / folder + if not local_config_folder.is_dir(): + continue + if not (local_config_folder / "dataset_config.yaml").exists(): + raise FileNotFoundError(f"Dataset test requires a config file: {folder}/dataset_config.yaml") + try: + task_config = load_yaml(local_config_folder / "task_config.yaml") + except FileNotFoundError as e: + raise FileNotFoundError( + f"Dataset test requires a task config file: {folder}/task_config.yaml" + ) from e + + dataset_cmd = task_config.get("anemoi_command", "anemoi-datasets create") + create_dataset = DatasetTask(folder, config, dataset_cmd=dataset_cmd) + check_dataset = DatasetCheck(folder, create_dataset.output_path) + create_dataset >> check_dataset + + +class DatasetTask(pf.Task): + def __init__(self, name: str, suite_config: dict, dataset_cmd: str = "anemoi-datasets create"): + config_file_path = STATIC_DATA_DIR / "datasets" / name / "dataset_config.yaml" + self.output_path = RESULTS_DIR_DATASETS / (name + ".zarr") + + create_command = dataset_cmd + f" {config_file_path} {self.output_path} --overwrite" + + super().__init__( + name=name.replace("-", "_"), + script=[suite_config.tools.load("datasets_env"), create_command], + ) + + +class DatasetCheck(pf.Task): + def __init__(self, name: str, dataset_path: Path): + check_if_dataset_exists = [f"test -d {dataset_path}", f"test -f {dataset_path}/.zattrs"] + super().__init__(name="check_" + name.replace("-", "_"), script=check_if_dataset_exists) + + +class TrainingFamily(pf.AnchorFamily): + def __init__(self, config, **kwargs): + super().__init__(name="training", **kwargs) + training_config_dir = SUITE_DIR / "configs/training" + + with self: + for folder in os.listdir(training_config_dir): + config_folder = training_config_dir / folder + if not config_folder.is_dir(): + continue + if not (config_folder / "training_config.yaml").exists(): + raise FileNotFoundError(f"Training test requires a config file: {folder}/training_config.yaml") + try: + task_config = load_yaml(config_folder / "task_config.yaml") + except FileNotFoundError as e: + raise FileNotFoundError( + f"Training test requires a task config file: {folder}/task_config.yaml" + ) from e + + training_cmd = task_config.get("anemoi_command", "anemoi-training train") + training = TrainingTask(folder, config, training_cmd=training_cmd) + check_training = TrainingCheck("check_" + folder, RESULTS_DIR_TRAINING / folder / "checkpoint") + training >> check_training + + # Attach required datasets to the training task to set triggers in main family + training.required_datasets = task_config.get("datasets", []) + + +class TrainingTask(pf.Task): + def __init__(self, folder: str, suite_config: dict, training_cmd: str = "anemoi-training train"): + self.required_datasets: Optional[str] = None + + overrides = { + "--config-path": STATIC_DATA_DIR / "training" / folder, + "hardware.paths.output": str(RESULTS_DIR_TRAINING / folder) + + "/", # add trailing slash to ensure checkpoints are in ".../global/checkpoint" + "hardware.paths.data": RESULTS_DIR_DATASETS, + "training.max_epochs": 2, + } + + training_command = training_cmd + " --config-name=training_config " + dict_to_overrides_string(overrides) + + super().__init__( + name=folder, script=[suite_config.tools.load("training_env"), training_command], submit_arguments="gpu_job" + ) + + +class TrainingCheck(pf.Task): + def __init__(self, name: str, checkpoint_path: Path): + checkpoint_checks = pf.FileScript(SUITE_DIR / "configs/training/basic_check.sh") + checkpoint_checks.environment_variable("CHECKPOINT_DIR", str(checkpoint_path)) + super().__init__(name=name, script=checkpoint_checks) + + +class CleanupTask(pf.Task): + def __init__(self, **kwargs): + script = ["rm -rf $OUTPUT_ROOT"] + super().__init__(name="cleanup", script=script, **kwargs) + + +class InitFamily(pf.AnchorFamily): + def __init__(self, config, **kwargs): + super().__init__(name="init", **kwargs) + with self: + deploy_tools = wl.DeployToolsFamily(config.tools) + deploy_data = wl.DeployDataFamily(config.static_data) + + clean_up = CleanupTask() + clean_up >> deploy_tools + clean_up >> deploy_data + + +class MainFamily(pf.AnchorFamily): + def __init__(self, config, **kwargs): + super().__init__(name="main", **kwargs) + with self: + create_fam = CreateDatasetFamily(config) + training_fam = TrainingFamily(config) + + for training_task in [task for task in training_fam.all_tasks if isinstance(task, TrainingTask)]: + if not training_task.required_datasets: + raise ValueError( + f"Training task '{training_task.name}' requires datasets, but none are specified in task_config.yaml." + ) + for dataset in training_task.required_datasets: + dataset_task = dataset.replace("-", "_") + if dataset_task not in [task.name for task in create_fam.all_tasks]: + raise KeyError( + f"Dataset '{dataset}' in training task {training_task.name} not found in dataset test cases. Ensure that all datasets match the name of a dataset test case." + ) + create_fam.find_node(dataset_task) >> training_task + + clean_up = CleanupTask() + create_fam >> clean_up + training_fam >> clean_up # only run cleanup if all tests pass + + +class MainSuite(pf.Family): + def __init__(self, config, **kwargs): + super().__init__(defstatus=pf.state.suspended, **kwargs) + + with self: + f_init = InitFamily(config=config, inlimits=self.work) + f_main = MainFamily(config=config, inlimits=self.work) + f_init >> f_main diff --git a/tests/system-level/build.sh b/tests/system-level/build.sh new file mode 100755 index 0000000..432f8d6 --- /dev/null +++ b/tests/system-level/build.sh @@ -0,0 +1,7 @@ +#!/bin/bash + + +module load wellies/${WELLIES_VERSION:-1.2.0} +module list + +./deploy.py user "$@" diff --git a/tests/system-level/configs/data.yaml b/tests/system-level/configs/data.yaml new file mode 100644 index 0000000..8a26d3f --- /dev/null +++ b/tests/system-level/configs/data.yaml @@ -0,0 +1,7 @@ +static_data: + anemoi_test_configs: + type: git + source: "git@github.com:ecmwf/anemoi-docs.git" + branch: "{anemoi_docs_branch}" + files: + - "tests/system-level/anemoi_test/configs/*" diff --git a/tests/system-level/configs/host.yaml b/tests/system-level/configs/host.yaml new file mode 100644 index 0000000..a0016fb --- /dev/null +++ b/tests/system-level/configs/host.yaml @@ -0,0 +1,19 @@ +host: + hostname: "hpc" + user: "{USER}" + log_directory: "%ECF_HOME%" + workdir: "$TMPDIR" + submit_arguments: + defaults: + memory_per_cpu: 8Gb + job_name: "%FAMILY1:NOT_DEF%_%TASK%" + tmpdir_size: 20Gb + cpu_job: + queue: nf + total_tasks: 1 + gpu_job: + queue: ng + time: "00:20:00" + RAW_PRAGMA: + - "#SBATCH --gpus=1" + - "#SBATCH --gres=gpu:1" diff --git a/tests/system-level/configs/tools.yaml b/tests/system-level/configs/tools.yaml new file mode 100644 index 0000000..b35789c --- /dev/null +++ b/tests/system-level/configs/tools.yaml @@ -0,0 +1,46 @@ +# Put here configuration related to environment and tools +# For a complete list of options, see: +# https://pyflow-wellies.readthedocs.io/latest/config/tools_config/#tool-types +ecflow_variables: + MODULES_VERSION: "new" + +tools: +# ---------- System modules --------------------------------------------------- + modules: + python: + name: python3 + version: 3.10.10-01 + uv: + name: uv + version: 0.6.14 + + packages: + anemoi_datasets: + type: git + source: https://github.com/ecmwf/anemoi-datasets.git + branch: "{anemoi_datasets_branch}" + depends: [uv, UV_CACHE_DIR] + post_script: "uv pip install --upgrade pip && uv pip install -e ." + anemoi_training: + type: git + source: https://github.com/ecmwf/anemoi-core.git + branch: "{anemoi_training_branch}" + depends: [uv, UV_CACHE_DIR] + post_script: "git fetch --unshallow && uv pip install --upgrade pip && uv pip install -e ./training[all,tests] -e ./graphs[all,tests] -e ./models[all,tests]" + + # --------Virtual envs -------------------------------------------------- + environments: + datasets_env: + type: venv + packages: [anemoi_datasets] + depends: [python] + training_env: + type: venv + packages: [anemoi_training] + depends: [python] + + # --------Environment variables ----------------------------------------- + + env_variables: + UV_CACHE_DIR: + value: "{OUTPUT_ROOT_SUITE}/uv_cache" diff --git a/tests/system-level/configs/user.yaml b/tests/system-level/configs/user.yaml new file mode 100644 index 0000000..fdcb9e2 --- /dev/null +++ b/tests/system-level/configs/user.yaml @@ -0,0 +1,37 @@ +# Configuration file for pyflow suite. + +# This file only contains a selection of most common options. +# For more information: https://pyflow-wellies.readthedocs.io/latest/configurations/ + +# -------- Main Suite configuration ------------------------------------------- +group: "anemoi_tests" +name: "{USER}" +labels: + info: Anemoi system-level tests suite, deployed by {USER} for {name} +limits: + work: 20 + +# -------- Ecflow server ------------------------------------------------------ +ecflow_server: + hostname: "localhost" + user: "{USER}" + deploy_dir: "{HOME}/pyflow/{group}/{name}" + +# -------- Global variables ------------------------------------ +# this group can be redefined in other config files. The resulting option will be the concatenation of all +# the variables. Be careful with the order of the variables and the order of the configuration files, +# as the variables will be overriden if they have the same name. + +OUTPUT_ROOT_SUITE: "/lus/h1resw02/project/prepml/ecflow_server/workdirs" +# For the ci workflow, the output root is the prepml workdir so that it is in the same place as prepml suites +# and can share the same uv cache. + +ecflow_variables: + OUTPUT_ROOT: "{OUTPUT_ROOT_SUITE}/testing/{group}/{name}" + LIB_DIR: "%OUTPUT_ROOT%/local" + DATA_DIR: "%OUTPUT_ROOT%/data" + RESULTS_DIR: "%OUTPUT_ROOT%/results" + +anemoi_datasets_branch: "main" # Branch for anemoi-datasets +anemoi_training_branch: "main" # Branch for anemoi-training +anemoi_docs_branch: "main" # Branch for anemoi-docs to be used to deploy the anemoi_tests suite diff --git a/tests/system-level/deploy.py b/tests/system-level/deploy.py new file mode 100755 index 0000000..eeb9941 --- /dev/null +++ b/tests/system-level/deploy.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# isort: skip_file +"""Deploy script for anemoi_test suite. +This has been generated by the wellies-quickstart and the following versions: +{ + "ecflow": "5.13.8", + "pyflow": "3.5.1", + "tracksuite": "0.4.2", + "wellies": "1.2.0" +} + +for options, see: +./deploy.py --help +""" +import logging +from anemoi_test.config import Config +from anemoi_test.nodes import MainSuite + +import pyflow as pf +import wellies as wl +from wellies.show_versions import show_versions + +logger = logging.getLogger("anemoi_test") + + +if __name__ == "__main__": + + # Create parser and parse arguments + parser = wl.get_parser() + args = parser.parse_args() + + # Set log level from args + logging.basicConfig(level=args.log_level, format="[%(levelname)s]%(name)s: %(message)s") + logger.debug(f"deploying with versions: \n{show_versions(as_dict=True)}") + + # Create config object and suite + logger.debug("Creating Config instance") + config = Config(args) + ecflow_server = config.ecflow_server + + logger.debug("Initialising Suite instance") + # Empty top level suite to allow group/user structure on ecflow server + with pf.Suite(config.group, files=ecflow_server.deploy_dir) as suite: + MainSuite( + config, + name=config.name, + host=config.host, + variables=config.suite_variables, + limits=config.limits, + labels=config.labels, + ) + + # Deploy suite scripts and definition file + logger.info(f"Deploying suite to {ecflow_server.hostname}:{ecflow_server.deploy_dir}") + wl.deploy_suite( + suite, + name=config.name, + hostname=ecflow_server.hostname, + user=ecflow_server.user, + deploy_dir=ecflow_server.deploy_dir, + backup_deploy=config.backup_deploy, + build_dir=args.build_dir, + no_prompt=args.y, + no_deploy=args.no_deploy, + message=args.message, + ) diff --git a/tests/system-level/profiles.yaml b/tests/system-level/profiles.yaml new file mode 100644 index 0000000..09ba329 --- /dev/null +++ b/tests/system-level/profiles.yaml @@ -0,0 +1,6 @@ +# This file contains the main configuration profiles +user: + - configs/user.yaml + - configs/host.yaml + - configs/tools.yaml + - configs/data.yaml