foldspace 0.0.1__py3-none-any.whl

foldspace/args.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+# Manage command line options for Foldspace package
+# Authors: Fabio Pitari, Lucia Rodriguez Muñoz, Antonio Memmolo
+
+import argparse, logging, sys, socket
+from datetime import datetime
+if sys.version_info < (3, 10):
+    import importlib_resources as resources
+else:
+    import importlib.resources as resources
+
+import yaml
+from rich.logging import RichHandler
+from .utils import render_j2, rprint  
+from .jobs import list_queued_jobs
+from .prompts import ask_if_resume, ask_resume_id
+#from importlib import metadata # if retrieving version is needed
+
+def get_arguments():
+    """Retrieve variable from option passed at execution time. All the
+    unrecognized options are stored to be passed to the scheduler submit
+    command.
+
+    Returns
+    -------
+    args : object
+        ArgParse object containing native Foldspace options
+    batch_options : list
+        List containing all the unrecognized options (and values), to be
+        passed to submit command
+    """
+    # N.B. do not start using logging from this function since this would
+    # force default logging settings, hiding all the other loggings. You can
+    # use it starting from check_local() function
+
+    parser = argparse.ArgumentParser(add_help=False, description="Foldspace provides access to HTTP services ('tools') on HPC clusters via queued jobs and SSH tunnels. The user supplies cluster connection details (see the --user, --host, and/or --cluster options) and selects a tool available in a cluster directory ('toolbox') using the --tool option. A job is submitted to run the chosen tool and returns the information needed to connect. Default resource values are set for these job submissions, but users will likely want to change them; sbatch/qsub options can therefore be appended to the foldspace command to tune the job. The tool implementation can either be in a default path when --cluster is used, or provided from a custom toolbox by specifying the --toolbox option (mandatory when using --host without --cluster). Foldspace is developed and mantained by Cineca HPC department, please check https://gitlab.hpc.cineca.it/interactive_computing/foldspace for more detailed documentation.")
+
+    parser.add_argument('--help', action='help', # not automated to avoid -h which conflicts with qsub -h
+                        help='Print this help and exit')
+
+    #parser.add_argument('--version', action='version', version=f'%(prog)s {metadata.version("foldspace")}', # retrieve version from package
+    parser.add_argument('--version', action='version', version='', # relies on version printed in the logo
+                        help='Print Foldspace version')
+
+    parser.add_argument('--user', metavar='USERNAME',
+                        dest='user', required=False,
+                        help='Username to use on the HPC cluster')
+
+    parser.add_argument('--host', metavar='HOST',
+                        dest='host', required=False, default=None,
+                        help='Remote ssh host to connect with. You can either use any host configured in your local SSH config file or the login node url coupling it with --user.')
+
+    parser.add_argument('--toolbox', metavar='REMOTE_PATH',
+                        dest='toolbox', required=False, default=None,
+                        help='Remote path on the cluster to which foldspace client will interact. A default value is automatically set when using --cluster option. If you want to deploy your own toolbox on an HPC cluster please read the Foldspace repo documentation')
+
+    parser.add_argument('--tool', metavar='TOOL',
+                        dest='tool', required=False, default=None,
+                        help='Tool to start (use --list-tools to list the available values in the toolbox)')
+
+    parser.add_argument('--list-tools', action='store_true',
+                        dest='list_tools', required=False, default=False,
+                        help='List avaliable tools in the toolbox of the given cluster/host for the --tool option')
+
+    parser.add_argument('--navigator', metavar="SSH_NAVIGATOR",
+                        dest='navigator', required=False, default=None, # real default set in check_local()
+                        help='Choose ssh implementation. Supported profiles are: socket (ssh commands from your system shell, using ssh sockets), nosocket (ssh commands, without sockets), paramiko (experimental, using paramiko and derivative packages). Default is "socket" (but if you are using the --cluster option default might differ for different clusters)')
+
+    parser.add_argument('--cluster', metavar='DEFAULT_CLUSTER',
+                        dest='cluster', required=False, default=None,
+                        help='Alias for predefined set of options for specific clusters. It requires --user and/or --host. You can use --list-clusters for the possible values. Specific HPC clusters might require some preliminary operation from your side to connect via ssh (e.g. retrieving certificates, setting known hosts) which might be executed by default by Foldspace; you can disable this behaviour via --no-startup-script option')
+
+    parser.add_argument('--list-clusters', action='store_true',
+                        dest='list_clusters', required=False, default=False,
+                        help='List avaliable default clusters for the --cluster option')
+
+    parser.add_argument('--no-startup-script', action='store_false',
+                        dest='exec_startup_script', required=False, default=True,
+                        help='If coupled with --cluster, avoids to run the startup script associated with it (you can inspect the scripts using the --list-clusters option)')
+
+    parser.add_argument('--port', metavar='LOCAL_PORT',
+                        dest='port', required=False, default=None, type=int, # real default set in check_local()
+                        help='Local port to be used to expose and access the tool from your browser. Default is a random free port')
+
+    parser.add_argument('--keyfile', metavar="SSH_KEYFILE",
+                        dest='keyfile', required=False, default=None,
+                        help='Path for ssh private key to be used in the ssh connections')
+
+    parser.add_argument('--debug', action='store_true',
+                        dest='debug', required=False, default=False,
+                        help='Run in debug mode')
+
+    parser.add_argument("-/--<any batch option>", metavar="VALUE", nargs="*", default=[],
+                        help="You can add any option of sbatch/qsub command to overwrite/complement default values of the tool's job. Check 'man sbatch'/'man qsub' on your cluster for details.")
+
+    # All the values from known options are stored as attributes into
+    # args; all the unknown options, values and arguments are stored as
+    # strings into the batch_options list. N.B. this means that, in case
+    # of namesake options, foldspace ones prevails, thus keep them as
+    # unambiguous as possible
+    args, batch_options = parser.parse_known_args()
+
+    return args, batch_options
+
+
+def check_local(args, batch_args=[]):
+    """Elaborate all the options from get_arguments() to translate into
+    variables. It might also quit the code if the options are not enough
+    or inconsistent, or if they're simple enough not to require any
+    action from the main code (e.g. --list-clusters)
+
+    Parameters
+    ----------
+    args : object
+        ArgParse object containing parsed options
+    batch_args : list, optional
+        List of options and arguments to be passed to job scheduler at submit
+        time
+
+    Returns
+    -------
+    host
+        SSH cluster host (either just the url or an host in ~/.ssh/config.
+        String.
+    keyfile
+        Identity file to be used in ssh commands. Either string containing
+        path or None.
+    local_port
+        Local port on which expose tools. Integer.
+    batch_options
+        User-defined for sbatch/qsub command in Slurm/PBS schedulers. String.
+    SshNavigator
+        One of the classes from the 'navigators' module in Foldspace
+    startup_script_content
+        Python code to be run before SSH initialization for default
+        clusters (if defined). Either string or None
+    tool
+        Tool to run in the batch job. Either string or None
+    toolbox
+        Remote path containing the toolbox files. String.
+    list_tools
+        Defines if the user asked to list the toolbox tools or not. Boolean
+    user
+        SSH user. Either string or None (in case it is defined in
+        ~/.ssh/config)
+    """
+    # Enable debugging infos when requested
+    if args.debug == True:
+        logging.basicConfig(
+            format="{asctime} {message}",
+            style="{",
+            datefmt="%Y-%m-%d %H:%M",
+            level=logging.DEBUG,
+            handlers=[RichHandler()]
+        )
+
+    # N.B. you can use logging from now on
+    logging.info("Entering in check_local() function")
+
+    # Initialize variables
+    cluster, exec_startup_script, host, keyfile, list_tools, local_port, navigator, tool, toolbox, user = args.cluster, args.exec_startup_script, args.host, args.keyfile, args.list_tools, args.port, args.navigator, args.tool, args.toolbox, args.user
+    batch_options = " ".join(batch_args)
+    startup_script_content = None
+
+    # Debug
+    logging.info("Foldspace options/arguments parsed from options: {raw_options}".format(raw_options={ k:v for k,v in vars(args).items() if '<any batch option>' not in k}))
+    logging.info(f"Scheduler options/arguments parsed from options: {batch_options.strip() or 'none'}")
+
+    # In case --list-clusters option was used, print default clusters
+    # list and exit
+    if args.list_clusters:
+        with resources.files(__package__).joinpath("default_clusters.yml.j2").open("r") as d:
+            clusters = yaml.safe_load(d) # yaml to dictionary, not rendered
+        rprint("Available clusters for the --cluster option:", text_style="cyan", prepend="╰─┬─ ")
+        for cluster in clusters.keys():
+            if cluster != list(clusters.keys())[-1]:
+                rprint(cluster, prepend="  ├─── ")
+            else:
+                rprint(cluster, prepend="  ╰─── ")
+        for cluster in clusters.keys():
+            rprint(f"{clusters[cluster]['description'].strip()}", title=cluster, border_style="cyan")
+            if "startup_script" in clusters[cluster].keys():
+                rprint(f"Using --cluster={cluster} the following script will be automatically run on your computer before connecting to {cluster} via ssh (unless you add the --no-startup-script option):\n\n{clusters[cluster]['startup_script'].strip()}", title=f"Startup script for {cluster}", border_style="yellow")
+        sys.exit(0)
+
+    ### In any other case you need one between --host and --cluster
+    ### If neither --cluster and --host are specified
+    if cluster == None and host ==None:
+        raise InvalidOptions("Missing minimal options (--list-clusters, --cluster or --host), please check --help")
+
+    # In any other case you specified one of the two.
+    # If --cluster is specified (with or without --host)
+    elif cluster != None:
+        # User has to be either defined in ssh host (passed via --host)
+        # or passed via --user, thus exit if both are missing
+        if user == None and host == None:
+            raise InvalidOptions("The --cluster option requires the --user option and/or the --host option.")
+
+        # load yaml file containing default clusters
+        with resources.files(__package__).joinpath("default_clusters.yml.j2").open("r") as d:
+            default_clusters = yaml.safe_load(render_j2(d.read(), **locals())) # yaml to dictionary
+
+        # If cluster is known, set variables accordingly
+        if cluster in default_clusters:
+            # Use default host if not differently specified via
+            # --host option
+            if not host:
+                host = default_clusters[cluster]["host"].strip()
+                logging.info(f"Default host {host} for {cluster} will be used")
+            # Use default toolbox if not differently specified via
+            # --toolbox option
+            if not toolbox:
+                toolbox = default_clusters[cluster]["toolbox"].strip()
+                logging.info(f"Default toolbox {toolbox} for {cluster} will be used")
+            # Use default navigator if specified as cluster defaults and
+            # not differently specified via --navigator option
+            if not navigator and "navigator" in default_clusters[cluster].keys():
+                navigator = default_clusters[cluster]["navigator"].strip()
+                logging.info(f"Default navigator {navigator} for {cluster} will be used")
+            # Dump startup script content if present in cluster defaults
+            # and --no-startup-script is not specified
+            if exec_startup_script and "startup_script" in default_clusters[cluster].keys():
+                startup_script_content = default_clusters[cluster]["startup_script"].strip()
+                logging.info(f"Startup script for {cluster} will be run")
+        # If the cluster is unknown, exit
+        else:
+            raise InvalidOptions(f"Cluster {cluster} is unknown, check --list-clusters for availiable default clusters or specify both --host and --toolbox for custom clusters.")
+
+    # If just --host is specified
+    elif host != None:
+        # You can't go on if the toolbox of a custom cluster is
+        # not specified (all the variables remain unchanged; --user
+        # is not enforced since user might be already specified in a
+        # ~/.ssh/config host)
+        if toolbox is None:
+            raise InvalidOptions(f"You need to specify --toolbox too for the specific toolbox path on {host} host")
+
+    # Set random (free) local port if not specified by user
+    if not local_port:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("localhost", 0)) # 0 is a random free port
+            local_port = s.getsockname()[1]
+            logging.info(f"Local port {local_port} is free, chosen.")
+
+    # Set default value for navigator if not previously set via --cluster
+    # defaults or via --navigator option
+    if not navigator:
+        logging.info("Navigator not set so far, thus the default value 'socket' will be used")
+        navigator = "socket"
+
+    # Import navigator
+    if navigator == "nosocket":
+        from .navigators import NoSocketNavigator as SshNavigator
+    elif navigator == "socket":
+        from .navigators import SocketNavigator as SshNavigator
+    elif navigator == "paramiko":
+        from .navigators import ParamikoNavigator as SshNavigator
+    else:
+        raise InvalidOptions(f"Navigator {navigator} unknown, please check --help for supported values")
+
+    return batch_options, host, keyfile, list_tools, local_port, SshNavigator, startup_script_content, tool, toolbox, user
+
+
+def check_remote(list_tools, navigator, host, tool, toolbox):
+    """Elaborate all the options from get_arguments() and check_local
+    with respect to the remote side on the cluster, i.e. the toolbox,
+    and the user specific permissions to see tools.
+    It might also quit the code if the options are not enough or
+    inconsistent, or if they're simple enough not to require any action
+    from the main code (e.g. --list-tools)
+
+    Parameters
+    ----------
+    navigator : object
+        An object from one of the classes in the 'navigators' module
+    toolbox : str
+        The path of the toolbox on the cluster
+    tool : str or None
+        The tool requested (if requested, see 'list_tools')
+    list_tools : bool
+        Wether the user asked to just list the tools or not
+
+    Returns
+    -------
+    dict
+        Data retrieved remotely from toolbox.yml
+    int or None
+        When restoring, it's the job id to resume
+    """
+    logging.info("Entering in check_remote() function")
+
+    # Validate toolbox
+    run_toolbox = navigator.command(f"cat {toolbox}/toolbox.yml")
+    if run_toolbox['rc'] != 0:
+        rprint(f"Path {toolbox} on {host} doesn't seem to be a valid toolbox", text_style="red", prepend="╰──── ")
+        sys.exit(1)
+       
+    toolbox_data = yaml.safe_load(run_toolbox['out'])
+    
+    # Get user tool access for the given toolbox
+    allowed_tools_cmd = f"""
+    for d in {toolbox}/tools/*; do
+      if [ -d "$d" ] && [ -r "$d" ] && [ -x "$d" ]; then
+        printf '%s\n' "${{d##*/}}"
+      fi  
+    done
+    """
+    allowed_tools_cmd_out = navigator.command(allowed_tools_cmd)
+    allowed_tools = allowed_tools_cmd_out["out"].split()
+    logging.debug(f"Tools enabled for user in the given toolbox: {allowed_tools}")
+    
+    # Filter toolbox_data using user permissions 
+    toolbox_data["tools"] = {
+        tool_name: tool_data
+        for tool_name, tool_data in toolbox_data.get("tools", {}).items()
+        if tool_name in allowed_tools
+        }
+    logging.debug(f"Data parsed from toolbox.yml: {toolbox_data}")
+    
+    # In case --list-tools option was used, print tools list and exit
+    if list_tools:
+        rprint("Available tools for the --tool option:", text_style="cyan", prepend="╰─┬─ ")
+        for tool_name in toolbox_data['tools'].keys():
+            if tool_name != list(toolbox_data['tools'].keys())[-1]:
+                rprint(tool_name, prepend="  ├─── ")
+            else:
+                rprint(tool_name, prepend="  ╰─── ")
+        for tool_name, tool_description in toolbox_data['tools'].items():
+            rprint(tool_description, title=tool_name, border_style="cyan")
+        sys.exit(0)
+
+    # Retrieve the list of past jobs, and which one of them is active
+    resume_id, active_tools = None, {} # init
+    tools_history = yaml.safe_load(navigator.command(f"cat {toolbox_data['user_dir']}/tools_history.yml")['out']) # N.B. returns None if file doesn't exist
+    if tools_history:
+        active_jobs = list_queued_jobs(navigator, toolbox_data['scheduler'])
+        skipped_states = ["COMPLETING"]
+        active_tools = {job_id: state for job_id, state in active_jobs.items() if job_id in set(tools_history.keys()) and state not in skipped_states} # set() increase efficiency if the history is very long
+
+    # If there are no active tools and neither --list-tools nor --tool
+    # were specified, exit
+    if (not active_tools) and (not tool):
+        rprint(f"No tools currently queued on {toolbox_data['cluster_name']} and --tool option was not specified", text_style="red", prepend="╰─── ")
+        sys.exit(0)
+
+    # If the user survived until here, this means that either he/she is
+    # starting a new tool or restoring a tool which is still active.
+    # If user has tools already queued, check what to do
+    if active_tools:
+
+        # List active tools
+        rprint(f"You have tools already queued on the cluster", text_style="yellow", prepend = "├─┬─ ")
+        rprint("", text_style="cyan", prepend = "│ │ ")
+
+        for old_job_id, old_job_state in active_tools.items():
+            old_toolbox = tools_history[old_job_id]['toolbox']
+            old_tool = tools_history[old_job_id]['tool']
+            old_batch_options = tools_history[old_job_id]['batch_options']
+            old_submit_timestamp = datetime.fromisoformat(str(tools_history[old_job_id]['submit_timestamp'])).astimezone().strftime("%A, %B %d, %Y, %H:%M:%S %Z")
+
+            rprint(f"{old_tool} in job {old_job_id} (state: {old_job_state})", text_style="medium_purple", prepend="│ ├─┬─ ")
+            rprint(f"toolbox: [grey63]{old_toolbox}[/grey63]", prepend="│ │ ├─── ")
+            if old_batch_options:
+                rprint(f"batch options: [cyan]{old_batch_options}[/cyan]", prepend="│ │ ├─── ")
+            else:
+                rprint(f"batch options: [grey50]none[/grey50]", prepend="│ │ ├─── ")
+            rprint(f"submitted on: [green]{old_submit_timestamp}[/green]", prepend="│ │ ╰─── ")
+            rprint("", text_style="cyan", prepend = "│ │ ")
+
+        # Resume automatically when --tool was not passed; ask what to
+        # do if it was passed. If user decides not to resume tools,
+        # resume_id keeps staying to None
+        if not tool: # just for better graphic output
+            rprint(f"Your tool is going to be resumed", prepend="│ ╰─ ")
+        if (not tool) or (ask_if_resume() == True):
+            resume_id = ask_resume_id(tools_history, active_tools)
+
+    # If --tool was specified and there are no active jobs, user wants
+    # unequivocally submit a new tool job; in this case it will be
+    # checked if it's present in the toolbox
+    if tool and not resume_id:
+        if tool not in toolbox_data['tools'].keys():
+            rprint(f'Tool {tool} is not available in the toolbox. Available tools: {", ".join(list(toolbox_data["tools"].keys()))}. Please use the --list-tools option if you need details of these tools', text_style="red", prepend="╰──── ")
+            sys.exit(1)
+
+    return toolbox_data, resume_id
+
+
+class InvalidOptions(Exception):
+    "Placeholder to define a specific exception for unusable options"
+    pass

