File Coverage

blib/lib/NBI/Slurm.pm

Criterion	Covered	Total	%
statement	28	106	26.4
branch	4	40	10.0
condition	0	6	0.0
subroutine	8	13	61.5
pod	7	7	100.0
total	47	172	27.3

#ABSTRACT: NBI Slurm module
#
# NBI::Slurm - Main entry-point module for the NBI::Slurm package.
#
# DESCRIPTION:
#   Provides utility functions and constants used across the package:
#     - %FORMAT_STRINGS : squeue format codes (jobid, user, memory, etc.)
#     - load_config()   : reads ~/.nbislurm.config key=value settings
#     - has_squeue()    : checks whether the squeue binary is available
#     - queues()        : lists available SLURM partitions via sinfo
#     - valid_queue()   : validates a queue name against the cluster
#     - execute_command(): runs a shell command and captures stdout/stderr/exit
#     - timelog()       : returns a formatted timestamp string for logging
#     - days_since_update(): returns days since a file was last modified
#
# RELATIONSHIPS:
#   - Loads and re-exports NBI::Job and NBI::Opts so callers only need
#     "use NBI::Slurm".
#   - NBI::Queue and NBI::QueuedJob both import %FORMAT_STRINGS and
#     $NBI::Slurm::VERSION from this module.
#   - NBI::EcoScheduler uses $NBI::Slurm::VERSION and is called by
#     bin/runjob when eco scheduling is active.
#
use strict;
use warnings;

package NBI::Slurm;
use NBI::Job;
use NBI::Opts;
use base qw(Exporter);
our @ISA = qw(Exporter);
our @EXPORT = qw(Job Opts load_config has_squeue timelog execute_command %FORMAT_STRINGS);

$NBI::Slurm::VERSION = '0.17.2';




our %FORMAT_STRINGS = (
 'account'    => '%a',
 'jobid'      => '%A',
 'jobname'    => '%j',
 'cpus'       => '%C',
 'end_time'   => '%E',
 'start_time' => '%S',
 'total_time' => '%l',
 'time_left'  => '%L',
 'memory'     => '%m',
 'command'    => '%o',
 'queue'      => '%P',
 'reason'     => '%r',
 'status'     => '%T', # short: %t
 'workdir'    => '%Z',
 'user'       => '%u',
);

sub execute_command {
    my ($command) = @_;

    # Use File::Temp functionality available in core Perl
    use File::Temp qw(tempfile);

    # Create temporary files for capturing output
    my ($stdout_fh, $stdout_file) = tempfile(UNLINK => 1);
    my ($stderr_fh, $stderr_file) = tempfile(UNLINK => 1);
    close($stdout_fh);  # Close handles so system() can write to files
    close($stderr_fh);

    # Execute command with output redirection
    my $full_command = "$command >$stdout_file 2>$stderr_file";
    system($full_command);

    # Capture exit code (system() returns exit code << 8)
    my $exit_code = $? >> 8;

    # Read stdout
    my $stdout = '';
    if (open(my $stdout_read_fh, '<', $stdout_file)) {
        $stdout = do { local $/; <$stdout_read_fh> };
        close($stdout_read_fh);
    }

    # Read stderr
    my $stderr = '';
    if (open(my $stderr_read_fh, '<', $stderr_file)) {
        $stderr = do { local $/; <$stderr_read_fh> };
        close($stderr_read_fh);
    }

    # Files are automatically cleaned up due to UNLINK => 1

    return {
        stdout    => $stdout,
        stderr    => $stderr,
        exit_code => $exit_code
    };
}

sub load_config {
    my $filename = shift;
    if (! $filename) {
        $filename = "$ENV{HOME}/.nbislurm.config";
    }
    my $config = {};
    if (! -e "$filename") {
        say STDERR "# Config file not found: $filename" if ($ENV{"DEBUG"});
        return $config;
    }
    open(my $fh, "<", $filename) or die "Cannot open $filename: $!";
    while (<$fh>) {
        chomp;
        next if (/^\s*$/);
        next if (/^#/);
        next if (/^;/);
        my ($key, $value) = split(/=/, $_);
        # discard keys with spaces
        next if ($key =~ /\s/);
        $config->{$key} = $value;
    }
    close  $fh;
    return $config;
}


sub has_squeue {
    my $cmd = "squeue --version";
    my $output = `$cmd 2>&1`;
    if ($? == 0) {
        return 1;
    } else {
        return 0;
    }
}

sub queues {
  my $can_fail = shift;
  # Retrieve queues from SLURM
  my $has_sinfo = undef;
  eval {
    $has_sinfo = `sinfo --version > /dev/null 2>&1`;
  };

  chomp($has_sinfo) if defined $has_sinfo;
  if (not defined $has_sinfo and ! $can_fail) {
    Carp::croak "ERROR NBI::Slurm: sinfo failed. Are you in a SLURM cluster?\n";
  }
  my $cmd = "timeout 5s sinfo --format '%P' --noheader";
  my @output = `$cmd 2>/dev/null`;
  if ($? != 0 and ! $can_fail) {
    Carp::croak "ERROR NBI::Slurm: sinfo did not find queues. Are you in a SLURM cluster?\n";
  }

  chomp @output;
  return @output;
}

sub valid_queue {
  if ($ENV{'SKIP_SLURM_CHECK'}) {
    return 1;
  }
  my $queue = shift;
  my @queues = queues('CAN_FAIL');
  my @input_queues = split(/,/, $queue);

  if (scalar(@input_queues) == 0) {
    # Let's assume it exists... TODO CHECK
    return  1;
  }
  foreach my $input_queue (@input_queues) {
    if (! grep { $_ eq $input_queue } @queues) {
      return 0;
    }
  }
  return 1;
}


sub days_since_update {
    my $file_path = shift;

    # Check if the required modules can be loaded
    eval {
        require File::Spec;
        require Time::Piece;
        require Time::Seconds;
    };
    if ($@) {
        return -1;  # Failed to load required module(s)
    }

    # Check if the file exists
    unless (-e $file_path) {
        return -1;  # File not found
    }

    # Get the file's last modification time
    my $last_modified = (stat($file_path))[9];

    # Calculate the number of days since the last modification
    my $current_time = time();
    my $days_since_update = int(($current_time - $last_modified) / (24 * 60 * 60));

    return $days_since_update;
}

sub timelog {
    my $tag = shift;
    $tag = "nbi-slurm" unless defined $tag;
    my ($sec, $min, $hour, $day, $month, $year) = localtime(time);
    $year += 1900;  # Adjust year (localtime returns years since 1900)
    $month += 1;    # Adjust month (localtime returns 0-11)

    # Format with leading zeros
    return sprintf("[%s:%04d-%02d-%02d %02d:%02d:%02d]\t",
            $tag, $year, $month, $day, $hour, $min, $sec);
}

1;

__END__