Use slurm to turn on various apps

utils.sh

tensorboard_slurm() {
  local port="$1"
  local exp_path="$2"
  local session_tag="tensorboard_$(date +%s)"

  if [[ -z "$port" || -z "$exp_path" ]]; then
    echo "Usage: tensorboard_slurm <port> <path_to_experiment>"
    return 1
  fi

  tmux new-session -d -s tensorboard_session

  # Pane 0: Access the compute node via SLURM with a session tag
  tmux send-keys -t tensorboard_session:0 "srun --job-name=$session_tag --pty -N 1 -p debug /bin/bash -l" C-m
  sleep 2

  # Wait for SLURM allocation and get the compute node name
  while true; do
    compute_node=$(squeue -n "$session_tag" -o "%N" | tail -n 1)
    if [[ "$compute_node" != "NODELIST" && -n "$compute_node" ]]; then
      break
    fi
    echo "Waiting for allocation..."
    sleep 1
  done

  # Pane 1: Start TensorBoard with port forwarding
  tmux split-window -v -t tensorboard_session:0
  tmux send-keys -t tensorboard_session:0.1 "ssh -o StrictHostKeyChecking=no -L $port:localhost:$port $compute_node" C-m
  tmux send-keys -t tensorboard_session:0.1 "/home/kwanghec/StreamingDSU/envs/bin/tensorboard --logdir $exp_path --port $port" C-m

  # Attach to the session
  tmux attach-session -t tensorboard_session
}

jupyter_slurm() {
  local port="$1"
  local exp_path=$(pwd)
  local env_path="$2"
  local session_tag="jupyter_$(date +%s)"

  if [[ -z "$env_path" ]]; then
    if [[ -d "$(pwd)/envs" ]]; then
      env_path=$(pwd)/envs
    fi
  fi

  if [[ -z "$port" || -z "$env_path" ]]; then
    echo "Usage: jupyter_slurm <port> <path_to_jupyter>"
    return 1
  fi

  tmux new-session -d -s jupyter_session

  # Pane 0: Access the compute node via SLURM with a session tag
  tmux send-keys -t jupyter_session:0 "srun --job-name=$session_tag --pty -N 1 --mem 8G -p debug /bin/bash -l" C-m
  sleep 2

  # Wait for SLURM allocation and get the compute node name
  while true; do
    compute_node=$(squeue -n "$session_tag" -o "%N" | tail -n 1)
    if [[ "$compute_node" != "NODELIST" && -n "$compute_node" ]]; then
      break
    fi
    echo "Waiting for allocation..."
    sleep 1
  done

  # Pane 1: Start TensorBoard with port forwarding
  tmux split-window -v -t jupyter_session:0
  tmux send-keys -t jupyter_session:0.1 "ssh -o StrictHostKeyChecking=no -L $port:localhost:$port $compute_node" C-m
  tmux send-keys -t jupyter_session:0.1 "$env_path/bin/jupyter lab --notebook-dir=$exp_path --port=$port --ip=0.0.0.0 --no-browser" C-m

  # Attach to the session
  tmux attach-session -t jupyter_session
}