foldspace/default_clusters.yml.j2

@@ -0,0 +1,239 @@
+leonardo:
+  host: login.leonardo.cineca.it
+  toolbox: "/leonardo/prod/opt/tools/foldspace/1.0/"
+  description: |
+    Leonardo is the current Tier-0 system at Cineca. Please take a look to https://www.hpc.cineca.it/systems/hardware/leonardo/ for details.
+  navigator: nosocket
+  startup_script: |
+    import os
+    import shutil
+    import subprocess
+    import sys
+    from pathlib import Path
+
+    #SMALLSTEP_IDENTITY = "{{ user if user else host }}@foldspace"
+    SMALLSTEP_IDENTITY = "foldspace"
+
+    LEONARDO_ADDRESSES = [
+        "login01-ext.leonardo.cineca.it",
+        "login02-ext.leonardo.cineca.it",
+        "login05-ext.leonardo.cineca.it",
+        "login07-ext.leonardo.cineca.it",
+    ]
+
+    KEY_ALGORITHMS = ["rsa", "ecdsa", "ed25519", "dsa"]
+
+    WILDCARD_HOST = "login*.leonardo.cineca.it"
+
+
+    # Stop execution if smallstep binary is not present in the PATH environment variable
+    if shutil.which("step") is None:
+        print(
+            "Error: no Smallstep binary found (step command)",
+            file=sys.stderr,
+        )
+        sys.exit(127)
+
+    # Retrieve smallstep certificate if not already present
+    list_proc = subprocess.run(
+        ["step", "ssh", "list", "--raw", SMALLSTEP_IDENTITY],
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    inspect_proc = subprocess.run(
+        ["step", "ssh", "inspect"],
+        input=list_proc.stdout,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    if "Valid" not in inspect_proc.stdout:
+        print("Running smallstep client to allow ssh connection toward Cineca clusters...")
+        subprocess.run(
+            ["step", "ssh", "login", SMALLSTEP_IDENTITY, "--provisioner", "cineca-hpc"]
+        )
+
+    # Set ssh known hosts properly for leonardo login nodes
+    ssh_dir = Path.home() / ".ssh"
+    previous_umask = os.umask(0o077)
+    try:
+        ssh_dir.mkdir(parents=True, exist_ok=True)
+    finally:
+        os.umask(previous_umask)
+
+    known_hosts = ssh_dir / "known_hosts"
+
+    # remove previous values for leonardo
+    subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", "login.leonardo.cineca.it"])
+    for address in LEONARDO_ADDRESSES:
+        subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", address])
+
+    with open(known_hosts, "a") as known_hosts_file:
+        for keyalg in KEY_ALGORITHMS:
+            for address in LEONARDO_ADDRESSES:
+                scan_proc = subprocess.run(
+                    ["ssh-keyscan", "-t", keyalg, address],
+                    stdout=subprocess.PIPE,
+                    text=True,
+                )
+                known_hosts_file.write(scan_proc.stdout.replace(address, WILDCARD_HOST))
+
+
+g100:
+  host: login.g100.cineca.it
+  toolbox: "/cineca/prod/opt/tools/foldspace/1.0/"
+  description: |
+    Galileo 100 is the current Tier-1 system at Cineca. Please take a look to https://www.hpc.cineca.it/systems/hardware/galileo100/ for details.
+  navigator: nosocket
+  startup_script: |
+    import os
+    import shutil
+    import subprocess
+    import sys
+    from pathlib import Path
+
+    #SMALLSTEP_IDENTITY = "{{ user if user else host }}@foldspace"
+    SMALLSTEP_IDENTITY = "foldspace"
+
+    G100_ADDRESSES = [
+        "login01-ext.g100.cineca.it",
+        "login02-ext.g100.cineca.it",
+        "login03-ext.g100.cineca.it",
+    ]
+
+    KEY_ALGORITHMS = ["rsa", "ecdsa", "ed25519", "dsa"]
+
+    WILDCARD_HOST = "login*.g100.cineca.it"
+
+
+    # Stop execution if smallstep binary is not present in the PATH environment variable
+    if shutil.which("step") is None:
+        print(
+            "Error: no Smallstep binary found (step command)",
+            file=sys.stderr,
+        )
+        sys.exit(127)
+
+    # Retrieve smallstep certificate if not already present
+    list_proc = subprocess.run(
+        ["step", "ssh", "list", "--raw", SMALLSTEP_IDENTITY],
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    inspect_proc = subprocess.run(
+        ["step", "ssh", "inspect"],
+        input=list_proc.stdout,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    if "Valid" not in inspect_proc.stdout:
+        print("Running smallstep client to allow ssh connection toward Cineca clusters...")
+        subprocess.run(
+            ["step", "ssh", "login", SMALLSTEP_IDENTITY, "--provisioner", "cineca-hpc"]
+        )
+
+    # Set ssh known hosts properly for g100 login nodes
+    ssh_dir = Path.home() / ".ssh"
+    previous_umask = os.umask(0o077)
+    try:
+        ssh_dir.mkdir(parents=True, exist_ok=True)
+    finally:
+        os.umask(previous_umask)
+
+    known_hosts = ssh_dir / "known_hosts"
+
+    # remove previous values for g100
+    subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", "login.g100.cineca.it"])
+    for address in G100_ADDRESSES:
+        subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", address])
+
+    with open(known_hosts, "a") as known_hosts_file:
+        for keyalg in KEY_ALGORITHMS:
+            for address in G100_ADDRESSES:
+                scan_proc = subprocess.run(
+                    ["ssh-keyscan", "-t", keyalg, address],
+                    stdout=subprocess.PIPE,
+                    text=True,
+                )
+                known_hosts_file.write(scan_proc.stdout.replace(address, WILDCARD_HOST))
+
+
+pitagora:
+  host: login.pitagora.cineca.it
+  toolbox: "/pitagora/prod/opt/tools/foldspace/1.0/"
+  description: |
+    Pitagora is the new Tier-1 system hosted by CINECA and fully dedicated to the EUROFusion community. Please take a look to https://docs.hpc.cineca.it/hpc/pitagora.html for details.
+  navigator: nosocket
+  startup_script: |
+    import os
+    import shutil
+    import subprocess
+    import sys
+    from pathlib import Path
+
+    #SMALLSTEP_IDENTITY = "{{ user if user else host }}@foldspace"
+    SMALLSTEP_IDENTITY = "foldspace"
+
+    PITAGORA_ADDRESSES = [
+        "login01-ext.pitagora.cineca.it",
+        "login02-ext.pitagora.cineca.it",
+        "login03-ext.pitagora.cineca.it",
+        "login04-ext.pitagora.cineca.it",
+        "login05-ext.pitagora.cineca.it",
+        "login06-ext.pitagora.cineca.it",
+    ]
+
+    KEY_ALGORITHMS = ["rsa", "ecdsa", "ed25519", "dsa"]
+
+    WILDCARD_HOST = "login*.pitagora.cineca.it"
+
+
+    # Stop execution if smallstep binary is not present in the PATH environment variable
+    if shutil.which("step") is None:
+        print(
+            "Error: no Smallstep binary found (step command)",
+            file=sys.stderr,
+        )
+        sys.exit(127)
+
+    # Retrieve smallstep certificate if not already present
+    list_proc = subprocess.run(
+        ["step", "ssh", "list", "--raw", SMALLSTEP_IDENTITY],
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    inspect_proc = subprocess.run(
+        ["step", "ssh", "inspect"],
+        input=list_proc.stdout,
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    if "Valid" not in inspect_proc.stdout:
+        print("Running smallstep client to allow ssh connection toward Cineca clusters...")
+        subprocess.run(
+            ["step", "ssh", "login", SMALLSTEP_IDENTITY, "--provisioner", "cineca-hpc"]
+        )
+
+    # Set ssh known hosts properly for pitagora login nodes
+    ssh_dir = Path.home() / ".ssh"
+    previous_umask = os.umask(0o077)
+    try:
+        ssh_dir.mkdir(parents=True, exist_ok=True)
+    finally:
+        os.umask(previous_umask)
+
+    known_hosts = ssh_dir / "known_hosts"
+
+    # remove previous values for pitagora
+    subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", "login.pitagora.cineca.it"])
+    for address in PITAGORA_ADDRESSES:
+        subprocess.run(["ssh-keygen", "-f", str(known_hosts), "-R", address])
+
+    with open(known_hosts, "a") as known_hosts_file:
+        for keyalg in KEY_ALGORITHMS:
+            for address in PITAGORA_ADDRESSES:
+                scan_proc = subprocess.run(
+                    ["ssh-keyscan", "-t", keyalg, address],
+                    stdout=subprocess.PIPE,
+                    text=True,
+                )
+                known_hosts_file.write(scan_proc.stdout.replace(address, WILDCARD_HOST))