Debugging utility for VSCode and SLURM clusters

[frazane]

Is your feature request related to a problem? Please describe.

A common complaint is that while it is relatively easy to use the VSCode debugger from a login node on an HPC cluster, doing so from a compute node can be quite tricky, but there are ways to do it. One possibility is to start an interactive job on a compute node and open a VSCode server on it via remote SSH. However, this is often not possible or allowed, and other workarounds are needed, see for example here: https://github.com/KaijingOfficial/Vscode-Debug-Python-on-Slurm (credit to @dnerini for finding it and testing it).

It would be great if we could integrate this into anemoi-utils. Thoughts?

Describe the solution you'd like

Such a debugging functionality could be brought to the entire anemoi framework by leveraging the centralization of the CLI abstractions provided by anemoi-utils. We could for example inject the debugpy code somewhere in the Command base class or the cli_main and activate it behind a --debugpy flag or similar. Below I provided a small example of how we can inject the debugpy initialization before running any program.

This is an example script (it's just to convey the idea, it won't work in your setup unless you've already set up your launch.json file in a specific way - also, the actual implementation in anemoi-utils hopefully won't be this hacky 😉 ) that you can use to magically set up debugging for any anemoi-x command on a SLURM compute node:

"""Wrapper script to attach debugpy to anemoi programs running on SLURM clusters.

Description:
    This script sets up a debugpy server on the master node of a SLURM job
    and configures VSCode by modifying the .vscode/launch.json file to connect
    to the debugpy server. It then imports and runs the specified anemoi module's
    main function with the provided command-line arguments.

Example usage:
    srun -u python scripts/debugging/attach_slurm.py anemoi-training train ...

    
Inspired from: https://github.com/KaijingOfficial/Vscode-Debug-Python-on-Slurm
"""

import debugpy
import traceback
import time
import os
import json
import importlib
import sys
import argparse

def setup_debug(is_main_process, port=10099):
    master_addr = os.environ["SLURM_NODELIST"].split(",")[0]
    
    # Auto-configure .vscode/launch.json
    launch_json_path = os.path.join(".vscode", "launch.json")
    os.makedirs(os.path.dirname(launch_json_path), exist_ok=True)
        
    with open(launch_json_path, "r") as f:
        existing_config = json.load(f)
    
    if existing_config.get("configurations"):
        for i, config in enumerate(existing_config["configurations"]):
            if config.get("name") == "Debug: attach to SLURM job":
                existing_config["configurations"][i]["connect"].update({
                    "port": port,
                    "host": master_addr
                })
    else:
        raise ValueError("Expected an existing launch.json configuration")
    
    with open(launch_json_path, "w") as f:
        json.dump(existing_config, f, indent=4)

    if is_main_process:  # Master process handler
        print(f"Debug portal active on {master_addr}:{port}", flush=True)
        debugpy.listen((master_addr, port))
        debugpy.wait_for_client()
        print("Debugger linked!", flush=True)


if __name__ == "__main__":


    # parse the command line arguments (as they would normally be used)
    parser = argparse.ArgumentParser(description="Wrapper to attach debugpy to anemoi programs.")
    parser.add_argument("module", type=str, help="The anemoi module to run (e.g., anemoi-training, anemoi-inference, etc.)")
    parser.add_argument("subcommand", type=str, help="The subcommand to run (e.g., train, run, create, etc.)")
    parser.add_argument("extra_args", nargs=argparse.REMAINDER, help="Additional arguments for the subcommand")
    args = parser.parse_args()

    # import the correct entry point module
    module_name = args.module.replace("-", ".")
    entry_module = importlib.import_module(f"{module_name}.__main__")
    main_func = entry_module.main

    # determine if this is the main process
    rank = int(os.environ.get("SLURM_PROCID", "0"))
    is_main_process = (rank == 0)
    
    # set up debugpy
    print(f"Process rank {rank} - Debug mode: {'ON' if is_main_process else 'OFF'}", flush=True)
    setup_debug(is_main_process)

    # execute the main function with the provided arguments
    cmd_args = [args.subcommand] + args.extra_args
    sys.argv = [args.module] + cmd_args

    if is_main_process:
        # rank-0: let exceptions propagate
        main_func()
    else:
        try:
            main_func()
        except Exception as e:
            print(f"Rank {rank} encountered exception: {e}", flush=True)
            traceback.print_exc()
            print(f"Pausing process on rank {rank} to prevent job termination.", flush=True)
            while True:
                time.sleep(1)

