Files
agosti_j 9c6cfe3e21
Build Release Binaries / Build linux-amd64 (push) Successful in 9s
Build Release Binaries / Build linux-arm64 (push) Successful in 9s
Correct lanes create for versions below 5
2026-08-25 16:52:43 +02:00

1006 lines
22 KiB
Go

package command
import (
"archive/tar"
"compress/gzip"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net/http"
"os"
"os/exec"
"path/filepath"
"regexp"
"strconv"
"strings"
"time"
)
var ErrCancelled = errors.New("cancelled by user")
// =============================================================================
// Cryosparcm command
// =============================================================================
func cryosparcmCmd(
hostname,
cryosparcpath,
commandpath string,
help bool,
args ...string) *exec.Cmd {
var cmd *exec.Cmd
remoteCmd := []string{"cryosparcm"}
if cryosparcpath != "" {
remoteCmd = append(remoteCmd,
"--dir",
filepath.Join(cryosparcpath, "cryosparc_master"),
)
}
if commandpath != "" {
remoteCmd = append(remoteCmd,
"--cwd",
commandpath,
)
}
remoteCmd = append(remoteCmd, args...)
if help {
cmd = exec.Command(
filepath.Join(
cryosparcpath,
"cryosparc_master",
"bin",
"cryosparcm"),
"--help",
)
} else if len(args) > 1 && args[len(args)-1] == "--help" {
cmd = exec.Command(
filepath.Join(cryosparcpath,
"cryosparc_master",
"bin",
"cryosparcm"),
args...,
)
} else {
cmd = exec.Command(
"ssh",
fmt.Sprintf("%s@%s", os.Getenv("USER"), hostname),
strings.Join(remoteCmd, " "),
)
}
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}
// =============================================================================
// Cryosparcw command
// =============================================================================
func cryosparcwCmd(
cryosparcPath string,
arch string,
help bool,
args ...string,
) *exec.Cmd {
var cmdArgs []string
// Get local architecture.
localArch, err := exec.Command("uname", "-m").Output()
if err != nil {
localArch = nil
}
isLocal := strings.TrimSpace(string(localArch)) == arch
// CryoSPARC >= 5 supports architecture-specific worker
// installations. Older versions use the generic worker path.
workerDir := fmt.Sprintf("cryosparc_worker_%s", arch)
versionData, err := os.ReadFile(
filepath.Join(cryosparcPath, "cryosparc_master", "version"),
)
if err == nil {
version := strings.TrimPrefix(strings.TrimSpace(string(versionData)), "v")
major, _, _ := strings.Cut(version, ".")
if majorVersion, err := strconv.Atoi(major); err == nil &&
majorVersion < 5 {
workerDir = "cryosparc_worker"
}
}
workerCmd := filepath.Join(
cryosparcPath,
workerDir,
"bin",
"cryosparcw",
)
if !isLocal && arch == "aarch64" {
// Run on ARM nodes.
cmdArgs = []string{
"srun",
"--cluster", "gmerlin7",
"--partition", "gh-interactive",
"--time", "0-00:01:00",
workerCmd,
}
} else {
cmdArgs = []string{workerCmd}
}
if help {
cmdArgs = append(cmdArgs, "--help")
} else {
cmdArgs = append(cmdArgs, args...)
}
cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}
// =============================================================================
// Download Master
// =============================================================================
type progressReader struct {
io.Reader
Total int64
Current int64
Update func(float64)
}
func (r *progressReader) Read(p []byte) (int, error) {
n, err := r.Reader.Read(p)
if n > 0 {
r.Current += int64(n)
if r.Total > 0 {
r.Update(float64(r.Current) / float64(r.Total))
}
}
return n, err
}
func (i *Command) downloadCryosparc(
update func(float64),
downloaddir,
release,
license,
installation,
arch string) error {
url := fmt.Sprintf(
"https://get.cryosparc.com/download/%s-v%s/%s?arch=%s",
installation,
release,
license,
arch,
)
archive := filepath.Join(
downloaddir,
fmt.Sprintf("cryosparc_%s.tar.gz",
installation))
resp, err := http.Get(url)
if err != nil {
return fmt.Errorf(
"download CryoSPARC %s package: %w",
installation,
err)
}
defer func() {
if err := resp.Body.Close(); err != nil {
log.Printf("failed to close response body: %v", err)
}
}()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf(
"download failed (HTTP %d). "+
"Please check the CryoSPARC release, license ID, "+
"architecture, and internet connection",
resp.StatusCode,
)
}
f, err := os.Create(archive)
if err != nil {
return fmt.Errorf("create %q: %w", archive, err)
}
defer func() {
if err := f.Close(); err != nil {
log.Printf("failed to close file: %v", err)
}
}()
reader := &progressReader{
Reader: resp.Body,
Total: resp.ContentLength,
Update: update,
}
if _, err := io.Copy(f, reader); err != nil {
return fmt.Errorf("save %q: %w", archive, err)
}
return nil
}
// =============================================================================
// Directory Check
// =============================================================================
func directoryExistsAndNotEmpty(path string) (bool, error) {
info, err := os.Stat(path)
switch {
case errors.Is(err, os.ErrNotExist):
return false, nil
case err != nil:
return false, fmt.Errorf("stat %q: %w", path, err)
case !info.IsDir():
return false, fmt.Errorf("%q exists but is not a directory", path)
}
f, err := os.Open(path)
if err != nil {
return false, fmt.Errorf("open %q: %w", path, err)
}
defer func() {
if err := f.Close(); err != nil {
log.Printf("failed to close file: %v", err)
}
}()
_, err = f.Readdirnames(1)
switch {
case errors.Is(err, io.EOF):
return false, nil // directory exists but is empty
case err != nil:
return false, fmt.Errorf("read %q: %w", path, err)
default:
return true, nil // directory exists and contains at least one entry
}
}
// =============================================================================
// Create Directory
// =============================================================================
func (i *Command) createDirectory(_ func(float64), path string) error {
if err := os.MkdirAll(path, 0o755); err != nil {
return err
}
return nil
}
// =============================================================================
// Create backup
// =============================================================================
func (i *Command) backupCryosparcDatabase(
_ func(float64),
cryosparcpath,
dbpath,
backuppath string) error {
var version string
var destination string
date := time.Now().Format("2006-01-02T15:04:05")
baseName := filepath.Base(filepath.Clean(dbpath))
versionData, err := os.ReadFile(
filepath.Join(cryosparcpath, "cryosparc_master", "version"))
if err == nil {
version = strings.TrimSpace(string(versionData))
version = strings.TrimPrefix(version, "v")
destination = filepath.Join(
backuppath,
fmt.Sprintf("%s.v%s.%s.backup", baseName, version, date),
)
} else {
destination = filepath.Join(
backuppath,
fmt.Sprintf("%s.%s.backup", baseName, date),
)
}
if err := copyDir(dbpath, destination); err != nil {
return fmt.Errorf("create backup: %w", err)
}
return nil
}
func copyDir(src, dst string) error {
info, err := os.Stat(src)
if err != nil {
return err
}
if err := os.MkdirAll(dst, info.Mode().Perm()); err != nil {
return err
}
entries, err := os.ReadDir(src)
if err != nil {
return err
}
for _, entry := range entries {
srcPath := filepath.Join(src, entry.Name())
dstPath := filepath.Join(dst, entry.Name())
if entry.IsDir() {
if err := copyDir(srcPath, dstPath); err != nil {
return err
}
continue
}
if err := copyFile(srcPath, dstPath); err != nil {
return err
}
}
return nil
}
func copyFile(src, dst string) error {
in, err := os.Open(src)
if err != nil {
return err
}
defer func() {
if err := in.Close(); err != nil {
log.Printf("failed to close file: %v", err)
}
}()
info, err := in.Stat()
if err != nil {
return err
}
out, err := os.OpenFile(
dst,
os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
info.Mode().Perm(),
)
if err != nil {
return err
}
defer func() {
if err := out.Close(); err != nil {
log.Printf("failed to close file: %v", err)
}
}()
_, err = io.Copy(out, in)
return err
}
// =============================================================================
// Extract archive
// =============================================================================
func (i *Command) extractArchive(
_ func(float64),
downloaddir,
extractdir,
filename string) error {
archive := filepath.Join(
downloaddir,
filename,
)
file, err := os.Open(archive)
if err != nil {
return fmt.Errorf("open archive %q: %w", archive, err)
}
defer func() {
if err := file.Close(); err != nil {
log.Printf("failed to close file: %v", err)
}
}()
gz, err := gzip.NewReader(file)
if err != nil {
return fmt.Errorf("read gzip archive: %w", err)
}
defer func() {
if err := gz.Close(); err != nil {
log.Printf("failed to close gzip: %v", err)
}
}()
tarReader := tar.NewReader(gz)
for {
header, err := tarReader.Next()
switch {
case err == io.EOF:
return nil
case err != nil:
return fmt.Errorf("read archive: %w", err)
}
target := filepath.Join(
extractdir,
header.Name,
)
switch header.Typeflag {
case tar.TypeDir:
if err := os.MkdirAll(target, 0o755); err != nil {
return fmt.Errorf(
"create directory %q: %w",
target,
err,
)
}
case tar.TypeReg:
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
return fmt.Errorf(
"create parent directory %q: %w",
target,
err,
)
}
out, err := os.OpenFile(
target,
os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
os.FileMode(header.Mode),
)
if err != nil {
return fmt.Errorf(
"create file %q: %w",
target,
err,
)
}
_, copyErr := io.Copy(out, tarReader)
closeErr := out.Close()
if copyErr != nil {
return fmt.Errorf(
"extract %q: %w",
target,
copyErr,
)
}
if closeErr != nil {
return fmt.Errorf(
"close %q: %w",
target,
closeErr,
)
}
}
}
}
// =============================================================================
// Rename directory
// =============================================================================
func (i *Command) renameDirectory(from, to string) error {
if err := os.Rename(from, to); err != nil {
return fmt.Errorf("failed to rename directory %q to %q: %w", from, to, err)
}
return nil
}
// =============================================================================
// Install Master
// =============================================================================
func (i *Command) installMaster(
cryosparcpath,
license,
hostname,
dbpath,
ssdpath,
arch string,
baseport uint) *exec.Cmd {
installScript := filepath.Join(
cryosparcpath,
"cryosparc_master",
"install.sh",
)
installArgs := []string{
installScript,
"--license", license,
"--hostname", hostname,
"--dbpath", dbpath,
"--ssdpath", ssdpath,
"--port", strconv.Itoa(int(baseport)),
}
// Get local architecture.
localArch, err := exec.Command("uname", "-m").Output()
if err != nil {
localArch = nil
}
isLocal := strings.TrimSpace(string(localArch)) == arch
var args []string
if isLocal {
// Run locally.
args = installArgs
} else if arch == "aarch64" {
// Run on ARM nodes.
args = append([]string{
"srun",
"--cluster", "gmerlin7",
"--partition", "gh-interactive",
"--time", "0-00:30:00"},
installArgs...)
} else {
// Run on x86 nodes.
args = append([]string{
"srun",
"--cluster", "merlin7",
"--partition", "interactive",
"--reservation", "interactive",
"--time", "0-00:30:00"},
installArgs...)
}
versionData, err := os.ReadFile(
filepath.Join(cryosparcpath, "cryosparc_master", "version"))
if err == nil {
version := strings.TrimSpace(string(versionData))
version = strings.TrimPrefix(version, "v")
major, _, ok := strings.Cut(version, ".")
if !ok {
major = version
}
if majorVersion, err := strconv.Atoi(major); err == nil &&
majorVersion >= 5 {
args = append(args, "--ignore-port-conflicts")
}
}
cmd := exec.Command(args[0], args[1:]...)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}
// =============================================================================
// Install Worker
// =============================================================================
func (i *Command) installWorker(cryosparcpath, license, arch string) *exec.Cmd {
installScript := filepath.Join(
cryosparcpath,
"cryosparc_worker",
"install.sh",
)
var cmd *exec.Cmd
// Get local architecture
localArch, err := exec.Command("uname", "-m").Output()
if err != nil {
localArch = []byte("")
}
isLocal := strings.TrimSpace(string(localArch)) == arch
if isLocal {
// Run locally
cmd = exec.Command(
installScript,
"--license", license,
)
} else if arch == "aarch64" {
// Run on ARM nodes
cmd = exec.Command(
"srun",
"--cluster", "gmerlin7",
"--partition", "gh-interactive",
"--time", "0-00:30:00",
installScript,
"--license", license,
)
} else {
// Run on x86 nodes
cmd = exec.Command(
"srun",
"--cluster", "merlin7",
"--partition", "interactive",
"--reservation", "interactive",
"--time", "0-00:30:00",
installScript,
"--license", license,
)
}
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}
// =============================================================================
// Replace CryoSPARC license in config.sh
// =============================================================================
func (i *Command) replaceLicenseID(
_ func(float64),
installDir,
license string) error {
path := filepath.Join(installDir, "config.sh")
data, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("read config.sh %q: %w", path, err)
}
re := regexp.MustCompile(`(?m)^export CRYOSPARC_LICENSE_ID=".*"$`)
updated := re.ReplaceAllString(
string(data),
fmt.Sprintf(`export CRYOSPARC_LICENSE_ID="%s"`, license),
)
if err := os.WriteFile(path, []byte(updated), 0644); err != nil {
return fmt.Errorf("write config.sh %q: %w", path, err)
}
return nil
}
// =============================================================================
// Create new lane
// =============================================================================
type ClusterInfo struct {
Name string `json:"name"`
WorkerBinPath string `json:"worker_bin_path"`
CachePath string `json:"cache_path"`
SendCmdTpl string `json:"send_cmd_tpl"`
QsubCmdTpl string `json:"qsub_cmd_tpl"`
QstatCmdTpl string `json:"qstat_cmd_tpl"`
QdelCmdTpl string `json:"qdel_cmd_tpl"`
QinfoCmdTpl string `json:"qinfo_cmd_tpl"`
}
func (i *Command) CreateLane(
_ func(float64),
cryosparcpath,
name,
cachepath,
memory,
time,
partition,
gpus,
cpuspertask,
cluster,
arch string,
) error {
// Directory where the generated files will be placed
laneDir := filepath.Join(cryosparcpath, "lanes", name)
if err := os.MkdirAll(laneDir, 0755); err != nil {
return fmt.Errorf("create lane directory: %w", err)
}
// ====================================================================== //
// CryoSPARC installations before version 5 can only access the
// cryosparcw if the parent directory has the name 'cryosparc_worker'
// ====================================================================== //
workerDir := fmt.Sprintf("cryosparc_worker_%s", arch)
versionData, err := os.ReadFile(
filepath.Join(cryosparcpath, "cryosparc_master", "version"),
)
if err == nil {
version := strings.TrimPrefix(
strings.TrimSpace(string(versionData)), "v")
major, _, _ := strings.Cut(version, ".")
if majorVersion, err := strconv.Atoi(major); err == nil &&
majorVersion < 5 {
workerDir = "cryosparc_worker"
}
}
// ====================================================================== //
workerBinPath := filepath.Join(
cryosparcpath,
workerDir,
"bin",
"cryosparcw",
)
// ---------------------------------------------------------------------
// cluster_info.json
// ---------------------------------------------------------------------
info := ClusterInfo{
Name: name,
WorkerBinPath: workerBinPath,
CachePath: cachepath,
SendCmdTpl: "{{ command }}",
QsubCmdTpl: fmt.Sprintf(
"bash -c 'sbatch --parsable "+
"--cluster=%s \"{{ script_path_abs }}\" | cut -d \";\" -f 1'",
cluster,
),
QstatCmdTpl: fmt.Sprintf(
"squeue --cluster=%s -j {{ cluster_job_id }}",
cluster,
),
QdelCmdTpl: fmt.Sprintf(
"scancel --cluster=%s {{ cluster_job_id }}",
cluster,
),
QinfoCmdTpl: fmt.Sprintf(
"sinfo --cluster=%s",
cluster,
),
}
infoJSON, err := json.MarshalIndent(info, "", " ")
if err != nil {
return fmt.Errorf("marshal cluster info: %w", err)
}
infoPath := filepath.Join(laneDir, "cluster_info.json")
if err := os.WriteFile(infoPath, infoJSON, 0644); err != nil {
return fmt.Errorf("write %s: %w", infoPath, err)
}
// ---------------------------------------------------------------------
// cluster_script.sh
// ---------------------------------------------------------------------
var script string
// Set default cpus per task value for GPU nodes
if cpuspertask == "num_cpu" && strings.Contains(partition, "a100-") {
cpuspertask = "16 * num_gpu"
} else if cpuspertask == "num_cpu" && strings.Contains(partition, "gh-") {
cpuspertask = "72 * num_gpu"
}
if cluster == "gmerlin7" {
script = fmt.Sprintf(`#!/usr/bin/env bash
# cryoSPARC cluster submission script template for SLURM
# Lane: %s
{# This template uses jinja2 syntax. #}
# Available variables:
# script_path_abs={{ script_path_abs }}
# run_cmd={{ run_cmd }}
# num_cpu={{ num_cpu }}
# num_gpu={{ num_gpu }}
# ram_gb={{ ram_gb }}
# job_dir_abs={{ job_dir_abs }}
# project_dir_abs={{ project_dir_abs }}
# job_log_path_abs={{ job_log_path_abs }}
# worker_bin_path={{ worker_bin_path }}
# run_args={{ run_args }}
# project_uid={{ project_uid }}
# job_uid={{ job_uid }}
# job_creator={{ job_creator }}
# cryosparc_username={{ cryosparc_username }}
#SBATCH --job-name=cryosparc_{{ project_uid }}_{{ job_uid }}
#SBATCH --output={{ job_log_path_abs }}.out
#SBATCH --error={{ job_log_path_abs }}.err
#SBATCH --ntasks=1
#SBATCH --mem={{ (%s * 1000)|int }}M
#SBATCH --time=%s
#SBATCH --partition=%s
#SBATCH --cluster=%s
#SBATCH --gres=gpu:{{ %s }}
#SBATCH --cpus-per-task={{ %s }}
#SBATCH --hint=nomultithread
{%%- if num_gpu == 0 %%}
# Use CPU cluster
echo "Error: No GPU requested. Use a CPU lane instead." >&2
exit 1
{%%- else %%}
# Print hostname, for debugging
echo "Job Id: $SLURM_JOBID"
echo "Host: $SLURM_NODELIST"
module purge
srun {{ run_cmd }}
EXIT_CODE=$?
echo "Exit code: $EXIT_CODE"
exit $?
{%%- endif %%}
`, name, memory, time, partition, cluster, gpus, cpuspertask)
} else {
script = fmt.Sprintf(`#!/usr/bin/env bash
# cryoSPARC cluster submission script template for SLURM
# Lane: %s
{# This template uses jinja2 syntax. #}
# Available variables:
# script_path_abs={{ script_path_abs }}
# run_cmd={{ run_cmd }}
# num_cpu={{ num_cpu }}
# num_gpu={{ num_gpu }}
# ram_gb={{ ram_gb }}
# job_dir_abs={{ job_dir_abs }}
# project_dir_abs={{ project_dir_abs }}
# job_log_path_abs={{ job_log_path_abs }}
# worker_bin_path={{ worker_bin_path }}
# run_args={{ run_args }}
# project_uid={{ project_uid }}
# job_uid={{ job_uid }}
# job_creator={{ job_creator }}
# cryosparc_username={{ cryosparc_username }}
#SBATCH --job-name=cryosparc_{{ project_uid }}_{{ job_uid }}
#SBATCH --output={{ job_log_path_abs }}.out
#SBATCH --error={{ job_log_path_abs }}.err
#SBATCH --ntasks=1
#SBATCH --mem={{ (%s * 1000)|int }}M
#SBATCH --time=%s
#SBATCH --partition=%s
#SBATCH --cluster=%s
#SBATCH --cpus-per-task={{ %s }}
#SBATCH --hint=nomultithread
{%%- if num_gpu > 0 %%}
# Use GPU cluster
echo "Error: GPU requested. Use a GPU lane instead." >&2
exit 1
{%%- else %%}
# Print hostname, for debugging
echo "Job Id: $SLURM_JOBID"
echo "Host: $SLURM_NODELIST"
module purge
srun {{ run_cmd }}
EXIT_CODE=$?
echo "Exit code: $EXIT_CODE"
exit $?
{%%- endif %%}
`, name, memory, time, partition, cluster, cpuspertask)
}
scriptPath := filepath.Join(laneDir, "cluster_script.sh")
if err := os.WriteFile(scriptPath, []byte(script), 0755); err != nil {
return fmt.Errorf("write %s: %w", scriptPath, err)
}
return nil
}
// =============================================================================
// Set CryoSPARC base port
// =============================================================================
// setCryosparcBasePort will automatically set a base port for CryoSPARC based
// on the availabilty of a number contiguous ports on the server running
// CryoSPARC
func (i *Command) setCryosparcBasePort(
hostname string,
startport,
endport,
count uint) error {
remoteCmd := []string{
"findbaseport",
"--start", strconv.FormatUint(uint64(startport), 10),
"--end", strconv.FormatUint(uint64(endport), 10),
"--count", strconv.FormatUint(uint64(count), 10),
}
cmd := exec.Command(
"ssh",
fmt.Sprintf("%s@%s", os.Getenv("USER"), hostname),
strings.Join(remoteCmd, " "),
)
// Capture stdout instead of sending it directly to the screen.
output, err := cmd.Output()
if err != nil {
return err
}
// Remove whitespace/newline.
portStr := strings.TrimSpace(string(output))
// Convert the output to a number.
port, err := strconv.ParseUint(portStr, 10, 16)
if err != nil {
return fmt.Errorf(
"invalid port returned by findbaseport: %q: %w", portStr, err)
}
i.cfg.BasePort = uint(port)
return nil
}
// =============================================================================
// Check running CryoSPARC instances
// =============================================================================
func (i *Command) checkInstances(hostname string) *exec.Cmd {
remoteCmd := `uports --cmd "node dist/server/index.js"
| awk 'NR > 2 {print "http://` + hostname + `:" $1}'`
cmd := exec.Command(
"ssh",
fmt.Sprintf("%s@%s", os.Getenv("USER"), hostname),
remoteCmd,
)
cmd.Stdin = os.Stdin
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stderr
return cmd
}