File Coverage

bin/runjob

Criterion	Covered	Total	%
statement	212	312	67.9
branch	82	182	45.0
condition	28	82	34.1
subroutine	19	24	79.1
pod			n/a
total	341	600	56.8

#!/usr/bin/env perl
#ABSTRACT: Run a job in the cluster using NBI::Slurm
#PODNAME: runjob

use v5.12;
use warnings;
use Getopt::Long;
use FindBin qw($RealBin);
use Data::Dumper;
use File::Basename;
use File::Spec;
use Term::ANSIColor qw(:constants);
use Cwd qw(abs_path getcwd);

my $BIN = basename($0);
$Data::Dumper::Sortkeys = 1;

if (-e "$RealBin/../dist.ini") {
    say STDERR "[dev mode] Using local lib" if ($ENV{"DEBUG"});
    use lib "$RealBin/../lib";
}

use NBI::Slurm;
use NBI::EcoScheduler;
use Cwd;

my $user_home_dir = $ENV{HOME};
my @afterok;
my $user_current_dir = getcwd();
my $username = $ENV{USER};
my $config = NBI::Slurm::load_config("$user_home_dir/.nbislurm.config");
my $version = $NBI::Slurm::VERSION;
my $queue = $config->{'queue'} // 'qib-short';
my $gpuqueue = $config->{'gpuqueue'} // 'qib-gpu';
my $threads = $config->{'threads'} // 1;
my $memory = $config->{'memory'} // 8000;
my $time = $config->{'time'} // "2h";
my $tmpdir = $config->{'tmpdir'} // "/tmp";
my $name;
my $email_address = $config->{'email'} // undef;
my $mail_type = $config->{'email_type'} // "none";
my $opt_placeholder = $config->{'placeholder'} // "#FILE#";
my $params_array_file;
my $command;
my $verbose;
my $debug;
my $run;
my $gpu;
my $opt_eco;
my $opt_no_eco;
my @slurm_options;
GetOptions(
    'm|memory=s'    => \$memory,
    'c|cores|threads=i' => \$threads,
    'q|queue=s'     => \$queue,
    'T|time=s'      => \$time,
    'w|tmpdir=s'    => \$tmpdir,
    'n|name=s'      => \$name,
    'r|run'         => \$run,
    'a|email-address=s' => \$email_address,
    'e|mail-type=s' => \$mail_type, # 'BEGIN,END,FAIL,REQUEUE,ALL'
    'after=s'       => \@afterok,
    'f|files=s'       => \my @files,
    'p|params-array=s' => \$params_array_file,
    'placeholder=s' => \$opt_placeholder,
    'gpu'           => \$gpu,
    'option=s'      => \@slurm_options,
    'eco'           => \$opt_eco,
    'no-eco'        => \$opt_no_eco,
    'verbose'       => \$verbose,
    'version'       => sub { say "runjob v", $NBI::Slurm::VERSION; exit },
    'debug'         => \$debug,
    'help'          => sub { usage(1) },
) or usage(1);

if ($opt_eco && $opt_no_eco) {
    say STDERR RED, BOLD, "ERROR:", RESET, " --eco and --no-eco are mutually exclusive";
    exit 1;
}

$debug = 1 if ($ENV{"DEBUG"});
$verbose = 1 if ($debug);

# Update queue based on time
my $time_hours = NBI::Opts::_time_to_hour($time);
$queue = update_queue($queue, $time_hours);
$command = join(" ", @ARGV);
usage(1) unless ($command);

if (@files and defined $params_array_file) {
    say STDERR RED, BOLD, "ERROR:", RESET, " --files and --params-array are mutually exclusive";
    exit 1;
}

my @abs_files = get_abs_files(@files);
my ($params_array_abs, $params_rows) = get_params_array_info($params_array_file, $command);
my $total_array_tasks = @abs_files ? scalar(@abs_files) : $params_rows;
my $array_task_limit = get_array_task_limit($config);
my @array_chunks = build_array_chunks($total_array_tasks, $array_task_limit);