Organisation

MeteoSwiss

[HCookie]

I feel as if this is hyper specific to a particular use case, at ECMWF we are able to connect directly via ssh to a compute cluster and can thus use the debugger natively.
If this is needed, I would suggest raising with @ecmwf/anemoi_technical_subgroup

[frazane]

Yeah I agree this is quite dependent on the user's development environment. It could also just be included in some dev documentation as one of the options to use the VSCode debugger (along with the option so ssh into the compute cluster). I do think it's worth putting this somewhere for the benefit of the community, but no strong opinion on the how. So yes, maybe it could be mentioned in the ATS

[dietervdb-meteo]

I think this could be useful for us on LUMI as well.
Since this is relevant for certain users only, and also related to some particular development tool, could this be an example of something we should support through anemoi-plugins?

[bdvllrs]

This looks interesting! Moreover debugpy uses the Debugger Adapter Protocol, so it can be used with any editor that implement this protocol (even vim with a plugin!) and not just VSCode.

[aaron-hopkinson]

Personally I'm not averse to adding an (optional?) dependency on debugpy/something that implements the debugger adapter protocol, but it looks like the example that was posted mixes up three different things: 1) using debugpy; 2) getting the relevant address/port via slurm env vars (and ensuring it runs only on the main process); and 3) setting up some vscode specific configuration. I feel like these are three different issues, and we should try to isolate these as much as possible. For example, I don't think we should get the address/port inside anemoi (via a slurm environment variable), but instead we could pass this in via the command line to try to isolate the slurm dependency here. Similarly, I suspect that the vscode config injection can probably be done outside of the anemoi process (e.g. a wrapper or submission script).

[frazane]

Personally I'm not averse to adding an (optional?) dependency on debugpy/something that implements the debugger adapter protocol, but it looks like the example that was posted mixes up three different things: 1) using debugpy; 2) getting the relevant address/port via slurm env vars (and ensuring it runs only on the main process); and 3) setting up some vscode specific configuration. I feel like these are three different issues, and we should try to isolate these as much as possible. For example, I don't think we should get the address/port inside anemoi (via a slurm environment variable), but instead we could pass this in via the command line to try to isolate the slurm dependency here. Similarly, I suspect that the vscode config injection can probably be done outside of the anemoi process (e.g. a wrapper or submission script).

