File Coverage

bin/session

Criterion	Covered	Total	%
statement	58	107	54.2
branch	8	44	18.1
condition	9	18	50.0
subroutine	12	18	66.6
pod			n/a
total	87	187	46.5

#!/usr/bin/env perl
#ABSTRACT: Start an interactive session on the cluster
#PODNAME: session
use v5.16;
use Getopt::Long;
use File::Basename;
use Carp qw(confess);
use Term::ANSIColor qw(:constants);
use FindBin qw($RealBin);
if (-e "$RealBin/../dist.ini") {
    say STDERR "[dev mode] Using local lib" if ($ENV{"DEBUG"});
    use lib "$RealBin/../lib";
}
use NBI::Slurm;
my $BIN = basename($0);

main() unless caller;
1;

sub main {
    my $USER = $ENV{USER};
    my $config = NBI::Slurm::load_config("$USER/.nbislurm.config");
    my $DEFAULT_SESSION = $config->{'session'} // "";
    my $SPECIAL_SESSION = $config->{'special_session'} // "";
    my $DEFAULT_HOURS = $config->{'session_hours'} // 4;
    my $NAME = "session";
    my $MEM_MB = $config->{'session_memory'} // 1000;
    my $CORES = $config->{'session_cpus'} // 1;
    my $DAYS = 0;
    my $OPT_HOURS;
    my $OPT_TIME;
    my $OPT_INTEL;
    my $OPT_PARTITION;
    my $OPT_SPECIAL;
    my $OPT_VERBOSE;

    GetOptions(
        'c|threads=i' => \$CORES,
        'm|mem=s' => \$MEM_MB,
        'd|days=i' => \$DAYS,
        'h|hours=i' => \$OPT_HOURS,
        'T|t|time=s' => \$OPT_TIME,
        'n|name=s' => \$NAME,
        'q|queue=s' => \$OPT_PARTITION,
        's|special' => \$OPT_SPECIAL,
        'i|intel' => \$OPT_INTEL,
        'verbose' => \$OPT_VERBOSE,
        'version' => sub { say "session v", $NBI::Slurm::VERSION; exit(0); },
        'help' => sub { usage(); exit(0); },
    ) or usage(1);

    # Get memory
    $MEM_MB = parse_memory($MEM_MB);

    if ($MEM_MB < 350) {
        say STDERR RED, "[WARNING]", RESET, "Memory set to $MEM_MB Mb, automatically setting to $MEM_MB Gb", RESET;
        $MEM_MB *= 1000;
    }

    my $time = session_resolve_time(
        hours         => $OPT_HOURS,
        time_string   => $OPT_TIME,
        default_hours => $DEFAULT_HOURS,
        days          => $DAYS,
    );

    # Get queue
    if (not $OPT_PARTITION) {
        if ($config->{'queue'}) {
            $OPT_PARTITION = $config->{'queue'};
        } else {
            $OPT_PARTITION = partition_tag($time->{total_hours});
        }
    }

    my $PARAMS = $DEFAULT_SESSION;
    $PARAMS   .= $OPT_INTEL    ?  " --constraint=intel " : "";
    $PARAMS   .= $OPT_SPECIAL  ?  " $SPECIAL_SESSION "   : "";

    my $PARTITION //= $OPT_PARTITION;

    say STDERR "-" x 40;
    print STDERR YELLOW,
            "job name    ", RESET, $NAME, "\n",
    YELLOW, "memory      ", RESET, $MEM_MB, " Mb\n",
    YELLOW, "cores       ", RESET, $CORES, "\n",
    YELLOW, "time        ", RESET, $time->{slurm}, "\n",
    YELLOW, "partition   ", RESET, $PARTITION, "\n";

    say STDERR "-" x 40;
    my $cmd = qq(srun --job-name=$NAME --mem=$MEM_MB -c $CORES --time=$time->{slurm} -n 1 --disable-status -s --pty --partition=$PARTITION $PARAMS -u bash -li);

    say STDERR $cmd if ($OPT_VERBOSE);

    if (has_srun() == 0) {
        say STDERR RED, "[ERROR] srun not found... Are you on the cluster?", RESET;
    }
    exec($cmd);
}

sub partition_tag {
    my $hours = shift;
    if ($hours <= 2) {
        return "qib-short,nbi-short";
    } elsif ($hours <= 8) {
        return "qib-medium,nbi-medium";
    } else {
        return "qib-long,nbi-long";
    }
}

sub session_resolve_time {
    my %args = @_;
    confess "ERROR: --hours and --time are mutually exclusive\n"
        if defined $args{hours} && defined $args{time_string};

    my $hours = defined $args{time_string}
        ? NBI::Opts::_time_to_hour($args{time_string})
        : defined $args{hours}
            ? $args{hours}
            : $args{default_hours};

    $hours += ($args{days} // 0) * 24;

    return {
        total_hours => $hours,
        slurm       => session_hours_to_slurm($hours),
    };
}

sub hours_to_days_hours {
    my $hours = shift;
    my $days = int($hours / 24);
    my $hours = $hours % 24;
    return ($days, $hours);
}

sub session_hours_to_slurm {
    my $hours = shift;
    my $seconds = int($hours * 3600);
    my $days = int($seconds / 86400);
    $seconds %= 86400;
    my $hh = int($seconds / 3600);
    $seconds %= 3600;
    my $mm = int($seconds / 60);
    my $ss = $seconds % 60;
    return sprintf("%d-%02d:%02d:%02d", $days, $hh, $mm, $ss);
}


sub parse_memory {
    # Parse string return mb
    my $mem = shift;
    my $mb = 0;
    if ($mem =~ /^(\d+)$/) {
        $mb = $1;
    } elsif ($mem =~ /(\d+)\s?Gb?/i) {
        $mb = $1 * 1000;
    } elsif ($mem =~ /(\d+)\s?Mb?/i) {
        $mb = $1;
    } elsif ($mem =~ /(\d+)\s?Kb?/i) {
        $mb = int($1 / 1000);
    } else {
        $mb = $mem;
    }
    return $mb;
}
sub usage {
    print <<"EOF";
Usage: $BIN [options]

Options:
  -c, --threads   Number of CPU cores (default: 1)
  -m, --mem       Memory size in MB (default: 1000)
  -h, --hours     Number of hours (default: 4; mutually exclusive with --time)
  -t, --time      Time string like 12h, 1d, 2h30m
  -n, --name      Job name (default: session)
  -q, --queue     Queue/partition name (default: determined by hours)
  -s, --special   Enable the special parameter string from the config file (default: off)
  -i, --intel     Use Intel constraint (default: off)

EOF
  exit() if ($_[0]);
}

sub has_srun {
    # Check if srun is available
    my $srun = `srun --version 2> /dev/null`;
    chomp $srun;
    if ($srun eq "") {
        return 0;
    }
    return 1;
}

__END__