workflow.yaml

permissions:
  - '*'
sessions:
  dashboard:
    useTLS: false
    redirect: true

jobs:
  install_dask:
    steps:
      - name: Check Environment
        if: ${{ inputs.install_python_environment}}
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
          conda_env: ${{ inputs.conda_env }}
        run: |
          set -x
          conda_parent_dir=$(dirname ${{ inputs.conda_dir }})
          # Check if marker files exist
          ${sshcmd}  "[ -f ${conda_parent_dir}/.miniconda_installed ] && echo install_miniconda=false || echo install_miniconda=true"  >> $OUTPUTS
          ${sshcmd}  "[ -f ${conda_parent_dir}/.${conda_env}_installed ] && echo install_dask=false || echo install_dask=true"  >> $OUTPUTS
      - name: Install Miniconda
        if: ${{ inputs.install_python_environment }}
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
          install_miniconda: ${{ needs.install_dask.outputs.install_miniconda }}
          conda_dir: ${{ inputs.conda_dir }})
        run: |
          set -x
          if ! ${install_miniconda}; then
            exit 0
          fi
          ID=$(date +%s)-${RANDOM}
          conda_parent_dir=$(dirname ${conda_dir})
          ${sshcmd} << EOF
          wget --no-check-certificate https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda-${ID}.sh
          mkdir -p ${conda_parent_dir}
          bash /tmp/miniconda-${ID}.sh -b -p ${{ inputs.conda_dir }}
          touch ${conda_parent_dir}/.miniconda_installed
          EOF
      - name: Install Dask
        if: ${{ inputs.install_python_environment }}
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
          conda_env: ${{ inputs.conda_env }}
          conda_dir: ${{ inputs.conda_dir }}
          install_dask: ${{ needs.install_dask.outputs.install_dask }}
        run: |
          set -x
          if ! ${install_dask}; then
            exit 0
          fi
          conda_parent_dir=$(dirname ${conda_dir})
          ${sshcmd} << EOF
          source ${conda_dir}/etc/profile.d/conda.sh
          conda create --name ${conda_env} -y
          conda activate ${conda_env}
          conda install dask distributed dask-jobqueue -c conda-forge -y
          touch ${conda_parent_dir}/.${conda_env}_installed
          EOF
  prepare_work_directory:
    steps:
      - name: Transfer Data
        run: |
          # Create remote job directory
          ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }} mkdir -p ~/${PWD/#$HOME/}
          # Transfer dask script
          scp dask_script.py  ${{ inputs.resource.ip }}:~/${PWD/#$HOME/}
          # Transfer inputs.json
          scp inputs.json  ${{ inputs.resource.ip }}:~/${PWD/#$HOME/}
  create_session:
    needs:
      - install_dask
      - prepare_work_directory
    steps:
      - name: Reserve dashboard ports
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
        run: |
          local_port=$(pw agent open-port)
          echo "export local_port=${local_port}" >> ports
          echo "local_port=${local_port}" >> $OUTPUTS

          dashboard_port=$(${sshcmd} pw agent open-port)
          echo "export dashboard_port=${dashboard_port}" >> ports
          echo "dashboard_port=${dashboard_port}" >> $OUTPUTS

          nginx_proxy_port=$(${sshcmd} pw agent open-port)
          echo "export nginx_proxy_port=${nginx_proxy_port}" >> ports
          echo "nginx_proxy_port=${nginx_proxy_port}" >> $OUTPUTS
          
          cat ports
      - name: Get Controller Hostname
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
        run: |
          target_hostname=$(${sshcmd} hostname)
          echo "target_hostname=${target_hostname}"  | tee -a $OUTPUTS
      - name: Create session
        uses: parallelworks/update-session
        with:
          remotePort: "${{ needs.create_session.outputs.nginx_proxy_port }}"
          localPort: "${{ needs.create_session.outputs.local_port }}"
          remoteHost: "${{ needs.create_session.outputs.target_hostname }}"
          status: running
          slug: "status"
          target: ${{ inputs.resource.id }}
          name: ${{ sessions.dashboard }}


  run_dask:
    needs:
      - create_session
    steps:
      - name: Run Dask
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
          load_env: ${{ inputs.load_env }}
          conda_env: ${{ inputs.conda_env }}
          conda_dir: ${{ inputs.conda_dir }}
        run: |
          source ports
          set -x
          ${sshcmd} << EOF
          set -x
          if [ -z "${load_env}" ]; then
            source ${conda_dir}/etc/profile.d/conda.sh
            conda activate ${conda_env}
          else
            eval ${load_env}
          fi
          export dashboard_port_worker=${dashboard_port}
          export http_prefix=/me/session/${PW_USER}/${{ sessions.dashboard }}/
          cd ~/${PWD/#$HOME/}
          python -u dask_script.py &
          dask_script_pid=\$!
          echo \${dask_script_pid} > dask_script.pid
          wait \${dask_script_pid}
          rm dask_script.pid
          EOF
        
        cleanup: |
          sshcmd="ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}"
          container_name=$(echo ${PWD} | rev | cut -d'/' -f1-2 | rev | tr '/' '-')
          ${sshcmd} << EOF
          cd ~/${PWD/#$HOME/}
          sudo docker stop ${container_name}
          # Check if the file dask_script.pid exists
          if [ ! -f dask_script.pid ]; then
            exit 0
          fi
          dask_script_pid=\$(cat dask_script.pid)
          echo "Killing dask script PID \${dask_script_pid}"
          kill \${dask_script_pid}
          scancel \$(ls slurm-*.out | sed 's/slurm-\([0-9]\+\).out/\1/')
          EOF



  start_nginx_proxy:
    needs:
      - create_session
    steps:
      - name: Write Nginx config
        run: |
          source ports
          # Write config file
          cat >> config.conf <<HERE
          server {
            listen ${nginx_proxy_port};
            server_name _;

            index index.html index.htm index.php;

            # Allow CORS
            add_header 'Access-Control-Allow-Methods' 'GET, POST, OPTIONS';
            add_header 'Access-Control-Allow-Headers' 'Authorization, Content-Type, Accept, Origin, User-Agent, DNT, Cache-Control, X-Mx-ReqToken, Keep-Alive, X-Requested-With, If-Modified-Since';
            add_header X-Frame-Options "ALLOWALL";

            client_max_body_size 1000M;

            location / {
              proxy_pass http://127.0.0.1:${dashboard_port}/me/session/${PW_USER}/${{ sessions.dashboard }}/;  # Ensure this has a trailing slash
              proxy_http_version 1.1;
              proxy_set_header Upgrade \$http_upgrade;
              proxy_set_header Connection "upgrade";
              proxy_set_header X-Real-IP \$remote_addr;
              proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
              proxy_set_header Host \$http_host;
              proxy_set_header X-NginX-Proxy true;

              # Optionally, handle WebSocket connections
              proxy_set_header X-Forwarded-Proto \$scheme;
              proxy_set_header X-Forwarded-Host \$host;
              proxy_set_header X-Forwarded-Port \$server_port;
            }
          }

          HERE
          scp config.conf  ${{ inputs.resource.ip }}:~/${PWD/#$HOME/}
      - name: Start Nginx Docker Container
        env:
          sshcmd: ssh -o StrictHostKeyChecking=no ${{ inputs.resource.ip }}
        run: |
          set -x
          container_name=$(echo ${PWD} | rev | cut -d'/' -f1-2 | rev | tr '/' '-')
          ${sshcmd} << EOF
          cd ~/${PWD/#$HOME/}
          sudo service docker start
          touch empty
          touch nginx.logs
          # change ownership to nginx user
          sudo chown 101:101 nginx.logs 
          sudo docker run -d --rm --name ${container_name} -v $PWD/config.conf:/etc/nginx/conf.d/config.conf  -v $PWD/empty:/etc/nginx/conf.d/default.conf -v $PWD/nginx.logs:/var/log/nginx/access.log  -v $PWD/nginx.logs:/var/log/nginx/error.log --network=host nginxinc/nginx-unprivileged:1.25.3
          EOF
'on':
  execute:
    inputs:
      header:
        type: header
        text: Dask SLURM Demo
        size: 20
      resource:
        type: compute-clusters
        label: Select SLURM Cluster
        include-workspace: false
        provider:
          - gclusterv2
          - pclusterv2
          - azclusterv2
          - aws-slurm
          - google-slurm
          - azure-slurm
      partition:
        type: slurm-partitions
        label: SLURM partition
        resource: ${{ inputs.resource }}
      cores_per_job:
        label: Cores per job
        type: number
        min: 1
        max: 10
        default: 2
        tooltip: cores parameter in SLURMCluster object
      memory_per_job:
        label: Memory per job
        type: string
        default: 8GB
        tooltip: memory parameter in SLURMCluster object
      minimum_jobs:
        label: Minimum Jobs
        type: number
        min: 0
        max: 10
        default: 0
        tooltip: minimum parameter of the SLURMCluster.adapt method
      maximum_jobs:
        label: Maximum Jobs
        type: number
        min: 0
        max: 10
        default: 10
        tooltip: maximum parameter of the SLURMCluster.adapt method
      install_python_environment:
        label: Install Python Environment
        type: boolean
        default: true
      conda_dir:
        label: Path to install conda
        type: string
        default: ~/pw/software/miniconda3
        hidden: ${{ inputs.install_python_environment == false }}
        optional: ${{ .hidden }}
        ignore: ${{ .hidden }}
      conda_env:
        label: Conda environment
        type: string
        default: dask-slurm
        hidden: ${{ inputs.install_python_environment == false }}
        optional: ${{ .hidden }}
        ignore: ${{ .hidden }}
      load_env:
        label: Command to load the Python environment
        type: string
        hidden: ${{ inputs.install_python_environment == true }}
        optional: ${{ .hidden }}
        ignore: ${{ .hidden }}