if ($verbose) {
    say STDERR YELLOW, BOLD, "PARAMETERS:",RESET;
    say STDERR YELLOW, "Queue:      ", RESET, $queue // "(not set)";
    say STDERR YELLOW, "Memory:     ", RESET, $memory // "(not set)";
    say STDERR YELLOW, "Threads:    ", RESET, $threads // "(not set)";
    say STDERR YELLOW, "Time:       ", RESET, $time // "(not set)";
    say STDERR YELLOW, "Tmpdir:     ", RESET, $tmpdir // "(not set)";
    say STDERR YELLOW, "Name:       ", RESET, $name // "(not set)";
    say STDERR YELLOW, "Command:    ", RESET, $command // "(not set)";
    say STDERR YELLOW, "Params TSV: ", RESET, $params_array_abs if defined $params_array_abs;
    say STDERR YELLOW, "Array size: ", RESET, $total_array_tasks if $total_array_tasks;
    say STDERR YELLOW, "Array cap:  ", RESET, $array_task_limit if defined $array_task_limit;
}
# Check threadshours
if (-d $time) {
    say STDERR RED, BOLD, "ERROR:", RESET, " did you mean \"--tmpdir $time\"  instead of $threads threads?";
    usage();
}

# Sanitize tempdir
if (-d $tmpdir) {
    $tmpdir = File::Spec->rel2abs($tmpdir);
    say STDERR "[DEBUG] Using $tmpdir as temporary directory" if ($debug);
} else {
    # Check if it was meant to be a TIME instead
    if ($tmpdir =~/ / or $tmpdir =~/(\d+)[dh]/) {
        say STDERR RED, BOLD, "WARNING:", RESET, " $tmpdir looks like a time string, not a directory!";
    }
    eval {
        mkdir($tmpdir) or die "Cannot create $tmpdir: $!";
    };
    if ($@) {
        say STDERR RED, BOLD, "ERROR:", RESET, " Cannot create $tmpdir: $@";
        exit 1;
    } elsif ($verbose or $debug) {
        say STDERR "[INFO] Created $tmpdir";
    }
}

# Sanitize memory
if ($memory =~/^(\d+)$/ and $1 < 200) {
    say STDERR RED, BOLD, "WARNING:", RESET, " $memory is in Mb but <200, autoscaling to $memory GB";
    $memory = $1 * 1000;
}
$name = autoname($command) unless (defined $name);
say STDERR "[DEBUG] Using $name as job name" if ($debug);

my $afterok_string = defined $afterok[0] ? "-d afterok:" . join(":", @afterok) : undef;

# Prepare additional SLURM options
my @additional_opts = ();
push @additional_opts, $afterok_string if defined $afterok_string;
push @additional_opts, "--gres=gpu:1" if $gpu;

push @additional_opts, @slurm_options if @slurm_options;

$queue = $gpuqueue if $gpu;

#
if (not NBI::Slurm::valid_queue($queue)) {
    my @valid_queues = NBI::Slurm::queues('CAN FAIL');
    if (scalar @valid_queues) {
        say STDERR RED, BOLD, "WARNING:", RESET, " $queue is not a valid queue: ", join(", ", @valid_queues);
    } else {
        if ($run) {
            say STDERR RED, BOLD, "ERROR:", RESET, " You might be outside a slurm cluster. Run without --run to see the script", join(", ", @valid_queues);
            exit 1;
        } else {
            say STDERR RED, BOLD, "WARNING:", RESET, " You might be outside a slurm cluster", join(", ", @valid_queues);
        }


    }

}

# Eco scheduling
# eco_default from config defaults to 1 (on) if not explicitly set
my $eco_default = defined $config->{eco_default} ? $config->{eco_default} : $NBI::EcoScheduler::DEFAULTS{eco_default};
my $use_eco = $opt_no_eco ? 0
            : $opt_eco    ? 1
            :               $eco_default;

