04-create-dag-file-via-python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e619525e-be57-45a5-8af8-1790ebf2250b",
   "metadata": {},
   "source": [
    "# Creating HTCondor's DAG file via Python script for a single task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a72e58ee-e50e-4942-9451-37da35c25600",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "- Preparing Python script: `/mnt/nfs/scripts/print_hello.py` to run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b8d4467f-1d8a-427c-bc2c-bb8d7305d552",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Creating a new directory, if not exists\n",
    "os.makedirs(\"/mnt/nfs/scripts\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b08594a5-e4d0-4fe1-b8ed-82c2df6fd451",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /mnt/nfs/scripts/task1.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile /mnt/nfs/scripts/task1.py\n",
    "#!/usr/bin/env python3\n",
    "# task1.py\n",
    "import time\n",
    "\n",
    "print(\"Task 1 started\")\n",
    "time.sleep(3)  # Simulate a task that takes 2 seconds\n",
    "print(\"Task 1 completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "193956b8-7ce3-43e9-89fe-3326bec4c306",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /mnt/nfs/scripts/task2.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile /mnt/nfs/scripts/task2.py\n",
    "#!/usr/bin/env python3\n",
    "# task2.py\n",
    "import time\n",
    "\n",
    "print(\"Task 2 started\")\n",
    "time.sleep(5)  # Simulate a task that takes 2 seconds\n",
    "print(\"Task 2 completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3c148abc-f72a-4c47-be94-d8cb6910d683",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /mnt/nfs/scripts/task3.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile /mnt/nfs/scripts/task3.py\n",
    "#!/usr/bin/env python3\n",
    "# task3.py\n",
    "import time\n",
    "\n",
    "print(\"Task 3 started\")\n",
    "time.sleep(5)  # Simulate a task that takes 2 seconds\n",
    "print(\"Task 3 completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f8b0ff2e-aa21-462f-86d1-8cfb3bb83270",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /mnt/nfs/scripts/task4.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile /mnt/nfs/scripts/task4.py\n",
    "#!/usr/bin/env python3\n",
    "# task4.py\n",
    "import time\n",
    "\n",
    "print(\"Task 4 started\")\n",
    "time.sleep(3)  # Simulate a task that takes 2 seconds\n",
    "print(\"Task 4 completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76713ed1-6b8e-4dbb-94cf-f5edc15c4d31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# make sure the python script is executable\n",
    "!chmod 764 /mnt/nfs/scripts/task1.py\n",
    "!chmod 764 /mnt/nfs/scripts/task2.py\n",
    "!chmod 764 /mnt/nfs/scripts/task3.py\n",
    "!chmod 764 /mnt/nfs/scripts/task4.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8084768-dd05-4351-8cf3-02aeccc405d9",
   "metadata": {},
   "source": [
    "- Creating DAG file with Python Script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff3101b1-3c37-4ee6-a423-d5aaa74bcbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import htcondor\n",
    "import os\n",
    "from htcondor import dags\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "base_dir = \"/mnt/nfs\"\n",
    "\n",
    "# create the submit script for task 1\n",
    "sub1 = htcondor.Submit({\n",
    "    \"executable\": os.path.join(base_dir, \"scripts/task1.py\"),\n",
    "    \"request_cpus\": \"1\",\n",
    "    \"request_memory\": \"128MB\",\n",
    "    \"request_disk\": \"128MB\",\n",
    "    \"output\": os.path.join(base_dir, \"output/task1.out\"),\n",
    "    \"error\": os.path.join(base_dir, \"error/task1.err\"),\n",
    "    \"log\": os.path.join(base_dir, \"log/task1.log\"),\n",
    "})\n",
    "\n",
    "## Task 2 and Task 3 will be running concurrently\n",
    "# create the submit script for task 2\n",
    "sub2 = htcondor.Submit({\n",
    "    \"executable\": os.path.join(base_dir, \"scripts/task2.py\"),\n",
    "    \"request_cpus\": \"1\",\n",
    "    \"request_memory\": \"128MB\",\n",
    "    \"request_disk\": \"128MB\",\n",
    "    \"output\": os.path.join(base_dir, \"output/task2.out\"),\n",
    "    \"error\": os.path.join(base_dir, \"error/task2.err\"),\n",
    "    \"log\": os.path.join(base_dir, \"log/task2.log\"),\n",
    "})\n",
    "\n",
    "# create the submit script for task 3\n",
    "sub3 = htcondor.Submit({\n",
    "    \"executable\": os.path.join(base_dir, \"scripts/task3.py\"),\n",
    "    \"request_cpus\": \"1\",\n",
    "    \"request_memory\": \"128MB\",\n",
    "    \"request_disk\": \"128MB\",\n",
    "    \"output\": os.path.join(base_dir, \"output/task3.out\"),\n",
    "    \"error\": os.path.join(base_dir, \"error/task3.err\"),\n",
    "    \"log\": os.path.join(base_dir, \"log/task3.log\"),\n",
    "})\n",
    "\n",
    "# create the submit script for task 4\n",
    "sub4 = htcondor.Submit({\n",
    "    \"executable\": os.path.join(base_dir, \"scripts/task4.py\"),\n",
    "    \"request_cpus\": \"1\",\n",
    "    \"request_memory\": \"128MB\",\n",
    "    \"request_disk\": \"128MB\",\n",
    "    \"output\": os.path.join(base_dir, \"output/task4.out\"),\n",
    "    \"error\": os.path.join(base_dir, \"error/task4.err\"),\n",
    "    \"log\": os.path.join(base_dir, \"log/task4.log\"),\n",
    "})\n",
    "\n",
    "\n",
    "dag = dags.DAG()\n",
    "\n",
    "# Define job layers for each task\n",
    "task1_layer = dag.layer(\n",
    "    name='Task1',\n",
    "    submit_description=sub1\n",
    ")\n",
    "\n",
    "# Task 2 and Task 3 are children of Task 1, meaning they will run concurrently after Task 1 completes\n",
    "task2_layer = task1_layer.child_layer(\n",
    "    name='Task2',\n",
    "    submit_description=sub2\n",
    ")\n",
    "\n",
    "task3_layer = task1_layer.child_layer(\n",
    "    name='Task3',\n",
    "    submit_description=sub3\n",
    ")\n",
    "\n",
    "# Task 4 will run after both Task 2 and Task 3 complete\n",
    "task4_layer = dag.layer(\n",
    "    name='Task4',\n",
    "    submit_description=sub4\n",
    ")\n",
    "task4_layer.add_parents([task2_layer, task3_layer])\n",
    "\n",
    "# Set up the DAG directory\n",
    "# Write the DAG to disk\n",
    "dag_dir = os.path.abspath(\"/mnt/nfs/demo_dags/\")\n",
    "os.makedirs(dag_dir, exist_ok=True)\n",
    "dag_file = dags.write_dag(dag, dag_dir)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f01a8cf-e275-4989-ad24-4c9dd0ebe36f",
   "metadata": {},
   "source": [
    "- Check the DAG file created"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ad2608-575e-48a2-9cdb-58540df4e957",
   "metadata": {},
   "source": [
    "Based on these dependencies, the following execution order is implied:\n",
    "\n",
    "- Task1 executes\n",
    "- Task3 and Task2 execute concurrently (since they both depend on Task1)\n",
    "- Task4 executes only after both Task3 and Task2 have completed successfully"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "82a4ab45-5e8a-48e3-8e2a-d9a305dff378",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# BEGIN META\n",
      "# END META\n",
      "# BEGIN NODES AND EDGES\n",
      "JOB Task1:0 Task1.sub\n",
      "PARENT Task1:0 CHILD Task3:0\n",
      "PARENT Task1:0 CHILD Task2:0\n",
      "JOB Task3:0 Task3.sub\n",
      "PARENT Task3:0 CHILD Task4:0\n",
      "JOB Task2:0 Task2.sub\n",
      "PARENT Task2:0 CHILD Task4:0\n",
      "JOB Task4:0 Task4.sub\n",
      "# END NODES AND EDGES\n",
      "\n"
     ]
    }
   ],
   "source": [
    "with open(dag_file, \"r\") as file:\n",
    "    dagfile_content = file.read()\n",
    "    print(dagfile_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17569f6b-0d2b-473b-adcc-826c6f603c43",
   "metadata": {},
   "source": [
    "- Go to terminal, then type\n",
    "\n",
    "> cd /mnt/nfs/demo_dags\n",
    "\n",
    "> condor_submit_dag dagfile.dag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e4094d-b010-442d-879b-b87ae4f23a1d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}