{-# LANGUAGE TypeApplications #-}

module Hyperion.Slurm.Environment where

import           Control.Applicative       ((<|>))
import           Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import           Data.Maybe                (fromMaybe)
import           System.Environment        (lookupEnv)
import           System.Process            (readCreateProcess, shell)
import           Text.Read                 (readMaybe)

-- | Returns number of tasks per node by reading system environment variables.
-- If @SLURM_NTASKS_PER_NODE@ is defined, returns it. Otherwise, tries to compute
-- from @SLURM_NTASKS@ and @SLURM_JOB_NUM_NODES@. If this doens't work either,
-- fails with 'error'.
getNTasksPerNode :: IO Int
getNTasksPerNode :: IO Int
getNTasksPerNode =
  Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Could not determine NTASKS_PER_NODE") (Maybe Int -> Int) -> IO (Maybe Int) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  MaybeT IO Int -> IO (Maybe Int)
forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT (MaybeT IO Int
lookupNTasks MaybeT IO Int -> MaybeT IO Int -> MaybeT IO Int
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> MaybeT IO Int
computeNTasks)
  where
    lookupInt :: String -> MaybeT IO Int
    lookupInt :: [Char] -> MaybeT IO Int
lookupInt [Char]
name = IO (Maybe Int) -> MaybeT IO Int
forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT (IO (Maybe Int) -> MaybeT IO Int)
-> IO (Maybe Int) -> MaybeT IO Int
forall a b. (a -> b) -> a -> b
$ do
      Maybe [Char]
mStr <- [Char] -> IO (Maybe [Char])
lookupEnv [Char]
name
      Maybe Int -> IO (Maybe Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Int -> IO (Maybe Int)) -> Maybe Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ Maybe [Char]
mStr Maybe [Char] -> ([Char] -> Maybe Int) -> Maybe Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Read Int => [Char] -> Maybe Int
forall a. Read a => [Char] -> Maybe a
readMaybe @Int
    lookupNTasks :: MaybeT IO Int
lookupNTasks = [Char] -> MaybeT IO Int
lookupInt [Char]
"SLURM_NTASKS_PER_NODE"
    computeNTasks :: MaybeT IO Int
computeNTasks = do
      Int
nTasks <- [Char] -> MaybeT IO Int
lookupInt [Char]
"SLURM_NTASKS"
      Int
nNodes <- [Char] -> MaybeT IO Int
lookupInt [Char]
"SLURM_JOB_NUM_NODES"
      Int -> MaybeT IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
nTasks Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
nNodes)

-- | Returns the contents of @SLURM_JOB_NODELIST@ as a list of nodes names
getJobNodes :: IO [String]
getJobNodes :: IO [[Char]]
getJobNodes = ([Char] -> [[Char]]) -> IO [Char] -> IO [[Char]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Char] -> [[Char]]
lines (IO [Char] -> IO [[Char]]) -> IO [Char] -> IO [[Char]]
forall a b. (a -> b) -> a -> b
$
  CreateProcess -> [Char] -> IO [Char]
readCreateProcess ([Char] -> CreateProcess
shell [Char]
"scontrol show hostnames $SLURM_JOB_NODELIST") [Char]
""

-- | Returns the value of @SLURMD_NODENAME@
lookupHeadNode :: IO (Maybe String)
lookupHeadNode :: IO (Maybe [Char])
lookupHeadNode = [Char] -> IO (Maybe [Char])
lookupEnv [Char]
"SLURMD_NODENAME"