if ($use_eco) {
    if (defined $config->{start_time} || defined $config->{start_date}) {
        say STDERR YELLOW, "[eco] --begin already set explicitly, ignoring eco scheduling", RESET;
    }
}

sub xdump {
    my ($hash_of_hashes, $indent) = @_;
    $indent //= '  ';
    while (my ($key, $value) = each %$hash_of_hashes) {
        if (ref $value eq 'HASH') {
            print STDERR BOLD, "$indent$key :\n", RESET;
            xdump($value, "$indent    ");
            print "$indent\n";
        } else {
            print STDERR "$indent$key = $value\n";
        }
    }
}


if (@array_chunks > 1) {
    say STDERR YELLOW, "[array] Splitting $total_array_tasks tasks into ", scalar(@array_chunks), " jobs", RESET;
}

my @jobs;
my $name_width = scalar(@array_chunks) > 1 ? length(scalar(@array_chunks)) : 0;
for my $chunk_idx (0 .. $#array_chunks) {
    my ($array_offset, $chunk_tasks) = @{$array_chunks[$chunk_idx]};
    my @job_files = @abs_files;
    my $job_params_array = $params_array_abs;
    my $job_params_rows = $params_rows;
    my $job_array_offset = 0;

    if (@abs_files && @array_chunks > 1) {
        @job_files = @abs_files[$array_offset .. ($array_offset + $chunk_tasks - 1)];
    } elsif (defined $params_array_abs) {
        $job_array_offset = $array_offset;
    }

    my $job_name = $name;
    if (@array_chunks > 1) {
        $job_name = sprintf("%s.part%0*d", $name, $name_width, $chunk_idx + 1);
    }

    my $opts = NBI::Opts->new(
        -queue => $queue,
        -threads => $threads,
        -memory => $memory,
        -time   => $time,
        -tmpdir => $tmpdir,
        -email_address => $email_address,
        -email_type => $mail_type,
        -opts => \@additional_opts,
        -files => \@job_files,
        -params_array => $job_params_array,
        -params_rows => $job_params_rows,
        -placeholder => $opt_placeholder,
        -array_offset => $job_array_offset,
        -array_tasks => $chunk_tasks,
    );

    if ($use_eco) {
        my ($begin_epoch, $tier) = NBI::EcoScheduler::find_eco_begin($opts->hours, $config);
        if (defined $begin_epoch) {
            my $begin_str = NBI::EcoScheduler::epoch_to_slurm($begin_epoch);
            my $delay     = NBI::EcoScheduler::format_delay($begin_epoch);
            ($opts->{start_date}, $opts->{start_time}) = split /T/, $begin_str;
            my $tier_warn = $tier == 2 ? " (job may overrun eco window)"
                          : $tier == 3 ? " (job may overlap peak hours — best available)"
                          :              "";
            say STDERR GREEN, "[eco] Job scheduled for $begin_str (in $delay)$tier_warn", RESET if $chunk_idx == 0;
        } else {
            say STDERR YELLOW, "[eco] No eco slot found in lookahead period, submitting immediately", RESET if $chunk_idx == 0;
        }
    }

    if ($verbose) {
        say STDERR GREEN, "CONFIG:\n", RESET if $chunk_idx == 0;
        xdump($config) if $chunk_idx == 0;
        print STDERR "\n" if $chunk_idx == 0;
        say STDERR $opts->view();
    }

    my $job = NBI::Job->new(
        -name => $job_name,
        -command =>  "cd \"$user_current_dir\"",
        -opts => $opts,
    );
    if ($debug) {
        say STDERR YELLOW, "JOB:\n", RESET, Dumper($job);
    }
    $job->append_command($command);
    push @jobs, $job;
}


