MODULE job

!-----------------------------------------------------------------------
! NAME
!     job
!
! DESCRIPTION
!     Tools to retieve job information of the program.
!
! AUTHORS & DATE
!     JB Clement, 10/06/2024
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! PARAMETERS
! ----------
real            :: timelimit = 86400. ! Time limit for the job: 86400 s = 24 h by default
real, parameter :: antetime  = 3600.  ! Anticipation time to prevent reaching the job time limit: 3600 s = 1 h by default
logical         :: timewall = .false. ! Flag to enable the program self-termination before timeout

contains
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

!=======================================================================
SUBROUTINE get_job_id(jobid)
!-----------------------------------------------------------------------
! NAME
!     get_job_id
!
! DESCRIPTION
!     Retrieves the ID for a given job ID via SLURM or PBS/TORQUE.
!
! AUTHORS & DATE
!     JB Clement, 10/06/2024
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
character(256), intent(out) :: jobid

! LOCAL VARIABLES
! ---------------
integer        :: sts
character(256) :: tmp

! CODE
! ----
! Try SLURM
call get_environment_variable("SLURM_JOBID",tmp,status = sts)
if (sts == 0 .and. len_trim(tmp) > 0) then
    jobid = trim(adjustl(tmp))
    return
end if

! Try PBS/TORQUE
call get_environment_variable("PBS_JOBID",tmp,status = sts)
if (sts == 0 .and. len_trim(tmp) > 0) then
    ! Extract only the part before the point in '12345.master.cluster' to get the numeric job ID
    jobid = trim(adjustl(tmp(1:index(tmp,'.') - 1)))
    return
end if

error stop 'Error: neither SLURM_JOB_ID nor PBS_JOBID found in environment!'

END SUBROUTINE get_job_id
!=======================================================================

!=======================================================================
SUBROUTINE get_job_timelimit(jobid)
!-----------------------------------------------------------------------
! NAME
!     get_job_timelimit
!
! DESCRIPTION
!     Retrieves the time limit (in seconds) for a given job ID via SLURM
!     or PBS/TORQUE.
!
! AUTHORS & DATE
!     JB Clement, 10/06/2024
!
! NOTES
!
!-----------------------------------------------------------------------

! DECLARATION
! -----------
implicit none

! ARGUMENTS
! ---------
character(*), intent(in) :: jobid

! LOCAL VARIABLES
! ---------------
logical        :: num_str
character(256) :: chtimelimit ! Time limit for the job outputted by the SLURM or PBS/TORQUE command
integer        :: i, ierr, funit, cstat, days, hours, minutes, seconds
character(1)   :: sep

! CODE
! ----
! Check that the job ID is numeric
num_str = .true.
do i = 1,len_trim(jobid)
    if (jobid(i:i) < '0' .or. jobid(i:i) > '9') then
        num_str = .false.
        exit
    end if
end do
if (.not. num_str) error stop "Error: job ID must be numeric!"

! Try SLURM (squeue)
call execute_command_line('squeue -j '//trim(adjustl(jobid))//' -h --Format TimeLimit > tmp_timelimit.txt',cmdstat = cstat)
if (cstat /= 0) then
    ! On failure, try PBS/TORQUE (qstat)
    call execute_command_line('qstat -f '//trim(adjustl(jobid))//' | grep "Walltime" | awk ''{print $3}'' > tmp_timelimit.txt',cmdstat = cstat)
    if (cstat > 0) then
        error stop 'Error: command execution failed!'
    else if (cstat < 0) then
        error stop 'Error: command execution not supported (neither SLURM nor PBS/TORQUE is installed)!'
    end if
end if

! Read the output
open(newunit = funit,file = 'tmp_timelimit.txt',status = 'old',form = 'formatted',action = 'read',iostat = ierr)
if (ierr /= 0) error stop 'Error: error opening file "tmp_timelimit.txt"!'
read(funit,'(a)') chtimelimit
close(funit)
chtimelimit = trim(adjustl(chtimelimit))

! Remove temporary file
call execute_command_line('\rm -f tmp_timelimit.txt',cmdstat = cstat)
if (cstat > 0) then
    error stop 'Error: command execution failed!'
else if (cstat < 0) then
    error stop 'Error: command execution not supported!'
end if

! Parse D-HH:MM:SS, HH:MM:SS or MM:SS
if (index(chtimelimit,'-') > 0) then ! Format "D-HH:MM:SS"
    read(chtimelimit,'(i1,a1,i2,a1,i2,a1,i2)') days, sep, hours, sep, minutes, sep, seconds
    timelimit = real(days)*86400. + real(hours)*3600. + real(minutes)*60. + real(seconds)
else if (index(chtimelimit,':') > 0 .and. len_trim(chtimelimit) > 5) then ! Format "HH:MM:SS"
    read(chtimelimit,'(i2,a1,i2,a1,i2)') hours, sep, minutes, sep, seconds
    timelimit = real(hours)*3600. + real(minutes)*60. + real(seconds)
else ! Format "MM:SS"
    read(chtimelimit,'(i2,a1,i2)') minutes, sep, seconds
    timelimit = real(minutes)*60. + real(seconds)
end if

timewall = .true.

END SUBROUTINE get_job_timelimit
!=======================================================================

END MODULE job