Lmod SitePackage.lua

SitePackage.lua

-- -*- lua -*-
-- vim:ft=lua:et:ts=4
require("sandbox")
require("strict")

local Dbg = require("Dbg")
local dbg = Dbg:dbg()
local unpack = unpack or table.unpack
local hook = require("Hook")

sitePkgRoot = os.getenv("LMOD_PKG_ROOT") or "/apps"
SiteRootDir = os.getenv("LMOD_SITE_ROOT_DIR") or "/apps"
DefaultsDir = os.getenv("LMOD_DEFAULTS_DIR") or "/apps/modulefiles/Defaults"

--[[

loadPkgDefaults:

Pulled from Lmod's contrib directory, with added argument "hierType".

By default hierType is "A", but can be set to "B" to use the hierarchyB function.

]]
function loadPkgDefaults(levels, hierType)
    local pkg = {}
    local hier
    local status
    local msg
    local whole

    hierType = hierType or "A"

    ------------------------------------------------------------
    -- Fill default values
    pkg.name         = myModuleName()
    pkg.version      = myModuleVersion()
    pkg.id           = myModuleFullName()
    pkg.display_name = pkg.name
    pkg.url          = ""
    pkg.license      = ""
    pkg.category     = ""
    pkg.keywords     = ""
    pkg.description  = ""
    pkg.help         = ""

    ------------------------------------------------------------
    -- build package prefix from modulefile location
    if (hierType == "A") then
        hier  = hierarchyA(pkg.id, levels)
    else
        hier  = hierarchyB(pkg.id, levels) 
    end

    local a          = {}
    a[#a+1]          = SiteRootDir
    for i = levels,1,-1 do
       a[#a+1] = hier[i]:gsub("/","-")
    end
    a[#a+1]          = pkg.id

    if (hierType == "B") then
        local b = {}
        local n = #a

        -- split last element of 'a' that contains sub-version
        for dir in a[n]:split("/") do
          b[#b + 1] = dir
        end

        -- save the actual version
        local b_last = b[n]
        -- save the sub-version
        local b_second_last = b[n-1]

        -- swap so foo/bar/1.0 becomes foo/1.0/bar
        b[n] = b_second_last
        b[n-1] = b_last

        -- replace last element with new swapped path
        a[n] = pathJoin(unpack(b))
    end

    pkg.prefix       = pathJoin(unpack(a))

    ------------------------------------------------------------
    -- determine compiler and mpi depedencies
    pkg.compiler     = hier[1]
    if (hier[2]) then
        if (hier[2] ~= "MPI")   then pkg.mpi = hier[2]  end
        if (hier[2] ~= "Lang")  then pkg.lang = hier[2] end
    end

    ------------------------------------------------------------
    -- Read default package description file

    local fn         = pathJoin(DefaultsDir, pkg.name .. ".lua")

    if (not isFile(fn)) then
      return pkg
    end

    local f          = io.open(fn)
    local whole      = false
    local status     = false
    local msg        = "Empty file"
    if (f) then
        whole = f:read("*all")
        f:close()
    end

    ------------------------------------------------------------
    -- Evaluate string from package description file through
    -- sandbox_run for safety checks.

    if (whole) then
        status, msg = sandbox_run(whole)
    end

    if (not status) then
        LmodError("Unable to load file: ", fn, ": ", msg, "\n")
    end

    for k,v in pairs(msg) do
       pkg[k] = v
    end
    return pkg
end

function setPkgInfo(pkg)
    help(pkg.help)
    whatis("Name: "        .. pkg.display_name)
    whatis("Version: "     .. pkg.version)
    whatis("Module: "      .. pkg.id)
    whatis("Category: "    .. pkg.category)
    whatis("Keyword: "     .. pkg.keywords)
    whatis("URL: "         .. pkg.url)
    whatis("License: "     .. pkg.license)
    whatis("Description: " .. pkg.description)
end

function checkRestrictedGroup(pkg, group)
    dbg.start{"checkRestrictedGroup(pkg, \"",group,"\")"}
    if (mode() ~= "load") then return true end
    if (group == nil)     then return true end
    local err_message = "Only users in group \'" .. group ..
        "\' can access module \'" .. pkg.id .. "\'"
    local found = false
    local grps = capture("groups")

    for g in grps:split("[ \n]") do
      if (g == group) then
         dbg.fini()
         return true
      end
    end
    LmodError(err_message,"\n")
    dbg.fini()
    return false
end

function logUsage(pkg)
    dbg.start{"logUsage(pkg)"}
    if (mode() ~= "load") then return true end
    local user = os.getenv("USER")
    local jobid = os.getenv("PBS_JOBID")
    local msg = ""
    dbg.print{"user: ",user," jobid: ",jobid,"\n"}
    if jobid == nil then
      msg = string.format("user=%s,app=%s", user, pkg.id)
    else
      msg = string.format("user=%s,app=%s,job=%s",
                          user, pkg.id, jobid)
    end
    local cmd = "logger -t lmod -p local0.info " .. msg
    os.execute(cmd)
    dbg.fini()
end

function prependModulePath(subdir)
    local mroot = os.getenv("MODULEPATH_ROOT")
    local mdir  = pathJoin(mroot, subdir)
    prepend_path("MODULEPATH", mdir)
end

function appendModulePath(subdir)
    local mroot = os.getenv("MODULEPATH_ROOT")
    local mdir  = pathJoin(mroot, subdir)
    append_path("MODULEPATH", mdir)
end

function dump(o)
    if type(o) == 'table' then
        local s = '{ '
        for k,v in pairs(o) do
            if type(k) ~= 'number' then k = '"'..k..'"' end
            s = s .. '['..k..'] = ' .. dump(v) .. ','
        end
        return s .. '} '
    else
        return tostring(o)
    end
end

--[[
hierarchyB

Designed to provide same functionality as hierarchyA
but designed when there is a sub-version

Example:

${MODULEPATH_ROOT}/Compiler/gcc/4.8.2/acml/gfortran64/5.3.1.lua

The above would be set using hierarchyB("acml/gfortran64_mp/5.3.1", 1) and
would return the expected output as if using hierarchyA

]]
function hierarchyB(pkgName, levels)
    local fn  = myFileName():gsub("%.lua$","")
    if (levels < 1) then
        return {}
    end

    -- Remove pkgName from end of string by using the
    -- "plain" matching via string.find function
    pkgName = path_regularize(pkgName:gsub("%.lua$",""))
    local i,j = fn:find(pkgName,1,true)
    if (j == fn:len()) then
      fn = fn:sub(1,i-1)
    end

    fn = path_regularize(fn)
    local j          = 0
    local numEntries = 0
    while (j) do
      j          = pkgName:find("/",j+1)
      numEntries = numEntries + 1
    end
    numEntries = numEntries - 1

    local a = {}

    for dir in fn:split("/") do
      a[#a + 1] = dir
    end

    local b = {}
    local n = #a

    for i = 1, levels do
      local bb = {}
      for j = 1, numEntries do
         local idx = n - numEntries + j
         bb[j] = a[idx]
      end
      b[i] = table.concat(bb,'/')
      n = n - numEntries
    end

    return b
end

--[[
setenv_ifunset

Sets environment variable only if not already set.  Prints a warning message
so user is aware a default was set.

Example:

setenv_ifunset("OMP_NUM_THREADS", "1")

This will set OMP_NUM_THREADS=1 if not already set in the user's environment.

]]
function setenv_ifunset(name, value)
  if (mode() == "load") then
    if (not os.getenv(name)) then
      local msg = string.format("WARNING: %s is not set.  Setting default of %s=%s", name, name, value)
      LmodMessage(msg)
      setenv(name, value)
    end
  end
end

function load_hook(t)
   -- the arg t is a table:
   --     t.modFullName:  the module full name: (i.e: gcc/4.7.2)
   --     t.fn:           The file name: (i.e /apps/modulefiles/Core/gcc/4.7.2.lua)

   -- Your site can use this any way that suits.  Here are some possibilities:
   --  a) Write this information into a file in your users directory (say ~/.lmod.d/.save).
   --     Then once a day/week/month collect this data.
   --  b) have this function call syslogd to register that this module was loaded by this
   --     user
   --  c) Write the same information directly to some database.

   -- This is an example writing to syslogd:

    if (mode() ~= "load") then return end
    local user  = os.getenv("USER")
    local jobid = os.getenv("SBATCH_JOBID") or "unknown"
    local msg   = string.format("user=%s,module=%s,fn=%s,job=%s", user, t.modFullName, t.fn, jobid)
    os.execute("logger -t lmod -p local0.info " .. msg)
end

--hook.register("load",load_hook)

------------------------------------------------------------------------
-- Any function that is called by a modulefile must be registered with
-- the sandbox as shown below.  Remember to use curly braces "{}" and
-- not parens "()" as you are sending a table as an argument.


sandbox_registration{ loadPkgDefaults      = loadPkgDefaults,
                      setPkgInfo           = setPkgInfo,
                      checkRestrictedGroup = checkRestrictedGroup,
                      logUsage             = logUsage,
                      prependModulePath    = prependModulePath,
                      appendModulePath     = appendModulePath,
                      setenv_ifunset       = setenv_ifunset
                    }