Yes, these are all reasonable points. That was just a very simple example to convey the idea (it won't even work in your setup unless you have already some entries in the launch.json). But in general, yes, we mix up VSCode debugging (debugpy by default) and SLURM as it is mentioned in the title. The reason for this is that these are, to my knowledge, used by the vast majority of the anemoi community. I fear that trying to make it more general is a bit of a useless effort if in the end almost everyone will use it in the same way (VSCode + SLURM).

[dnerini]

Internally we already have an implementation of this debugging feature, and we’re happy to continue using it within MeteoSwiss. The key point is that thanks to anemoi-utils there is an opportunity to make this debugging capability available to the wider community. Perhaps the better question is whether others outside MeteoSwiss would find it useful. Rather than raising this directly in ATS, I would suggest proposing it for the next Anemoi Community Meeting, so that we can get broader input. If there is sufficient interest, we’d be happy to follow up with a proper PR.

[frazane]

If you are interested in testing this for yourself, here's how to do it.

1. Create an attach_slurm.py script

"""Wrapper script to attach debugpy to anemoi programs running on SLURM clusters.

Description:
    This script sets up a debugpy server on the master node of a SLURM job
    and configures VSCode by modifying the .vscode/launch.json file to connect
    to the debugpy server. It then imports and runs the specified anemoi module's
    main function with the provided command-line arguments.

Example usage:
    srun -u python scripts/debugging/attach_slurm.py anemoi-training train ...

    
Inspired from: https://github.com/KaijingOfficial/Vscode-Debug-Python-on-Slurm
"""

import debugpy
import traceback
import time
import os
import json
import importlib
import sys
import argparse

def setup_debug(is_main_process, port=10099):
    master_addr = os.environ["SLURM_NODELIST"].split(",")[0]
    
    # Auto-configure .vscode/launch.json
    launch_json_path = os.path.join(".vscode", "launch.json")
    os.makedirs(os.path.dirname(launch_json_path), exist_ok=True)
        
    with open(launch_json_path, "r") as f:
        existing_config = json.load(f)
    
    if existing_config.get("configurations"):
        for i, config in enumerate(existing_config["configurations"]):
            if config.get("name") == "Debug: attach to SLURM job":
                existing_config["configurations"][i]["connect"].update({
                    "port": port,
                    "host": master_addr
                })
    else:
        raise ValueError("Expected an existing launch.json configuration")
    
    with open(launch_json_path, "w") as f:
        json.dump(existing_config, f, indent=4)

    if is_main_process:  # 🎯 Master process handler
        print(f"🚨 Debug portal active on {master_addr}:{port}", flush=True)
        debugpy.listen((master_addr, port))
        debugpy.wait_for_client()
        print("🔗 Debugger linked!", flush=True)


if __name__ == "__main__":


    # parse the command line arguments (as they would normally be used)
    parser = argparse.ArgumentParser(description="Wrapper to attach debugpy to anemoi programs.")
    parser.add_argument("module", type=str, help="The anemoi module to run (e.g., anemoi-training, anemoi-inference, etc.)")
    parser.add_argument("subcommand", type=str, help="The subcommand to run (e.g., train, run, create, etc.)")
    parser.add_argument("extra_args", nargs=argparse.REMAINDER, help="Additional arguments for the subcommand")
    args = parser.parse_args()

    # import the correct entry point module
    module_name = args.module.replace("-", ".")
    entry_module = importlib.import_module(f"{module_name}.__main__")
    main_func = entry_module.main

    # determine if this is the main process
    rank = int(os.environ.get("SLURM_PROCID", "0"))
    is_main_process = (rank == 0)
    
    # set up debugpy
    print(f"Process rank {rank} - Debug mode: {'ON' if is_main_process else 'OFF'}", flush=True)
    setup_debug(is_main_process)

    # execute the main function with the provided arguments
    cmd_args = [args.subcommand] + args.extra_args
    sys.argv = [args.module] + cmd_args

    if is_main_process:
        # rank-0: let exceptions propagate
        main_func()
    else:
        try:
            main_func()
        except Exception as e:
            print(f"⚠️ Rank {rank} encountered exception: {e}", flush=True)
            traceback.print_exc()
            print(f"⚠️ Pausing process on rank {rank} to prevent job termination.", flush=True)
            while True:
                time.sleep(1)

2. In your .vscode/launch.json file, add this debugging configuration

        {
            "name": "Debug: attach to SLURM job",
            "type": "debugpy",
            "request": "attach",
            "justMyCode": true,
            "connect": {
                "port": 99999,
                "host": "placeholder"
            }
        }

3. Prepend python attach_slurm.py to any anemoi command in your SLURM submission script, for example:

#!/bin/bash

#SBATCH --job-name="anemoi-training-era5-o96"
#SBATCH --partition=debug
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=24
#SBATCH --mem=256G
#SBATCH --time=00:30:00
#SBATCH --exclusive
#SBATCH --output=slogs/%x/%j.out
#SBATCH --error=slogs/%x/%j.err


source .venv/bin/activate

srun -u python attach_slurm.py anemoi-training train --config-path /path/to/configs/training --config-name config.yaml

4. Submit the script, wait for the job to start, then in the VSCode debugger select the "Debug: attach to SLURM job" option and start it. If it fails, you just have to wait a few more seconds. Then wait until you reach a breakpoint.

[frazane]

Note that if implemented directly in anemoi-utils this would be a lot simpler than the provided script. It could be just an extra function conditionally executed conditionally under

anemoi-utils/src/anemoi/utils/cli.py

Lines 212 to 214 in b22444f

	def cli_main(
	version: str, description: str, commands: dict[str, Command], test_arguments: list[str] | None = None
	) -> None:

so, something close to

if args.setup_debugpy:
    setup_debugpy()

still, the question remains whether this can be made less specific to the "VSCode + SLURM" setup.