notebooks/sample_submit_ptm_linux_batch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import dmsbatch\n",
    "from dmsbatch import create_batch_client, create_blob_client\n",
    "import datetime\n",
    "import logging\n",
    "#logger = logging.getLogger()\n",
    "#logger.setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First create a batch client from the config file\n",
    "\n",
    "The config file is described in the [README](../README.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = create_batch_client('../tests/data/dmsbatch.config')\n",
    "blob_client = create_blob_client('../tests/data/dmsbatch.config')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Application packages\n",
    "To copy large files and programs it is best to zip (or targz) them and upload them as application packages\n",
    "\n",
    "Application packages are setup separately in either azure management apis or from the web console or cli tool\n",
    "\n",
    "These are referenced here by their name and version\n",
    "e.g. DSM2, python and other programs\n",
    "\n",
    "One extra field (last one) is the path within the zip file where the executables can be found. These are used later to setup the PATH varible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app_pkgs = [('dsm2linux', '8.2.8449db2', 'DSM2-8.2.8449db2-Linux/bin'), ('dsm2_linux_rpms','1.0.0','')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show vms available\n",
    "\n",
    "https://docs.microsoft.com/en-us/azure/virtual-machines/fsv2-series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#display(client.skus_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create or resize existing pool\n",
    "If the pool doesn't exist it will create it\n",
    "If the pool exists, it will resize to the second arg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_name = 'ptmlinuxpool'\n",
    "# Intel repos deny request if multiple machines try to run these install scripts at the same time.\n",
    "# Create an application package with the rpms zipped\n",
    "# yum-config-manager --add-repo https://yum.repos.intel.com/2019/setup/intel-psxe-runtime-2019.repo\n",
    "# rpm --import https://yum.repos.intel.com/2019/setup/RPM-GPG-KEY-intel-psxe-runtime-2019\n",
    "# yumdownloader --resolve --destdir rpms glibc.i686 libstdc++.i686 glibc.x86_64 libstdc++.x86_64 intel-icc-runtime-32bit intel-ifort-runtime-32bit libgfortran4\n",
    "# zip -r dsm2_linux_rpms.zip rpms\n",
    "# Upload the dsm2_linux_rpms.zip to the Application Package using Batch Explorer or via Web Portal\n",
    "pool_start_cmds = ['printenv', 'yum localinstall --nogpgcheck $AZ_BATCH_APP_PACKAGE_dsm2_linux_rpms_1_0_0/rpms/*.rpm -y']\n",
    "#'yum install -y glibc.i686 libstdc++.i686 glibc.x86_64 libstdc++.x86_64',# --setopt=protected_multilib=false',\n",
    "#'yum-config-manager --add-repo https://yum.repos.intel.com/2019/setup/intel-psxe-runtime-2019.repo',\n",
    "#'rpm --import https://yum.repos.intel.com/2019/setup/RPM-GPG-KEY-intel-psxe-runtime-2019',\n",
    "#'yum install -y intel-icc-runtime-32bit intel-ifort-runtime-32bit']\n",
    "client.create_pool(pool_name,\n",
    "                    1,\n",
    "                    app_packages=[(app,version) for app,version,_ in app_pkgs], \n",
    "                    vm_size='standard_f32s_v2', \n",
    "                    tasks_per_vm=32,\n",
    "                    os_image_data=('openlogic', 'centos', '7_8'),\n",
    "                    start_task_cmd=client.wrap_commands_in_shell(pool_start_cmds, ostype='linux'),\n",
    "                    start_task_admin=True,\n",
    "                    elevation_level='admin'\n",
    "                    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autoscaling Formula to use for pool\n",
    "This can be added manually via the console or batch explorer in the resizing section.\n",
    "```\n",
    "// In this example, the pool size is adjusted based on the number of tasks in the queue. Note that both comments and line breaks are acceptable in formula strings.\n",
    "numCores = 32;\n",
    "// Get pending tasks for the past 15 minutes.\n",
    "$samples = $ActiveTasks.GetSamplePercent(TimeInterval_Minute * 1);\n",
    "// If we have fewer than 70 percent data points, we use the last sample point, otherwise we use the maximum of last sample point and the history average.\n",
    "$tasks = $samples < 70 ? max(0, $ActiveTasks.GetSample(1)) : max( $ActiveTasks.GetSample(1), avg($ActiveTasks.GetSample(TimeInterval_Minute * 1)));\n",
    "// If number of pending tasks is not 0, set targetVM to pending tasks, otherwise half of current dedicated nodes.\n",
    "$targetVMs = $tasks > 0 ? $tasks : max(0, $TargetLowPriorityNodes / numCores);\n",
    "// The pool size is capped at 20, if target VM value is more than that, set it to 20. This value should be adjusted according to your use case.\n",
    "cappedPoolSize = 60;\n",
    "$TargetLowPriorityNodes = max(0, min($targetVMs, cappedPoolSize));\n",
    "// Set node deallocation mode - keep nodes active only until tasks finish\n",
    "$NodeDeallocationOption = taskcompletion;\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create job on pool or fail if it exists\n",
    "Jobs are containers of tasks (things that run on nodes (machines) in the pool). If this exists, the next line will fail"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study_dir = 'X:/Share/xwang/DCP/neutrally_buoyant_particles/pa6k_2020'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tidefile = 'X:/Share/DSM2/v821/studies_dcp_2020/pa6k_2020/output/DCP_PA6K_2020.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study_prefix = os.path.basename(os.path.dirname(study_dir))+'/'+os.path.basename(study_dir)\n",
    "study_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "container_name='ptmbatch'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(container_name,'%s/DCP_EX.h5'%(study_prefix),tidefile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UPLOAD=True\n",
    "if UPLOAD:\n",
    "    # slow - 9 mins so use max_connections > 2 (default). Using 12 which seems to be a good fit here\n",
    "    blob_client.upload_file_to_container(container_name,'%s/DCP_EX.h5'%(study_prefix),tidefile,max_connections=10)\n",
    "#input_tidefile = client.create_input_file_spec('ptmnbjob',blob_prefix='DCP_EX.h5',file_path='.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = study_prefix.replace('/','_')\n",
    "job_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "copy_tidefile_task = client.create_task_copy_file_to_shared_dir(container_name,'%s/DCP_EX.h5'%(study_prefix),file_path='.',ostype='linux')\n",
    "client.create_job(job_name,pool_name,prep_task=copy_tidefile_task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if UPLOAD:\n",
    "    blob_client.zip_and_upload(container_name,study_prefix,study_dir,30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a task\n",
    "This uses the application package as pre -set up. If not, create one https://docs.microsoft.com/en-us/azure/batch/batch-application-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_ptm_single_task(task_name, run_no, task_prefix, study_prefix, envvars):\n",
    "    input_file = client.create_input_file_spec(container_name,blob_prefix='%s/%s.zip'%(study_prefix,os.path.basename(study_dir)),file_path='.')\n",
    "    #std_out_files = client.create_output_file_spec(\n",
    "    #    '../std*.txt', output_dir_sas_url, blob_path=f'{study_prefix}/{task_prefix}/{task_name}')\n",
    "    permissions = dmsbatch.commands.azureblob.BlobPermissions.WRITE\n",
    "    output_dir_sas_url = blob_client.get_container_sas_url(container_name, permissions)\n",
    "    output_dir = client.create_output_file_spec(\n",
    "        f'{run_no}/*', output_dir_sas_url, blob_path=f'{study_prefix}/{task_prefix}/{task_name}/{run_no}')\n",
    "    set_path_string = client.set_path_to_apps(app_pkgs, ostype='linux')\n",
    "    zip_fname = os.path.basename(study_dir)+'.zip'\n",
    "    cmd_string = client.wrap_cmd_with_app_path(\n",
    "        f\"\"\"\n",
    "        source /opt/intel/psxe_runtime/linux/bin/compilervars.sh ia32;\n",
    "        {set_path_string};\n",
    "        cd {study_prefix};\n",
    "        unzip {zip_fname}; \n",
    "        rm *.zip; \n",
    "        cd studies; \n",
    "        export TIDEFILE_LOC=$AZ_BATCH_NODE_SHARED_DIR; \n",
    "        sed -i 's+./output/DCP_EX.h5+${{TIDEFILE_LOC}}/DCP_EX.h5+g' planning_ptm.inp;\n",
    "        ptm planning_ptm.inp; \n",
    "        rm output/trace.out;\n",
    "        mkdir -p $AZ_BATCH_TASK_WORKING_DIR/{run_no};\n",
    "        mv output/* $AZ_BATCH_TASK_WORKING_DIR/{run_no};\n",
    "        \"\"\", app_pkgs,ostype='linux')\n",
    "    #print(cmd_string)\n",
    "    ptm_task = client.create_task(f'{task_prefix}_{task_name}_{run_no}', cmd_string,\n",
    "                                  resource_files=[input_file],\n",
    "                                  output_files=[output_dir],\n",
    "                                  env_settings=envvars)\n",
    "    return ptm_task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create all tasks\n",
    "This function looks at the insertion location file and the simulation years and months to create an array of tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import logging\n",
    "import datetime\n",
    "import os\n",
    "def create_tasks(study_prefix, \n",
    "               insertion_file='run_number_loc.txt',\n",
    "               simulation_start_year=1923,\n",
    "               simulation_end_year=2015,\n",
    "               simulation_start_day=1,\n",
    "               simulation_month=[1, 2, 3, 4, 5, 6],\n",
    "               simulation_days=92,\n",
    "               duration='1485minutes',\n",
    "               delay='0day'):\n",
    "    tasks = []\n",
    "    study_folder, study_name = study_prefix.split('/')\n",
    "    study_name,_ = study_name.split('_')\n",
    "    tsnow = str(datetime.datetime.now().timestamp()).split('.')[0]\n",
    "    userid = os.getlogin()\n",
    "    with open(insertion_file, 'r') as input:\n",
    "        for row in csv.DictReader(input):  # run#,particle#,node\n",
    "            run_no = row['run#']\n",
    "            particle_no = row['particle#']\n",
    "            insertion_node = row['node']\n",
    "            job_name_prefix = 'ptm-%s-%s-%s' % (\n",
    "                study_folder[0:5], study_name, run_no)\n",
    "            #\n",
    "            sim_days = datetime.timedelta(days=simulation_days)\n",
    "            for y in range(simulation_start_year, simulation_end_year+1):\n",
    "                for m in simulation_month:\n",
    "                    s_day = datetime.date(y, m, simulation_start_day)\n",
    "                    e_day = s_day + sim_days\n",
    "                    ptm_start_date = s_day.strftime(\"%d%b%Y\")\n",
    "                    ptm_end_date = e_day.strftime(\"%d%b%Y\")\n",
    "                    particle_insertion_row = '%s %s %s %s' % (\n",
    "                        insertion_node, particle_no, delay, duration)\n",
    "                    envvars = {'RUN_NO': '%s' % run_no,\n",
    "                               'PTM_START_DATE': '%s' % ptm_start_date,\n",
    "                               'PTM_END_DATE': '%s' % ptm_end_date,\n",
    "                               'PARTICLE_INSERTION_ROW': '%s' % particle_insertion_row,\n",
    "                               'DSM2_STUDY_NAME': 'DCP_%s_%sP' % (study_name, study_folder[0:1])\n",
    "                               }\n",
    "                    task = create_ptm_single_task(ptm_start_date, run_no, f'{userid}_{tsnow}', f'{study_prefix}', envvars)\n",
    "                    tasks.append(task)\n",
    "    logging.info('All done!')\n",
    "    return tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "study_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = create_tasks(study_prefix, insertion_file='d:/dev/ptm_batch/run_number_loc.txt',simulation_start_year=1923,simulation_end_year=2015)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next submit the task and wait \n",
    "Azure batch limits to submitting 100 tasks at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,round(len(tasks)/100)):\n",
    "    client.submit_tasks(job_name,tasks[i*100:i*100+100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finally resize the pool to 0 to save costs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#client.resize_pool(pool_name,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "9e539dc02fc56d55a22d79a0646788fa38cb6ebb3e85aa69306aae5c2f643a8f"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}