1006 lines
22 KiB
Go
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
|
|
}
|