setup_computer_daint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Piz Daint with AiiDA "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "<img width=\"400px\" src=\"https://www.cscs.ch/fileadmin/_processed_/c/3/csm_Piz_Daint_upgrade_26.04.2017_003_ec93735bd1.jpg\"><br>\n",
    "For more information visit <a target=\"_blank\" href=\"https://www.cscs.ch/computers/piz-daint/\">cscs.ch</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "from aiida import load_dbenv, is_dbenv_loaded\n",
    "from aiida.backends import settings\n",
    "if not is_dbenv_loaded():\n",
    "    load_dbenv(profile=settings.AIIDADB_PROFILE)\n",
    "\n",
    "from aiida.orm.querybuilder import QueryBuilder\n",
    "from aiida.orm import load_node, Code, Computer\n",
    "from aiida.common.exceptions import NotExistent, MultipleObjectsError\n",
    "from aiida.backends.utils import get_automatic_user\n",
    "#from aiida.backends.sqlalchemy.models.authinfo import DbAuthInfo\n",
    "from aiida.backends.djsite.db.models import DbAuthInfo\n",
    "\n",
    "import ipywidgets as ipw\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "from subprocess import check_call, check_output, call\n",
    "from os import path\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layout = ipw.Layout(width=\"400px\")\n",
    "style = {\"description_width\":\"150px\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Setup ssh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ssh_keygen():\n",
    "    fn = path.expanduser(\"~/.ssh/id_rsa\")\n",
    "    if not path.exists(fn):\n",
    "        print(\"Creating ssh key pair\")\n",
    "        !  mkdir -p /project/.ssh\n",
    "        ! ssh-keygen -f {fn} -t rsa -N ''\n",
    "    \n",
    "def is_host_known(hostname):\n",
    "    fn = path.expanduser(\"~/.ssh/known_hosts\")\n",
    "    if not path.exists(fn):\n",
    "        return False\n",
    "    return call([\"ssh-keygen\", \"-F\", hostname]) == 0\n",
    "\n",
    "def make_host_known(hostname, proxycmd=[]):\n",
    "    fn = path.expanduser(\"~/.ssh/known_hosts\")\n",
    "    print(\"Adding keys from %s to %s\"%(hostname, fn))\n",
    "    hashes = check_output(proxycmd+[\"ssh-keyscan\", \"-H\", hostname])\n",
    "    with open(fn, \"a\") as f:\n",
    "        f.write(hashes)\n",
    "\n",
    "def can_login(hostname, username):\n",
    "    userhost = username+\"@\"+hostname\n",
    "    print(\"Trying ssh \"+userhost+\"... \", end='')\n",
    "    ret = call([\"ssh\", userhost, \"true\"])\n",
    "    print(\"Ok\" if ret==0 else \"Failed\")\n",
    "    return ret==0\n",
    "\n",
    "def print_pubkey(hostname, username):\n",
    "    fn = path.expanduser(\"~/.ssh/id_rsa.pub\")\n",
    "    key = open(fn).read().strip()\n",
    "    print('ssh '+username+'@'+hostname+' \\'echo \"'+key+'\" >> ~/.ssh/authorized_keys\\'')\n",
    "\n",
    "def is_in_config(hostname):\n",
    "    fn = path.expanduser(\"~/.ssh/config\")\n",
    "    if not path.exists(fn):\n",
    "        return False\n",
    "    cfglines = open(fn).read().split(\"\\n\")\n",
    "    return \"Host \"+hostname in cfglines\n",
    "\n",
    "def write_ssh_config(targethost, username, proxyhost):\n",
    "    fn = path.expanduser(\"~/.ssh/config\")\n",
    "    print(\"Adding section to \"+fn)\n",
    "    with open(fn, \"a\") as f:\n",
    "        f.write(\"Host \"+targethost+\"\\n\")\n",
    "        f.write(\"User \"+username+\"\\n\")\n",
    "        f.write(\"ProxyCommand ssh \"+username+\"@\"+proxyhost+\" netcat %h %p\\n\")\n",
    "        \n",
    "def fetch_slurm_accounts(hostname, username):\n",
    "    userhost = username+\"@\"+hostname\n",
    "    print(\"Fetching slurm accounts\")\n",
    "    \n",
    "    sbucheck = check_output([\"ssh\", username+\"@daint.cscs.ch\", \"sbucheck\"])\n",
    "    slurm_accounts = re.findall(\"\\* (\\w+):.*constraints: (\\w+)\", sbucheck)\n",
    "    options = [(\"Please select a slurm account\", False)]\n",
    "    for l, c in slurm_accounts:\n",
    "        options.append([\"%s (%s)\"%(l,c), (l,c)])\n",
    "    drop_account.options = options\n",
    "\n",
    "def setup_ssh(proxyhost, targethost, username):\n",
    "    ssh_keygen()\n",
    "    \n",
    "    # setup proxy first\n",
    "    if not is_host_known(proxyhost):\n",
    "        make_host_known(proxyhost)\n",
    "\n",
    "    if not can_login(proxyhost, username):\n",
    "        print(\"Please execute the following command and then try again:\\n\")\n",
    "        print_pubkey(proxyhost, username)\n",
    "        return False\n",
    "                \n",
    "    # now setup target host\n",
    "    if not is_host_known(targethost):\n",
    "        make_host_known(targethost, proxycmd=['ssh', username+\"@\"+proxyhost])\n",
    "    \n",
    "    if not is_in_config(targethost):\n",
    "        write_ssh_config(targethost, username, proxyhost)\n",
    "         \n",
    "    # final check\n",
    "    if can_login(targethost, username):\n",
    "        fetch_slurm_accounts(targethost, username)\n",
    "        print(\"Automatic ssh setup successful :-)\")\n",
    "        return True\n",
    "    else:\n",
    "        print(\"Automatic ssh setup failed, sorry :-(\")\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_setup_ssh(b):\n",
    "    with setup_ssh_out:\n",
    "        clear_output()\n",
    "        if len(inp_username.value.strip())==0:\n",
    "            print(\"Please enter your CSCS username\")\n",
    "            return\n",
    "        setup_ssh(\"ela.cscs.ch\", \"daint.cscs.ch\", inp_username.value)\n",
    "\n",
    "inp_username = ipw.Text(description=\"CSCS user name:\", layout=layout, style=style)\n",
    "btn_setup_ssh = ipw.Button(description=\"Setup ssh\")\n",
    "btn_setup_ssh.on_click(on_setup_ssh)\n",
    "setup_ssh_out = ipw.Output()\n",
    "\n",
    "display(inp_username,btn_setup_ssh, setup_ssh_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Setup AiiDA Computer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/aiidateam/aiida_core/blob/develop/aiida/cmdline/commands/computer.py#L400\n",
    "\n",
    "def setup_computer(computer_name, username, account, partition):\n",
    "    try:\n",
    "        computer = Computer.get(computer_name)\n",
    "        print(\"A computer called {} already exists.\".format(computer_name))\n",
    "        return\n",
    "    except NotExistent:\n",
    "        pass\n",
    "\n",
    "    print(\"Creating new computer with name '{}'\".format(computer_name))\n",
    "    computer = Computer(name=computer_name)\n",
    "    computer.set_hostname(\"daint.cscs.ch\")\n",
    "    computer.set_description(\"Named after Piz Daint, a prominent peak in Grisons that overlooks the Fuorn pass, this supercomputer is a hybrid Cray XC50/XC40 system and is the flagship system for national HPC Service.\")\n",
    "    computer.set_enabled_state(True)\n",
    "    computer.set_transport_type(\"ssh\")\n",
    "    computer.set_scheduler_type(\"slurm\")\n",
    "    computer.set_workdir(\"/scratch/snx3000/\"+username+\"/aiida_run_\"+computer_name)\n",
    "    # set_mpirun_command() must be called after set_scheduler_type()\n",
    "    #cmd = \"srun -n {tot_num_mpiprocs} -c $SLURM_CPUS_PER_TASK --cpu_bind=rank --hint=nomultithread\"\n",
    "    cmd = \"srun -n {tot_num_mpiprocs}\"\n",
    "    computer.set_mpirun_command(cmd.split())\n",
    "    ncpus = 12 if partition == \"gpu\" else 36\n",
    "    computer.set_default_mpiprocs_per_machine(ncpus)\n",
    "    prepend_text  = \"### computer prepend_text start ###\\n\"\n",
    "    prepend_text += \"#SBATCH --account=%s\\n\"%account\n",
    "    prepend_text += \"#SBATCH --constraint=%s\\n\"%partition\n",
    "    prepend_text += \"#SBATCH --cpus-per-task=1\\n\"\n",
    "    prepend_text += \"export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK\\n\"\n",
    "    prepend_text += \"source $MODULESHOME/init/bash\\n\"\n",
    "    prepend_text += \"module load daint-%s\\n\"%partition\n",
    "    prepend_text += \"ulimit -s unlimited\\n\"\n",
    "    if partition == \"gpu\":\n",
    "        prepend_text += \"export CRAY_CUDA_MPS=1\\n\"\n",
    "    prepend_text += \"### computer prepend_text end ###\\n\"\n",
    "    computer.set_prepend_text(prepend_text)\n",
    "    computer.store()\n",
    "    \n",
    "    # create DbAuthInfo\n",
    "    authparams = {\n",
    "        'compress': True,\n",
    "        'gss_auth': False,\n",
    "        'gss_deleg_creds': False,\n",
    "        'gss_host': 'daint.cscs.ch',\n",
    "        'gss_kex': False,\n",
    "        'key_policy': 'RejectPolicy',\n",
    "        'load_system_host_keys': True,\n",
    "        'port': 22,\n",
    "        'proxy_command': 'ssh '+username+'@ela.cscs.ch netcat daint.cscs.ch 22',\n",
    "        'timeout': 60,\n",
    "        'username': username,\n",
    "    }\n",
    "    aiidauser = get_automatic_user()\n",
    "    authinfo = DbAuthInfo(dbcomputer=computer.dbcomputer, aiidauser=aiidauser)\n",
    "    authinfo.set_auth_params(authparams)\n",
    "    authinfo.save()\n",
    "    \n",
    "    ! verdi computer show {computer_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_account_change(c):\n",
    "    if drop_account.value:\n",
    "        inp_compname.value = inp_compname.value.rsplit(\"-\",1)[0] + \"-\" + drop_account.value[0]\n",
    "\n",
    "def on_setup_computer(b):\n",
    "    with setup_comp_out:\n",
    "        clear_output()\n",
    "        if len(inp_compname.value.strip())==0:\n",
    "            print(\"Please enter computer name\")\n",
    "            return\n",
    "        if not drop_account.value:\n",
    "            print(\"Plese select a slurm account\")\n",
    "            return\n",
    "        account, partition = drop_account.value\n",
    "        setup_computer(inp_compname.value, inp_username.value, account, partition)\n",
    "\n",
    "    \n",
    "inp_compname = ipw.Text(description=\"Computer name:\", value=\"daint\", layout=layout, style=style)\n",
    "drop_account = ipw.Dropdown(description=\"Slurm Account:\", options={\"first run ssh setup\":False}, style=style, layout=layout )\n",
    "drop_account.observe(on_account_change, names='value')\n",
    "btn_setup_comp = ipw.Button(description=\"Setup Computer\")\n",
    "btn_setup_comp.on_click(on_setup_computer)\n",
    "setup_comp_out = ipw.Output()\n",
    "display(inp_compname, drop_account, btn_setup_comp, setup_comp_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Test Computer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def on_test_computer(b):\n",
    "    with test_out:\n",
    "        clear_output()\n",
    "        if len(inp_compname.value.strip())==0:\n",
    "            print(\"Please enter computer name\")\n",
    "            return\n",
    "        ! verdi computer test {inp_compname.value}\n",
    "\n",
    "test_out = ipw.Output()\n",
    "btn_test_comp = ipw.Button(description=\"Test Computer\")\n",
    "btn_test_comp.on_click(on_test_computer)\n",
    "display(btn_test_comp, test_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Setup codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_codes(computer_name, partition):\n",
    "    computer = Computer.get(computer_name)\n",
    "    \n",
    "    qe_version = \"6.1\"\n",
    "    if partition == \"mc\":\n",
    "        qe_path = \"/apps/daint/UES/jenkins/6.0.UP04/mc/easybuild/software/QuantumESPRESSO/6.1.0-CrayIntel-17.08/bin/\"\n",
    "        cp2k_path = \"/apps/daint/UES/jenkins/6.0.UP04/mc/easybuild/software/CP2K/5.1-CrayGNU-17.08/bin/\"\n",
    "    elif partition == \"gpu\":\n",
    "        qe_path = \"/apps/daint/UES/jenkins/6.0.UP04/gpu/easybuild/software/QuantumESPRESSO/6.1.0-CrayIntel-17.08/bin/\"\n",
    "        cp2k_path = \"/apps/daint/UES/jenkins/6.0.UP04/gpu/easybuild/software/CP2K/5.1-CrayGNU-17.08-cuda-8.0/bin/\"\n",
    "    else:\n",
    "        raise(\"unkown partition \"+partition)\n",
    "    \n",
    "    # CP2K\n",
    "    cp2k_version = '6.0'\n",
    "    cp2k_revision = '18265'\n",
    "    code_label = \"cp2k_{}_{}\".format(cp2k_version, cp2k_revision)\n",
    "    code_full_name = \"{}@{}\".format(code_label, computer_name)\n",
    "    try:\n",
    "        Code.get_from_string(code_full_name)\n",
    "    except (NotExistent, MultipleObjectsError):            \n",
    "        code = Code(remote_computer_exec=(computer, cp2k_path + \"cp2k.psmp\"))\n",
    "        code.label = code_label\n",
    "        code.description = \"CP2K Code\"\n",
    "        code.set_input_plugin_name(\"cp2k\")\n",
    "        prepend_text  = \"### code prepend_text start ###\\n\"\n",
    "        prepend_text += \"module load CP2K\\n\"\n",
    "        prepend_text += \"export CP2K_DATA_DIR=/apps/daint/UES/jenkins/6.0.UP04/mc/easybuild/software/CP2K/5.1-CrayGNU-17.08/data\\n\"\n",
    "        prepend_text += \"### code prepend_text end ###\\n\"\n",
    "        code.set_prepend_text(prepend_text)\n",
    "        code.set_append_text(\"\")\n",
    "        code._reveal()\n",
    "        code.store()\n",
    "        ! verdi code show \"{code_full_name}\"\n",
    "    else:\n",
    "        print(\"Code '{}' already existent\".format(code_full_name))        \n",
    "    \n",
    "    # QuantumESPRESSO pw\n",
    "    for codename in ['pw', 'pp', 'ph', 'projwfc', 'matdyn', 'q2r']:\n",
    "        code_label = \"{}-{}\".format(codename, qe_version)\n",
    "        code_full_name = \"{}@{}\".format(code_label, computer_name)\n",
    "        try:\n",
    "            Code.get_from_string(code_full_name)\n",
    "        except (NotExistent, MultipleObjectsError):            \n",
    "            code = Code(remote_computer_exec=(computer, qe_path + \"{}.x\".format(codename)))\n",
    "            code.label = code_label\n",
    "            code.description = \"Quantum ESPRESSO {}.x version {}\".format(codename, qe_version)\n",
    "            code.set_input_plugin_name(\"quantumespresso.{}\".format(codename))\n",
    "            prepend_text  = \"### code prepend_text start ###\\n\"\n",
    "            prepend_text += \"module load QuantumESPRESSO/6.1.0-CrayIntel-17.08\\n\"\n",
    "            prepend_text += \"### code prepend_text end ###\\n\"\n",
    "            code.set_prepend_text(prepend_text)\n",
    "            code.set_append_text(\"\")\n",
    "            code._reveal()\n",
    "            code.store()\n",
    "            ! verdi code show \"{code_full_name}\"\n",
    "        else:\n",
    "            print(\"Code '{}' already existent\".format(code_full_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def on_setup_codes(b):\n",
    "    with setup_code_out:\n",
    "        clear_output()\n",
    "        partition = drop_account.value[1]\n",
    "        setup_codes(inp_compname.value, partition)\n",
    "\n",
    "setup_code_out = ipw.Output()\n",
    "btn_setup_codes = ipw.Button(description=\"Setup Codes\")\n",
    "btn_setup_codes.on_click(on_setup_codes)\n",
    "display(btn_setup_codes, setup_code_out)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}