packages feed

drmaa-0.2.0: src/DRMAA.hs

{-# LANGUAGE QuasiQuotes     #-}
{-# LANGUAGE TemplateHaskell #-}
module DRMAA
    ( withSession
    , initSession
    , exitSession
    , DrmaaAttribute(..)
    , defaultDrmaaConfig
    , drmaaRun
    ) where

import           Control.Exception     (bracket_)
import           Control.Monad         (when)
import           Foreign.C.String
import           Foreign.Marshal.Alloc
import           Foreign.Marshal.Array
import           Foreign.Ptr
import qualified Language.C.Inline     as C
import           System.Directory      (getCurrentDirectory)
import           Text.Printf           (printf)

C.include "stddef.h"
C.include "stdio.h"
C.include "drmaa.h"

withSession :: IO a -> IO a
withSession = bracket_ initSession exitSession

{-
drmaaScript :: String -> DrmaaAttribute -> IO ()
drmaaScript script config = bracket
    (shelly $ fmap (T.unpack . head . T.lines) $ silently $ run "mktemp" [template])
    (shelly . rm . fromText . T.pack)
    $ \tmpFl -> do
        writeFile tmpFl $ "#!/bin/sh\n" ++ script
        shelly $ run_ "chmod" ["+x", T.pack tmpFl]
        drmaaRun tmpFl [] config
  where
    template = T.pack $ drmaa_wd config ++ "/drmaa_script.XXXXXXXX.delete.me.sh"
    -}

-- | Initialize a session
initSession :: IO ()
initSession = alloca $ \ptr -> do
    status <- [C.block| int {
        int errnum = 0;
        errnum = drmaa_init (NULL, $(char* ptr), DRMAA_ERROR_STRING_BUFFER);
        if (errnum != DRMAA_ERRNO_SUCCESS) {
            return 1;
        }
        return 0;
        }|]
    when (status /= 0) $ peekCString ptr >>= error
{-# INLINE initSession #-}

exitSession :: IO ()
exitSession = do
    r <- [C.block| int {
        char error[DRMAA_ERROR_STRING_BUFFER];
        int errnum = 0;
        errnum = drmaa_exit (error, DRMAA_ERROR_STRING_BUFFER);
        if (errnum != DRMAA_ERRNO_SUCCESS) {
            fprintf (stderr, "Could not shut down the DRMAA library: %s\n", error);
            return 1;
        }
        return 0;
        }|]
    when (r /= 0) $ error "Exit 1"
{-# INLINE exitSession #-}

data DrmaaAttribute = DrmaaAttribute
    { drmaa_wd     :: !(Maybe FilePath)
    , drmaa_env    :: ![(String, String)]
    , drmaa_native :: !String
    } deriving (Show, Read)

defaultDrmaaConfig :: DrmaaAttribute
defaultDrmaaConfig = DrmaaAttribute
    { drmaa_wd = Nothing
    , drmaa_env = [ ("DRMAA_ENV_HAS_SET", "True") ]
    , drmaa_native = ""
    }

drmaaRun :: FilePath -> [String] -> DrmaaAttribute -> IO ()
drmaaRun exec args config = do
    c_exec <- newCString exec
    c_args <- mapM newCString args

    -- options
    wd <- get_wd (drmaa_wd config) >>= newCString
    env <- mapM (\(a,b) -> newCString $ a ++ "=" ++ b) $ drmaa_env config
    native <- newCString $ drmaa_native config

    e <- withArray (c_args++[nullPtr]) $ \aptr ->
        withArray (env ++ [nullPtr]) $ \envPtr -> do
            [C.block| int {
            int exception = 0;
            char error[DRMAA_ERROR_STRING_BUFFER];
            int errnum = 0;
            drmaa_job_template_t *jt = NULL;

            errnum = drmaa_allocate_job_template (&jt, error, DRMAA_ERROR_STRING_BUFFER);

            if (errnum != DRMAA_ERRNO_SUCCESS) {
                fprintf (stderr, "Could not create job template: %s\n", error);
            } else {
                /* set options */
                errnum = drmaa_set_attribute (jt, DRMAA_WD, $(char* wd),
                                             error, DRMAA_ERROR_STRING_BUFFER);
                errnum = drmaa_set_vector_attribute (jt, DRMAA_V_ENV, $(const char** envPtr),
                                             error, DRMAA_ERROR_STRING_BUFFER);
                errnum = drmaa_set_attribute (jt, DRMAA_NATIVE_SPECIFICATION, $(char* native),
                                             error, DRMAA_ERROR_STRING_BUFFER);

                errnum = drmaa_set_attribute (jt, DRMAA_REMOTE_COMMAND, $(char* c_exec),
                                             error, DRMAA_ERROR_STRING_BUFFER);
                if (errnum != DRMAA_ERRNO_SUCCESS) {
                    fprintf (stderr, "Could not set attribute \"%s\": %s\n",
                            DRMAA_REMOTE_COMMAND, error);
                } else {
                    errnum = drmaa_set_vector_attribute (jt, DRMAA_V_ARGV, $(const char** aptr), error,
                                                        DRMAA_ERROR_STRING_BUFFER);
                }

                if (errnum != DRMAA_ERRNO_SUCCESS) {
                    fprintf (stderr, "Could not set attribute \"%s\": %s\n",
                            DRMAA_REMOTE_COMMAND, error);
                } else {
                    char jobid[DRMAA_JOBNAME_BUFFER];
                    char jobid_out[DRMAA_JOBNAME_BUFFER];
                    int status = 0;
                    drmaa_attr_values_t *rusage = NULL;

                    errnum = drmaa_run_job (jobid, DRMAA_JOBNAME_BUFFER, jt, error,
                                           DRMAA_ERROR_STRING_BUFFER);

                    if (errnum != DRMAA_ERRNO_SUCCESS) {
                        fprintf (stderr, "Could not submit job: %s\n", error);
                        exception = 1;
                    } else {
                        printf ("Your job has been submitted with id %s\n", jobid);

                        errnum = drmaa_wait (jobid, jobid_out, DRMAA_JOBNAME_BUFFER, &status,
                                            DRMAA_TIMEOUT_WAIT_FOREVER, &rusage, error,
                                            DRMAA_ERROR_STRING_BUFFER);

                        if (errnum != DRMAA_ERRNO_SUCCESS) {
                            fprintf (stderr, "Could not wait for job: %s\n", error);
                            exception = 1;
                        } else {
                            char usage[DRMAA_ERROR_STRING_BUFFER];
                            int aborted = 0;

                            drmaa_wifaborted(&aborted, status, NULL, 0);

                            if (aborted == 1) {
                                printf("Job %s never ran\n", jobid);
                                exception = 1;
                            } else {
                                int exited = 0;

                                drmaa_wifexited(&exited, status, NULL, 0);

                                if (exited == 1) {
                                    int exit_status = 0;

                                    drmaa_wexitstatus(&exit_status, status, NULL, 0);
                                    printf("Job %s finished regularly with exit status %d\n", jobid, exit_status);
                                    exception = exit_status;
                                } else {
                                    int signaled = 0;

                                    drmaa_wifsignaled(&signaled, status, NULL, 0);

                                    if (signaled == 1) {
                                        char termsig[DRMAA_SIGNAL_BUFFER+1];

                                        drmaa_wtermsig(termsig, DRMAA_SIGNAL_BUFFER, status, NULL, 0);
                                        printf("Job %s finished due to signal %s\n", jobid, termsig);
                                    } else {
                                        printf("Job %s finished with unclear conditions\n", jobid);
                                    }
                                } /* else */
                            } /* else */

                            /*
                            printf ("Job Usage:\n");

                            while (drmaa_get_next_attr_value (rusage, usage, DRMAA_ERROR_STRING_BUFFER) == DRMAA_ERRNO_SUCCESS) {
                                printf ("  %s\n", usage);
                            }
                            */

                            drmaa_release_attr_values (rusage);
                        } /* else */
                    } /* else */
                } /* else */

                errnum = drmaa_delete_job_template (jt, error, DRMAA_ERROR_STRING_BUFFER);

                if (errnum != DRMAA_ERRNO_SUCCESS) {
                    fprintf (stderr, "Could not delete job template: %s\n", error);
                }
            } /* else */

            return exception;
            }|]
    when (e/=0) $ error $ printf
        "Job failed (status=%d). Please see SGE log for details"
        (fromIntegral e :: Int)

get_wd :: Maybe FilePath -> IO FilePath
get_wd Nothing  = getCurrentDirectory
get_wd (Just x) = return x
{-# INLINE get_wd #-}