my $err = 0;
if ($run) {
    my @job_ids;
    for my $job (@jobs) {
        my $j = $job->run();
        if ($j) {
            if ($verbose) {
                say STDERR $job->view();
            }
            push @job_ids, $j;
        } else {
            say STDERR RED, "Job not submitted: $j", RESET;
            if ($verbose) {
                say STDERR $job->view();
            }
            $err = 1;
            last;
        }
    }
    if (!$err) {
        say $_ for @job_ids;
    }
} else {
    for my $idx (0 .. $#jobs) {
        say STDERR GREEN, "JOB SCRIPT\n", RESET if @jobs == 1;
        say STDERR GREEN, "JOB SCRIPT ", ($idx + 1), "/", scalar(@jobs), "\n", RESET if @jobs > 1;
        say $jobs[$idx]->script();
    }
}


if ($err) {
    exit 1;
} elsif ($run) {
    my @submitted = map { $_->jobid } @jobs;
    say STDERR GREEN, "Jobs submitted: ", join(",", @submitted), RESET;
}
sub usage {
    say STDERR <

 Options:
    -n, --name       Job name [optional]
    -q, --queue      Queue name [default: nbi-short]
    -m, --memory     Memory to use [default: 8Gb]
    -c, --cores      Number of threads [default: 1]
    -T, --time       Time string [default: "0d 8h"]
    --after  INT     Job ID to wait for [can be used multiple times]
    -w, --tmpdir     Temporary directory [default: /tmp]
    -r, --run        Run the job (otherwise, just print the script)
    -f, --files      Input files for array (either specify -f as many times as needed or
                     use a _quoted_ pattern like "*.fasta" to include multiple files)
    -p, --params-array
                     TSV file for array jobs; refer to columns as ##1##, ##2##, ...
                     Mutually exclusive with --files
    --placeholder    Placeholder for array input files [default: #FILE#]
                     Ignored when --params-array is used
    --gpu            Request GPU resources (adds --gres=gpu:1)
    --option VALUE   Additional SLURM options (can be used multiple times)
    --eco            Schedule job in next low-energy window (overrides eco_default=0)
    --no-eco         Submit immediately, ignoring eco scheduling (overrides eco_default=1)
    --verbose        Verbose output
    --help           This help message
 ----------------------------------------------------------
END
    exit() if ($_[0]);
}

sub autoname {
    my $string = shift;
    my @parts = split(/\s+/, $string);
    my @ints = ("bash", "perl", "python", "python3", "R", "Rscript", "sh", "zsh");
    my @subs = ("bwa", "samtools", "seqfu", "seqkit", "bedtools", "taxonkit", "kmcp", "seqtk", "usearch", "vsearch");

    # From each @parts, replace it with the first [A-Za-z0-9]+ part
    for my $i (0..$#parts) {
        $parts[$i] =~ /^([A-Za-z0-9]+)/;
        $parts[$i] = $1;
    }
    # check if $fisrt is part of the @int array of strings
    if (grep {$_ eq $parts[0]} @ints and defined $parts[1]) {
        return $parts[1];
    } elsif (grep {$_ eq $parts[0]} @subs and defined $parts[1]) {
        return $parts[0] . "-" . $parts[1];
    } else {
        return $parts[0];
    }

}


sub update_queue {
    my ($queue, $time) = @_;
    # If the queue has a star, it can be updated
    if ($queue !~ /\*/) {
        return $queue;
    }
    # Get time in hours
    my $mock_opt = NBI::Opts->new(-time => $time);
    my $time_h = $mock_opt->hours;

    if ($time <= 2 ) {
        # Replace * with short
        $queue =~ s/\*/short/g;
    } elsif ($time_h <= 48) {
        # Replace * with medium
        $queue =~ s/\*/medium/g;
    } else {
        # Replace * with long
        $queue =~ s/\*/long/g;
    }
    return $queue;
}


sub expand_pattern_to_abs_paths {
    my ($pattern) = @_;

    # Get the current working directory
    my $current_dir = getcwd();

    # Expand the pattern
    my @files = glob($pattern);

    # Convert to absolute paths
    my @abs_paths = map { File::Spec->rel2abs($_, $current_dir) } @files;

    return @abs_paths;
}

sub get_abs_files {
    my @files = @_;
    for my $file (@files) {

        if (-e $file) {
            push @abs_files, abs_path($file);
            next;
        } else {
            my @tmpfiles = ();
            push @tmpfiles, expand_pattern_to_abs_paths($file);
            if (scalar @tmpfiles) {
                push @abs_files, @tmpfiles;
            } else {
                say STDERR RED, BOLD, "ERROR:", RESET, " $file does not exist!";
                exit 1;
            }
        }

    }
    return @abs_files;
}

sub get_params_array_info {
    my ($file, $command) = @_;
    return (undef, 0) unless defined $file;

    unless (-e $file) {
        say STDERR RED, BOLD, "ERROR:", RESET, " $file does not exist!";
        exit 1;
    }

    my $abs_file = abs_path($file);
    open(my $fh, "<", $abs_file) or do {
        say STDERR RED, BOLD, "ERROR:", RESET, " Cannot read $abs_file: $!";
        exit 1;
    };

    my %placeholders;
    while ($command =~ /##(\d+)##/g) {
        $placeholders{$1} = 1;
    }
    my $max_placeholder = 0;
    for my $n (keys %placeholders) {
        $max_placeholder = $n if $n > $max_placeholder;
    }
    if ($max_placeholder == 0) {
        say STDERR RED, BOLD, "ERROR:", RESET, " --params-array requires placeholders like ##1## in the command";
        exit 1;
    }

    my $rows = 0;
    while (my $line = <$fh>) {
        chomp $line;
        $line =~ s/\r$//;
        next if $line =~ /^\s*$/;
        next if $line =~ /^\s*#/;
        my @fields = split /\t/, $line, -1;
        if (scalar(@fields) < $max_placeholder) {
            say STDERR RED, BOLD, "ERROR:", RESET, " $abs_file row ", ($rows + 1), " has ", scalar(@fields), " columns but the command requires ##$max_placeholder##";
            exit 1;
        }
        $rows++;
    }
    close $fh;

    if ($rows == 0) {
        say STDERR RED, BOLD, "ERROR:", RESET, " $abs_file does not contain any usable TSV rows";
        exit 1;
    }

    return ($abs_file, $rows);
}

sub build_array_chunks {
    my ($total_tasks, $task_limit) = @_;
    return ([0, 0]) unless $total_tasks;
    return ([0, $total_tasks]) unless defined $task_limit && $task_limit > 0 && $total_tasks > $task_limit;

    my @chunks;
    for (my $offset = 0; $offset < $total_tasks; $offset += $task_limit) {
        my $chunk_tasks = $task_limit;
        my $remaining = $total_tasks - $offset;
        $chunk_tasks = $remaining if $remaining < $chunk_tasks;
        push @chunks, [$offset, $chunk_tasks];
    }
    return @chunks;
}

sub get_array_task_limit {
    my ($config) = @_;

    if (defined $ENV{NBI_MAX_ARRAY_SIZE} && $ENV{NBI_MAX_ARRAY_SIZE} =~ /^\d+$/ && $ENV{NBI_MAX_ARRAY_SIZE} > 0) {
        return $ENV{NBI_MAX_ARRAY_SIZE};
    }
    if (defined $config->{max_array_size} && $config->{max_array_size} =~ /^\d+$/ && $config->{max_array_size} > 0) {
        return $config->{max_array_size};
    }

    my $output = `scontrol show config 2>/dev/null`;
    return undef if $? != 0 || !$output;

    my @limits;
    if ($output =~ /^MaxArraySize\s*=\s*(\d+)/m) {
        push @limits, $1 if $1 > 0;
    }
    if ($output =~ /max_array_tasks=(\d+)/) {
        push @limits, $1 if $1 > 0;
    }
    return undef unless @limits;

    my $limit = shift @limits;
    for my $candidate (@limits) {
        $limit = $candidate if $candidate < $limit;
    }
    return $limit;
}

__END__