#!/usr/bin/env python

"""
 h2jupynb - Hoffman2 Jupyter launcher
 
 
 This script allows users to start a jupyter notebooks/labs session
running on the UCLA Hoffman2 Cluster (Slurm, hoffman2.oarc.ucla.edu)
 while displaying it on their local browser.

 Users download this script on their local machine and run it on a local
 terminal via:

 python3 h2jupynb -u $H2-USERNAME [other options]
 
 ***The script prompts the users twice for their Hoffman2 password***

 Jupyter is installed in the python/3.12.13 module (the default). With other
 python versions, or to use your own packages in the notebook, install jupyter
 in your own python virtual environment (-k) or conda environment (-j), or use
 a container that has it (-f).


 To see a full list of current available options:

 python3 h2jupynb --help

 See also:

 https://www.hoffman2.idre.ucla.edu/Using-H2/Connecting/Connecting.html#jupyter-notebook-lab-with-h2jupynb

 For questions open a ticket on:

 https://support.idre.ucla.edu/helpdesk/Tickets/New?categoryId=390006
 
 or file an issue or discussion on:

 https://github.com/rdauria/jupyter-notebook

 Author: RD (hpc@ucla.edu) / Adapted from: https://github.com/pyHPC/ipynbhpc

 Updates includes robustness fixes for Unicode decoding and Ctrl-C handling on Windows PowerShell

"""

# --- imports ---
from __future__ import print_function
from subprocess import Popen, PIPE, call
import subprocess
import sys
import ctypes
import webbrowser
from getopt import getopt
import time
import os
import platform
import random
from datetime import datetime
import sys, getopt
import signal
import time
import tempfile
import shutil
import atexit
from threading import Thread

__version__="4.3"
"""
 Changelog:
 Version: 4.3
 Date: 2026-10-20
  Changes:
   - -e distributed: -s requests tasks that may span nodes (-N <nodes> to fix the
     number of nodes), for parallel work launched with srun from the notebook.
     The session is an interactive step (salloc + srun --interactive) that holds no CPUs
 Version: 4.2
 Date: 2026-10-19
  Changes:
   - Removed -e (parallel environment): Slurm has none; -s requests CPU cores on one node
   - Jupyter is in python/3.12.13; otherwise users install it in their own environment
   - SIGTERM/SIGHUP (e.g. closing the terminal) clean up like Ctrl-C
 Version: 4.1
 Date: 2026-10-19
  Changes:
   - One SSH connection (ControlMaster) for the session and the tunnel, so the
     password is typed only once; -n no (and Windows) uses two connections as before
   - The tunnel listens only on localhost (dropped ssh -g)
 Version: 4.0
 Date: 2026-10-19
  Changes:
   - Ported to the Slurm cluster (hoffman2.oarc.ucla.edu): qrsh -> srun --pty
   - Added -q <QOS> (highp, pi_shared24, campus24) and -A <account>; -o yes = --qos=highp
   - CPU/GPU types, GPU memory, python, conda and cuda versions updated to the Slurm cluster
   - CPU type, GPU card and GPU memory are requested as Slurm features (-C)
   - Checks that the session is inside a Slurm job and that jupyter is available
   - Apptainer containers get --nv on GPU nodes
 Version: 3.6
 Date: 2026-03-25
  Changes:
   - Added modular help system (--help-gpu, --help-apptainer, etc.)
   - Updated code for better flow
   - Added Apptainer container support via -f
"""

# Inline Countdown so this distributable is self-contained (no external import).
class Countdown(Thread):
  """Background countdown timer that prints remaining time to the terminal.
  """
  def __init__(self, total_seconds, interval=10):
    super(Countdown, self).__init__()
    self.total = int(total_seconds)
    self.interval = max(1, int(interval))
    try:
      # Python 3/2 compatible Event
      from threading import Event
      self._stop_event = Event()
    except Exception:
      # fallback: simple flag
      self._stop_event = None
      self._stopped_flag = False
    # make daemon-compatible across versions
    try:
      self.daemon = True
    except Exception:
      pass
    self._last_len = 0

  def stop(self):
    if self._stop_event is not None:
      self._stop_event.set()
    else:
      self._stopped_flag = True

  def stopped(self):
    if self._stop_event is not None:
      return self._stop_event.is_set()
    return getattr(self, '_stopped_flag', False)

  def _format(self, secs):
    h = secs // 3600
    m = (secs % 3600) // 60
    s = secs % 60
    return "{:02d}:{:02d}:{:02d}".format(h, m, s)

  def _print_line(self, text):
    try:
      #sys.stdout.write('\r' + text + ' ' * max(0, self._last_len - len(text)))
      sys.stdout.write('\r' + text + ' ' * 2)  # pad extra spaces
      #sys.stdout.write('\r' + ' ' * self._last_len + '\r')
      sys.stdout.flush()
    except Exception:
      pass
    self._last_len = len(text)

  def run(self):
    remaining = self.total
    self._print_line("Time remaining: {}".format(self._format(remaining)))
    while remaining > 0:
      # wait with interval, but allow stop flag to be checked
      slept = 0
      while slept < self.interval:
        if self.stopped():
          break
        time.sleep(1)
        slept += 1
      if self.stopped():
        break
      remaining -= self.interval
      if remaining < 0:
        remaining = 0
      self._print_line("Time remaining: {}".format(self._format(remaining)))
    try:
      sys.stdout.write('\n')
      sys.stdout.flush()
    except Exception:
      pass

random.seed(datetime.now().strftime("%S"))

# Provide a small compatibility layer so the script can run under Python2 or Python3.
is_py2 = sys.version_info[0] == 2
# Ensure `xrange` is available on both versions (Python2 has it; Python3 will alias it to range)
if not is_py2:
  xrange = range


def _send_ctrl_break(p):
  """Try to interrupt a child process group on Windows using
  GenerateConsoleCtrlEvent; otherwise send SIGINT or kill.
  """
  try:
    if opsys == 'Windows':
      try:
        ctypes.windll.kernel32.GenerateConsoleCtrlEvent(1, int(p.pid))
        time.sleep(0.2)
      except Exception:
        try:
          p.kill()
        except Exception:
          pass
  except Exception:
    pass

# --- defaults ---
username = ''
username = str(username)
timeinhours = 2
timeinhours = int(timeinhours)
memoryingb = 5
memoryingb = int(memoryingb) 
numberofslots = 1
numberofslots = int(numberofslots)
parenv = 'shared'
numberofnodes = 0
port = random.randrange(6222, 10000, 2)
port = int(port)
directory = '~'
directory = str(directory)
opsys = platform.system()
opsys = str(opsys)
pythonver = 'python/3.12.13'
pythonver = str(pythonver)
usegpu = 'no'
gpu = ''
highp = 'no'
cudaver = '12.8'
exclusive = 'no'
arch = ''
mods = ''
lofmods = ''
modstoload = ''
cve = ''
pve = ''
#gputype = ''
gpumem = ''
sshlogs= 'no'
sharessh = 'yes'
gccver= ''
ptf = ''
hostname = ''
httpcommand = '' 
jobid = ''
account = ''
qos = ''

# --- Hoffman2 Cluster (Slurm) ---
H2_HOST = 'hoffman2.oarc.ucla.edu'
PYTHON_VERSIONS = ['3.10.20', '3.12.13']
CONDA_VERSIONS = ['miniforge', 'miniforge/24.3.0-0', 'miniforge/26.5.3-0']
CUDA_VERSIONS = ['11.7', '11.8', '12.2', '12.4', '12.8', '13.0', '13.2']
QOS_LIST = ['highp', 'pi_shared24', 'campus24']
# GPU card features and their GPU memory in GB (see: h2 gpus)
GPU_CARDS = {'P4': 8, 'K40m': 12, 'A100': 40, 'RTXPRO6000Bw': 97}
# CPU type features (see: sinfo -o "%f")
CPU_TYPES = ['amd-7542', 'amd-epyc-9554', 'amd-epyc-9655',
             'intel-E5-2650', 'intel-E5-2650v4', 'intel-E5-2670', 'intel-E5-2697Av4',
             'intel-gold-6140', 'intel-gold-6240', 'intel-gold-6342', 'intel-gold-6442Y',
             'intel-xeon6-6736P']


#print(f"h2jupynb version {__version__}")
print("h2jupynb version {}".format(__version__))
#print("h2jupynb version %s" % __version__)
print()

# Detect Windows robustly and show a short tip for PowerShell users
is_windows = (os.name == 'nt') or sys.platform.startswith('win') or opsys.lower().startswith('win')
if is_windows:
  print("Note for Windows users: if Ctrl-C doesn't stop the script while the countdown is running, try Ctrl+Break (Pause/Break) or Ctrl+Fn+B on some laptops.")

# Convenience flag for showing Windows-specific help and exiting early.
if '--windows-help' in sys.argv:
  print("Windows interruption tip: if Ctrl-C doesn't interrupt during the countdown, try Ctrl+Break (Pause/Break) or Ctrl+Fn+B. The script will also try to send a CTRL_BREAK_EVENT to child processes on Windows.")
  sys.exit(0)

# Exit cleanly (closing the ssh connection and removing the socket directory)
# when terminated or when the terminal is closed, as for Ctrl-C
def _exit_on_signal(signum, frame):
  sys.exit(1)
if not is_windows:
  signal.signal(signal.SIGTERM, _exit_on_signal)
  signal.signal(signal.SIGHUP, _exit_on_signal)

if sys.version_info<(2,6,0):
  sys.stderr.write("You need python 2.6 or later to run this script\n")
  exit(1)


if opsys == 'Windows':
  try:
    ret=call(['where','ssh.exe'])
    #mytest = 'MINGW' in os.environ['MSYSTEM']
  except:
    print("It appears that Git for Windows is not installed")
    print("or that you are not running this script from its")
    print("Git BASH terminal. Please install either Git for")
    print("Windows, availabe at: ")
    print("https://git-for-windows.github.io/")
    print("and run this script from its git BASH.")
    sys.exit(2)


# --- IO helpers ---

#def handler(signum, frame):
#  print("\nDetected Ctrl-c. Exiting...")
#  try:
#    _send_ctrl_break(pqsub)
#    _send_ctrl_break(ptunnel)
#  except Exception:
#    # pqsub or ptunnel may not be defined yet (need to add them to global in main!
#    pass
#  print("Succesfully cleaned up connections.")
#  #exit(1)
#  raise SystemExit

def gethostname(line):
    """
    Check a line from the remote output and extract the hostname.
    Returns True if hostname was found, False otherwise.
    """
    global hostname, jobid
    line = line.strip()
    # skip the terminal's echo of the srun command line ($(hostname), `hostname`)
    if line.startswith('HOSTNAME') and ('$' not in line) and ('`' not in line):
        # line format: HOSTNAME=<host> SLURM_JOB_ID=<id>
        fields = dict(f.partition('=')[::2] for f in line.split())
        hostname = fields.get('HOSTNAME', '').strip()
        jobid = fields.get('SLURM_JOB_ID', '').strip()
        return True
    return False

def gethttpcommand(line):
    global httpcommand
    if (line.find('http://%s:%s/'%(hostname,port))>0):
       httpcommand = line.split(' ')[-1].strip()
       return True
    if line.startswith('exitcode:'):
       # jupyter exited (or was not found) before printing its URL
       return True

def readwhile(stream,func):
  # Choose bytes type depending on Python version (Py2: str, Py3: bytes)
  bytes_type = str if is_py2 else bytes
  while True:
    line = stream.readline()
    # Support both bytes and str from subprocess streams. Decode bytes
    # using UTF-8 and replace undecodable bytes so we never raise a
    # UnicodeDecodeError while reading remote output.
    if isinstance(line, bytes_type):
      try:
        sline = line.decode('utf-8', errors='replace')
      except AttributeError:
        sline = line
    else:
      sline = line

    if sline != '':
      if sline.endswith('\n'):
        print(sline[:-1])
      else:
        print(sline)
      if func(sline):
        break
    else:
      raise Exception("Disconnected unexpectedly.")

    
def readwhile_threaded(stream, func):
  """Run readwhile(stream, func) in a background thread and allow the
  main thread to handle KeyboardInterrupt (Ctrl-C) on Windows PowerShell.

  This helper joins the reader thread in a loop so KeyboardInterrupt can be
  received and cause immediate cleanup.
  """
  exc = []

  def target():
    try:
      readwhile(stream, func)
    except Exception as e:
      exc.append(e)

  t = Thread(target=target)
  # Python 2's Thread ctor doesn't accept `daemon=`; set attribute for both versions
  try:
    t.daemon = True
  except Exception:
    pass
  t.start()
  try:
    while t.is_alive():
      t.join(timeout=0.1)
  except KeyboardInterrupt:
    # Best-effort cleanup here; pqsub/ptunnel may or may not be defined
    try:
      _send_ctrl_break(pqsub)
      _send_ctrl_break(ptunnel)
    except Exception:
      pass
    sys.exit(1)

  if exc:
    # Re-raise the first exception from the reader thread
    raise exc[0]

    
# --- SSH connection sharing (ControlMaster) ---
def make_control_path():
  """Return a path for the ControlMaster socket in a new private (0700)
  directory, or '' if connection sharing cannot be used. The directory is
  removed when the script exits.
  """
  try:
    cdir = tempfile.mkdtemp(prefix='h2j-')
    # Unix socket paths are limited to 104 (macOS) or 108 (Linux) bytes
    if len(cdir) > 90:
      shutil.rmtree(cdir, ignore_errors=True)
      cdir = tempfile.mkdtemp(prefix='h2j-', dir='/tmp')
  except Exception:
    return ''
  atexit.register(shutil.rmtree, cdir, True)
  return os.path.join(cdir, 'cm')

def ssh_control(cpath, op, extra=None):
  """Send a control command (-O check, forward, exit) to the master
  connection. Returns True on success; gives up after 10 seconds.
  """
  args = ['ssh', '-o', 'ControlPath=%s' % cpath] + (extra or []) + ['-O', op, '%s@%s' % (username, H2_HOST)]
  try:
    devnull = open(os.devnull, 'w')
    p = Popen(args, stdin=devnull, stdout=devnull, stderr=PIPE)
    for i in xrange(100):
      if p.poll() is not None:
        break
      time.sleep(0.1)
    if p.poll() is None:
      p.kill()
      return False
    err = p.stderr.read()
    if p.returncode != 0 and err:
      print(err.decode('utf-8', 'replace') if isinstance(err, bytes) else err)
    return p.returncode == 0
  except Exception:
    return False

# --- helper functions ---
def print_usage():
  print("Usage:")
  print("h2jupynb [-u <Hoffman2 user name>] [-v <python-version>] \n         [-t <time, integer number of hours>] [-m <memory, integer number of GB per core>] \n         [-s <number of CPU cores, on one node; tasks with -e distributed>] \n         [-e <distributed: tasks may span nodes>] [-N <number of nodes, with -e distributed>] \n         [-q <QOS: highp, pi_shared24 or campus24>] [-A <Slurm account>] \n         [-o <run on group owned nodes (= -q highp): yes/no>] [-x <run on an exclusively reserved node: yes/no>]  \n         [-a <CPU-type>] [-d <path to directory from which the jupyter NB/lab will start>] \n         [-g <run on a gpu node: yes/no>] [-c <gpu-card-type>] [-r <gpu-global-memory>] \n         [-l <cuda-version>]  [-p <port number>] [-j <conda virtual environment name>]\n         [-k <path/to/python/virtual-env>] \n         [-b <comma separated list of environmental modules>]\n         [-f </path/to/sif-file.sif> (only if python version = apptainer)]\n         [-z <write ssh debug files?: yes/no>] [-n <one shared ssh connection: yes/no>]\n         Use -h or --help for full details.")

def print_help():
    print_usage()
    print(
        "\nUse one of the following for detailed help:\n"
        "  --help_defaults     Default values\n"
        "  --help_python       Python versions\n"
        "  --help_apptainer    Apptainer usage\n"
        "  --help_envs         Environments (modules, conda, venv, directory)\n"
        "  --help_cpu          CPU options\n"
        "  --help_gpu          GPU options\n"
    )
    print("=== Notes ===\n")
    print("\t- Invalid parameters fall back to defaults")
    print("\t- Some combinations may lead to long wait times")
    print("\t- Your accounts, QOS and limits: h2 whoami (on the cluster)")
    print("\t- CPU overview: sinfo -o \"%20P %10c %10m %f\" (on the cluster)")
    print("\t- GPU overview: h2 gpus (on the cluster)\n")

    
def help_defaults():
    #print_usage()
    print("\n\t=== Defaults (when no arguments are provided) ===\n")
    print("\tUser name: same as client machine")
    print("\tTime: 2 hours")
    print("\tMemory per CPU core: 5 GB")
    print("\tCPU cores: 1, all on one node (no -e)")
    print("\tPython version: 3.12.13")
    print("\tPort: random, between 6222 and 10000")
    print("\tWorking directory: $HOME")
    print("\tGPU: no")
    print("\tCUDA version: 12.8 (if -g y)")
    print("\tExclusive node: no")
    print("\tQOS: cluster default (pi_shared24 for groups with nodes in Slurm, campus24 otherwise)")
    print("\tAccount: your default Slurm account (required with -A if you have more than one)")
    print("\tCPU: not specified")
    print("\tShared ssh connection (password typed once): yes, except on Windows\n")

def help_cpu():
    print_usage()
    print("\n\t=== CPU Help===\n")
    print("\tUse -a <arch-type> to select specific architecture (a Slurm feature, -C)")
    print("\tCPU types: " + ", ".join(CPU_TYPES))
    print("\tFamilies: -a intel, -a intel-gold or -a amd request any CPU of that family")
    print("\n\tCores: by default -s <N> requests N CPU cores on one node (--cpus-per-task)")
    print("\tWith -e distributed, -s <N> requests N tasks (--ntasks) that may be placed on")
    print("\tseveral nodes; add -N <nodes> to set the number of nodes. The notebook runs on the")
    print("\tfirst node: start parallel work with srun (e.g. srun ./my_mpi_program) from the")
    print("\tnotebook or its terminal. Memory (-m) is per core in both cases.")
    print("\n\tTo see available CPU types and their memory connect to the cluster")
    print("\tand issue the command: sinfo -o \"%20P %10c %10m %f\"\n")

def help_gpu():
    print_usage()
    print("\n\t=== GPU / CUDA Help===\n")
    print("\tUse -g yes to enable GPU")
    print("\tCUDA (-l): " + ", ".join(CUDA_VERSIONS) + " -- default is 12.8")
    print("\t  (K40m needs cuda 11.x; P4 needs cuda 12.x or earlier)")
    print("\tGPU cards (-c): " + ", ".join(sorted(GPU_CARDS)))
    print("\tGPU memory currently available are (in GB, -r): " +
          ", ".join("%d (%s)" % (m, c) for c, m in sorted(GPU_CARDS.items(), key=lambda x: x[1])))
    print("\tTo see available GPU cards and their memory connect to the cluster")
    print("\tand issue the command: h2 gpus\n")
    
def help_python():
    print_usage()
    print("\n\tUse -v <version> to select version of python")
    print("\t=== Available Python versions ===\n")
    print("\t" + ", ".join(PYTHON_VERSIONS))
    print("\t" + ", ".join(CONDA_VERSIONS) + ", apptainer\n")
    print("\te.g., version 3.12.13 loads python/3.12.13")
    print("\te.g., miniforge loads the default miniforge module (for conda envs, -j)")
    print("\tanaconda3/* and mamba/* are not available on the Slurm cluster: they are mapped to miniforge\n")
    print("\tJupyter is installed in python/3.12.13 (the default). With other versions, or to use")
    print("\tyour own packages in the notebook, install jupyter in your own virtual environment, e.g.:")
    print("\t  module load python/3.12.13; python3 -m venv $HOME/.virtualenvs/myenv")
    print("\t  $HOME/.virtualenvs/myenv/bin/pip install notebook <your packages>")
    print("\tand run with: -v 3.12.13 -k $HOME/.virtualenvs/myenv\n")

def help_envs():
    print_usage()
    print("\n\t=== Environments ===\n")
    print("\tModules: -b <comma-separated list> (e.g. -b gsl/2.6,curl/7.70.0)")
    print("\tConda env: -j <conda env name>, requires -v <miniforge version> (e.g. -v miniforge -j treemix)")
    print("\tVirtualenv: -k </PATH/TO/MYPYTHONENV>, requires -v <python version> (e.g. -v 3.12.13 -k .virtenvs/3.12.13/myenv)\n")
    print("\tInstall jupyter in the environment (pip install notebook, or conda install notebook): otherwise the")
    print("\tjupyter of the python/3.12.13 module, if loaded, runs and the notebook cannot use the environment's packages\n")
    print("\tRoot directory of jupyter session (defaults to $HOME): -d </PATH/TO/DIR>") 
    print("\t"r"(e.g., -d r$HOME/PATH/TO/DIR, or -d \$SCRATCH/MYDIR)")
    print()

def help_apptainer():
    print_usage()
    print("\n\t=== Apptainer ===\n")
    print("\tIf python version is set to 'apptainer' (-v apptainer), provide path to .sif container")
    print("\tThe container must have jupyter installed; on GPU nodes it is started with --nv")
    print("\t"r"e.g.: -f \$HOME/container.sif or -f \$H2_CONTAINER_LOC/rosetta-ml420-jupyter.sif")
    print()


def main(argv):
  # ------------------------------
  # Globals for options/args only
  # ------------------------------
  global  username, timeinhours, memoryingb, numberofslots
  global  pythonver, port, directory, usegpu, gpu, gpumem
  global  cudaver, highp, exclusive, arch, mods, parenv, numberofnodes
  global  cve, pve, lofmods, modstoload, sshlogs, gccver, ptf
  global  hostname, httpcommand, jobid, account, qos, sharessh

  # ------------------------------
  # Local variables
  # ------------------------------
  pqsub = None
  ptunnel = None
  
  # ------------------------------
  # Parse arguments
  # ------------------------------
  try:
    opts, args = getopt.getopt(
        sys.argv[1:],"hu:t:m:s:e:N:v:p:d:g:c:r:l:o:x:a:b:j:k:f:z:q:A:n:",
        [
                "help",
                "help_defaults",
                "help_python",
                "help_gpu",
                "help_cpu",
                "help_envs",
                "help_apptainer",
                "username=",
                "timeinhours=",
                "memoryingn=",
                "numberofslots=",
                "pythonver=",
                "port=",
                "dir=",
                "usegpu=",
                "gpu_card=",
                "gpu_mem",
                "cudaver=",
                "highp=",
                "exclusive=",
                "arch=",
                "modules=",
                "conda-virt-env=",
                "python-venv-path=",
                "sif-file-path=",
                "ssh-debug-files=",
                "qos=",
                "account=",
                "share-ssh=",
                "parallel=",
                "nodes="
        ]
    )

  except getopt.GetoptError:
    print_usage() 
    sys.exit(2)
  
  # ------------------------------
  # Process variables
  # ------------------------------
  for opt, arg in opts:
    #print(opt, arg)
    if opt in ("-h", "--help"):
      print_help()
      sys.exit(0)
    elif opt in ("--help_defaults"):
      help_defaults()
      sys.exit(0)
    elif opt in ("--help_python"):
      help_python()
      sys.exit(0)
    elif opt in ("--help_gpu"):
      help_gpu()
      sys.exit(0)
    elif opt in ("--help_cpu"):
      help_cpu()
      sys.exit(0)
    elif opt in ("--help_envs"):
      help_envs()
      sys.exit(0)
    elif opt in ("--help_apptainer"):
      help_apptainer()
      sys.exit(0)
    elif opt in ("-u", "--user"):
      username = arg
      username = str(username)
    elif opt in ("-w", "--what"):
      a_command = arg
      a_command = str(a_command)
      print("A_command= ",a_command)
    elif opt in ("-t", "--time"):
      timeinhours = arg
      try:
        timeinhours=int(timeinhours)
      except ValueError:
        print("The time entered is not in a valid format. Setting the time to 2 hours...")
        timeinhours = 2
    elif opt in ("-m", "--memory"):
      memoryingb = arg
      try:
        memoryingb=int(memoryingb)
      except ValueError:
        print("The memory entered is not in a valid format. Setting the memory to 5GB per computing core...")
        memoryingb = 5
    elif opt in ("-s", "--slots"):
      numberofslots = arg
      try:
        numberofslots=int(numberofslots)
      except ValueError:
        print("The requested number of CPU cores is invalid. Setting the number of CPU cores to 1...")
        numberofslots = 1
    elif opt in ("-e", "--parallel"):
      if str(arg).lower() in ('distributed', 'dist', 'd', '2'):
        parenv = 'distributed'
      else:
        if str(arg).lower() not in ('shared', 's', '1'):
          print("-e", arg, "is not valid (use: -e distributed). CPU cores will be on one node...")
        parenv = 'shared'
    elif opt in ("-N", "--nodes"):
      try:
        numberofnodes=int(arg)
      except ValueError:
        print("The requested number of nodes is invalid. Slurm will choose the number of nodes...")
        numberofnodes = 0
    elif opt in ("-v", "--version"):
      pythonver = arg
      pythonver = str(pythonver)
      if pythonver.startswith('anaconda3') or pythonver.startswith('mamba'):
        print("version, ", pythonver,", is not available on the Slurm cluster")
        print("setting python version to miniforge")
        pythonver = 'miniforge'
      if (pythonver in PYTHON_VERSIONS):
        pythonver = "python/"+pythonver
      elif (pythonver not in CONDA_VERSIONS) and (pythonver != 'apptainer'):
        print("version, ", pythonver,", of python not available")
        print("setting python version to 3.12.13")
        pythonver = 'python/3.12.13'
    elif opt in ("-p", "--port"):
      port = arg
      try:
        port=int(port)
      except ValueError:
        print("The network port requested is invalid. Setting the port to 8789...")
        port = 8789
    elif opt in ("-d", "--dir"):
      directory = arg
      directory = str(directory)
    elif opt in ("-g", "--usegpu"):
      usegpu = arg
      usegpu = str(usegpu)
      if (usegpu == 'Yes') or (usegpu == 'yes') or (usegpu == 'YES') or (usegpu == 'Y') or (usegpu == 'y'):
        #print("Jupyter notebook or lab will run on a GPU node")
        usegpu='yes'
      else:
        #print("TO RUN ON A GPU NODE USE: -g yes ")
        usegpu = 'no'
        #print("Jupyter notebook or lab will NOT run on a GPU node")
    elif opt in ("-c", "--gpu_card"):
      gpu = arg
      gpu = str(gpu)
      #print(f"GPU = {gpu}")
      if (gpu == 'K40'):
        gpu = 'K40m'
      if (gpu not in GPU_CARDS):
        print("version, ",gpu,", of GPU not available")
        print("setting GPU card to ANY")
        gpu = ''
      else:
        usegpu='yes'
    elif opt in ("-r", "--gpu_mem"):
      gpumem = arg
      try:
        gpumem=int(gpumem)
      except ValueError:
        print("GPU global memory should be entered as an integer number... ")
      if (gpumem in GPU_CARDS.values()):
        print("GPU global memory = ",gpumem,"GB available for card", [c for c in GPU_CARDS if GPU_CARDS[c] == gpumem][0])
        usegpu='yes'
        gpumem=str(gpumem)
      else:
        print("GPU global memory = ",gpumem,"GB not available")
        print("setting GPU global memory version to ANY")
        print("YOU SHOULD ENTER ONLY A NUMERIC VALUE")
        gpumem = ''
    elif opt in ("-l", "--cuda"):
      cudaver = arg
      cudaver = str(cudaver)
      if (cudaver not in CUDA_VERSIONS):
        print("version, ",cudaver,", not available")
        print("setting cuda version to 12.8")
        cudaver = '12.8'
    elif opt in ("-o", "--highp"):
      highp = arg
      highp = str(highp)
      #print("HIGHP= ",highp)
      if ( highp == 'NO') or (highp == 'No') or (highp == 'no') or (highp == 'nO') or (highp == 'N') or (highp == 'n'): 
        print("HIGHP= ",highp," Jupyter notebook or lab will *not* run on owned node(s)")
        highp = 'no'
      elif ( highp == 'Yes') or ( highp == 'yes') or ( highp == 'YES') or ( highp == 'Y') or ( highp == 'y'):
        print("HIGHP= ",highp," Jupyter notebook or lab will run on owned node(s)")
        highp='yes'
        qos='highp'
    elif opt in ("-x", "--exclusive"):
      exclusive = arg
      #print("EXCL= ",exclusive)
      exclusive = str(exclusive)
      #print("EXCL= ",exclusive)
      if ( exclusive == 'NO') or (exclusive == 'No') or (exclusive == 'no') or (exclusive == 'nO') or (exclusive == 'N') or (exclusive == 'n'):
        print("Jupyter notebook or lab will not run in exclusive mode")
        exclusive = 'no'
      elif (exclusive == 'Yes') or (exclusive == 'yes') or (exclusive == 'YES') or (exclusive == 'Y') or (exclusive == 'y'):
        print("Jupyter notebook or lab will run in exclusive mode")
        exclusive='yes'
      else:
        print("EXCLUSIVE=",exclusive)
    elif opt in ("-a", "--arch"):
      arch = arg
      arch = str(arch).replace('\\', '').rstrip('*')
      if (arch in ('intel', 'intel-gold', 'amd')):
        # any CPU of the family: OR of the matching features
        arch = '|'.join([a for a in CPU_TYPES if a.startswith(arch)])
      elif (arch not in CPU_TYPES):
        print("version, ",arch,", of CPU not available")
        print("no CPU-type will be requested")
        arch = ''
    elif opt in ("-b", "--modules"):
      mods = arg
      mods = str(mods)
      lofmods=mods.split(",")
      for l in lofmods:
        if ( modstoload == '' ):
          modstoload="module load " + l
        else:
          modstoload=modstoload + "; module load " + l
    elif opt in ("-j", "--conda-virt-env"):
      cve = arg
      cve = str(cve)
      print("cve=",cve)
    elif opt in ("-k", "--python-venv-path"):
      pve = arg
      pve = str(pve)
      print("pve=",pve)
    elif opt in ("-f", "--sif-file-path"):
      ptf = arg
      ptf = str(ptf)
      print("ptf=",ptf)
    elif opt in ("-q", "--qos"):
      qos = arg
      qos = str(qos)
      if (qos not in QOS_LIST):
        print("QOS, ",qos,", not available. Choose one of: " + ", ".join(QOS_LIST))
        print("the cluster default QOS will be used")
        qos = ''
      highp = 'yes' if qos == 'highp' else 'no'
    elif opt in ("-A", "--account"):
      account = arg
      account = str(account)
    elif opt in ("-n", "--share-ssh"):
      sharessh = 'no' if str(arg).lower() in ('no', 'n') else 'yes'
    elif opt in ("-z", "--ssh-debug-files"):
      sshlogs = arg
      sshlogs = str(sshlogs)
      if ( sshlogs == 'NO') or (sshlogs == 'No') or (sshlogs == 'no') or (sshlogs == 'nO') or (sshlogs == 'N') or (sshlogs == 'n'): 
        print("SSHLOGS= ",sshlogs," SSH logfiles will not be written.")
        sshlogs = 'no'
      elif ( sshlogs == 'Yes') or ( sshlogs == 'yes') or ( sshlogs == 'YES') or ( sshlogs == 'Y') or ( sshlogs == 'y'):
        print("SSHLOGS= ",sshlogs," SSH logfiles will be written.")
        sshlogs='yes'
    else:
      print_usage()
      sys.exit(1)
  
  # ------------------------------
  # Windows bits (probably not needed)
  # ------------------------------    
  if not username:
    if opsys == 'Windows':
      username = str(os.environ["USERNAME"])
    else:
      username = str(os.environ["USER"])

  # ------------------------------
  # Print out summary/arguments
  # ------------------------------
  print("") 
  print("The requested Hoffman2 Cluster user name is:", username )
  print("The requested runtime is:", timeinhours,"hours" )
  print("The requested memory per CPU core is:", memoryingb,"GB")
  if (numberofnodes > 0) and (parenv != 'distributed'):
    print("-N is used only with -e distributed: the CPU cores will be on one node")
    numberofnodes = 0
  if (parenv == 'distributed'):
    print("The number of tasks (may span nodes) is:", numberofslots)
    if (numberofnodes > 0):
      print("The number of nodes is:", numberofnodes)
    print("Start parallel work with srun from the notebook or its terminal (the notebook runs on the first node)")
  elif (numberofslots > 1):
    print("The number of CPU cores is:", numberofslots)
  print("The version of python for the notebook is:", pythonver)
  if (qos != ''):
    print("The requested QOS is:", qos)
  if (account != ''):
    print("The requested Slurm account is:", account)
  if (highp == 'yes'):
    print("Your Jupyter NB or lab will run on owned resources")
  # K40m (compute capability 3.5) is not supported by cuda >= 12, P4 (6.1) by cuda >= 13
  if (gpu == 'K40m' or gpumem == '12') and not cudaver.startswith('11.'):
    print("cuda version",cudaver,"does not support K40m cards: setting cuda version to 11.8")
    cudaver = '11.8'
  if (gpu == 'P4' or gpumem == '8') and cudaver.startswith('13.'):
    print("cuda version",cudaver,"does not support P4 cards: setting cuda version to 12.8")
    cudaver = '12.8'
  if (usegpu == 'yes'):
    if (gpu=='') and (gpumem == ''):
      print("Your Jupyter NB or lab will run on a GPU node and load cuda version",cudaver)
    else:
      if (gpumem == ''):
        print("Your Jupyter NB or lab will run on GPU card ",gpu," and load cuda version",cudaver)
      else:
        print("Your Jupyter NB or lab will run on GPU card ",gpu," with ",gpumem,"GB of global GPU memory and load cuda version",cudaver)
  if (port != ''):
    print("Your Jupyter NB or lab will run on port ",port)
  if (exclusive == 'yes'):
    print("Your Jupyter NB or lab will reserve a node exclusively ")
  if (directory == '~'):
    print("The directory on Hoffman2 is $HOME")
  else:
    print("The directory on Hoffman2 is", directory)
  if (arch != ''):
    print("Your Jupyter NB or lab will run on ",arch," CPU")
  if (mods != ''):
    print("Your Jupyter NB or lab will load these environmental modules ",mods)
  if (cve != ''):
    print("Your Jupyter NB or lab will load this conda virtual environment ",cve)  
  if (pve != ''):
    print("Your Jupyter NB or lab will load this python virtual environment ",pve)
  if (ptf != ''):
    print("Your Jupyter NB will run off the apptainer container ",ptf)    
  if (sshlogs != ''):
    print("SSH log files will be written ",sshlogs)

  #### moved here
  print("usegpu= ",usegpu,"; highp= ",highp,"; qos= ",qos,"; account= ",account,"; exclusive= ",exclusive, "; gputype= ",gpu, "; gpu_mem= ",gpumem, "; arch= ",arch, "; time= ",timeinhours, "; mem= ",memoryingb, "; PE= ", parenv, "; NCORES= ",numberofslots, "; NNODES= ",numberofnodes,"sshlogs",sshlogs)

  # --- runtime limits of the QOS (no QOS = pi_shared24 or campus24, 24 hours) ---
  maxhours = 14*24 if (qos == 'highp') else 24
  if (timeinhours > maxhours):
    print("The requested runtime exceeds the", maxhours, "hours allowed. Setting the time to", maxhours, "hours...")
    timeinhours = maxhours
  # add a space after ptf if it exists
  #ptf_part = f"{ptf} " if ptf else ''
  ptf_part = ptf + " " if ptf else ''

  # ------------------------------
  # Build SRUN_COMMAND depending on options
  # ------------------------------
  # distributed: salloc gets the job; the session is an interactive step
  # (srun --interactive) that holds none of its CPUs, so srun from the
  # notebook can use all the tasks
  srun_args = ['salloc' if (parenv == 'distributed') else 'srun', '-J', 'JUPYNB', '--time=%d:00:00' % timeinhours]
  if (qos != ''):
    srun_args.append('--qos=%s' % qos)
  if (account != ''):
    srun_args.append('-A %s' % account)
  if (parenv == 'distributed'):
    srun_args.append('--ntasks=%d' % numberofslots)
    if (numberofnodes > 0):
      srun_args.append('--nodes=%d' % numberofnodes)
  else:
    srun_args += ['--nodes=1', '--ntasks=1', '--cpus-per-task=%d' % numberofslots]
  srun_args.append('--mem-per-cpu=%dG' % memoryingb)
  if (exclusive == 'yes'):
    srun_args.append('--exclusive')
  if (usegpu == 'yes'):
    srun_args.append('--gres=gpu:1')
  # node features: GPU card, GPU memory and CPU type
  features = []
  if (gpu != ''):
    features.append(gpu)
  if (gpumem != ''):
    features.append('vmem%sG' % gpumem)
  if (arch != ''):
    features.append('(%s)' % arch if ('|' in arch and features) else arch)
  if features:
    srun_args.append("-C '%s'" % '&'.join(features))
  # the session reports where it runs, then becomes a login shell
  if (parenv == 'distributed'):
    srun_args += ['srun', '--interactive', '--preserve-env']
  srun_args += ['--pty', 'bash', '-c', "'echo HOSTNAME=$(hostname) SLURM_JOB_ID=$SLURM_JOB_ID; exec bash -l'"]
  SRUN_COMMAND = ' '.join(srun_args)
  # --- module activation templates ---
  MODULE_LOADCUDA_TEMP = "module load cuda/%s"
  MODULE_LOAD_TEMP = "module load %s"
  CONDA_ACTIVATE = "conda activate %s"
  PYTHON_ACTIVATE= "source %s/bin/activate"
  APPTAINER_ACTIVATE= "apptainer exec --nv " if (usegpu == 'yes') else "apptainer exec "
  CD_TEMP = "cd %s"
  
  # print interactive session summary ---
  print()
  print("THE FOLLOWING INTERACTIVE SESSION WILL BE REQUESTED:")
  print()
  print(SRUN_COMMAND+"\n")
  
  print("\nIF THE REQUESTED RESOURCES DO NOT LOOK CORRECT INTERRUPT W/ `Control C`")
  print("RUN: `h2jupynb --help` TO SEE A FULL LIST OF OPTIONS") 
  time.sleep(2) 
  #exit()
  
  # --- jupyter notebook template ---
  STARTNBCMD = "jupyter notebook"
  print('STARTNBCMD = ',STARTNBCMD)
  
  # ------------------------------
  # Set up SSH session
  # ------------------------------
  ssh_args_common = [
    'ssh',
    '-o', 'ServerAliveCountMax=5',
    '-o', 'IPQoS=throughput',
    '-o', 'ServerAliveInterval=30',
    '-X', '-Y', '-t', '-t', '-4',
    '%s@%s' % (username, H2_HOST)
  ]

  if sshlogs == 'no':
    ssh_args = ssh_args_common
  elif sshlogs == 'yes':
    ssh_args = ssh_args_common[:-1] + ['-vvv', '-E', 'myssh.log', ssh_args_common[-1]]

  # Make this connection a ControlMaster, so the tunnel can be added to it
  # later without authenticating again. ControlPersist=no: the master (and
  # its socket) goes away with the session. Not supported on Windows.
  cpath = ''
  if (sharessh == 'yes') and not is_windows:
    cpath = make_control_path()
  if cpath:
    ssh_args = ssh_args[:1] + ['-o', 'ControlMaster=yes', '-o', 'ControlPath=%s' % cpath,
                               '-o', 'ControlPersist=no'] + ssh_args[1:]

  # Open pipes; on Python3 we use text mode with explicit UTF-8 decoding
  # and errors='replace'. Under Python2 we open binary pipes and rely on
  # readwhile to decode bytes (Python2's Popen doesn't support text/encoding args).
  if is_py2:
    if opsys == 'Windows':
      pqsub = Popen(ssh_args,stdin=PIPE,stdout=PIPE,stderr=PIPE,creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
    else:
      pqsub = Popen(ssh_args,stdin=PIPE,stdout=PIPE,stderr=PIPE)
  else:
    if opsys == 'Windows':
      pqsub=Popen(ssh_args,stdin=PIPE,stdout=PIPE,stderr=PIPE,text=True,encoding='utf-8',errors='replace',creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
    else:
      pqsub=Popen(ssh_args,stdin=PIPE,stdout=PIPE,stderr=PIPE,text=True,encoding='utf-8',errors='replace')
  # Flush pipes once
  pqsub.stdin.flush()
  pqsub.stdout.flush()

    
  # ------------------------------
  # Set up interactive session (pqsub)
  # ------------------------------    
  # srun waits until the resources are allocated. It discards input typed
  # ahead, so the commands for the compute node are sent only once it has
  # reported its host name. If srun fails, the login shell runs the echo below.
  pqsub.stdin.write(SRUN_COMMAND+"\n")
  # printenv rather than $SLURM_JOB_ID: tcsh aborts on undefined variables
  pqsub.stdin.write('echo HOSTNAME=`hostname` SLURM_JOB_ID=`printenv SLURM_JOB_ID`\n')
  pqsub.stdin.flush()
  
  # wait and capture hostname
  readwhile_threaded(pqsub.stdout, gethostname)
  
  # check that we are inside the Slurm job and not on a login node
  if ('login' in hostname) or (jobid == ''):
    print("\n ############################################")
    print(" ******We seem to be on the wrong host!******")
    print(" Most likely the Jupyter session cannot start")
    print(" because the interactive session cannot be   ")
    print(" allocated (check the srun error above). If  ")
    print(" you belong to more than one Slurm account,  ")
    print(" select one with -A <account> (see: h2 whoami")
    print(" on the cluster). Otherwise open a ticket at:")
    print(" https://support.idre.ucla.edu/helpdesk      ")
    print(" providing the exact command you used to     ")
    print(" start the session. Copying and pasting the  ")
    print(" the output of your terminal from the issuing of ")
    print(" the h2jupynb command will expedite matters. ")
    print(" **NOTE: Check GPU availability with: h2 gpus")
    print(" Exiting...                                  ")
    print(" #############################################\n")
    try:
      pqsub.kill()
    except Exception:
      pass
    quit()
  print("Running in Slurm job", jobid, "on", hostname)
  
  if (usegpu == 'yes'):
    pqsub.stdin.write(MODULE_LOADCUDA_TEMP % (cudaver) +"\n")
  
  if (mods != ''):
    pqsub.stdin.write(modstoload +"\n")

  pqsub.stdin.write(MODULE_LOAD_TEMP % (pythonver) +"\n")
  
  if (cve != ''):
    pqsub.stdin.write(CONDA_ACTIVATE % (cve) +"\n")
  if (pve != ''):
    pqsub.stdin.write(PYTHON_ACTIVATE % (pve) +"\n")
  
  # send initial commands to pqsub
  pqsub.stdin.write('module li\n')
  pqsub.stdin.write(CD_TEMP % (directory) +"\n")
  pqsub.stdin.flush()
  
  # ------------------------------
  # Start jupyter session
  # ------------------------------      
  if ptf == '' and pythonver != 'apptainer':
    pqsub.stdin.write(
      "command -v jupyter >/dev/null || echo 'JUPYTER NOT FOUND: use -v 3.12.13 or install jupyter in your own virtual environment (see: h2jupynb --help_envs)'\n"
    )
    pqsub.stdin.write(
      "{} --no-browser --ip=0.0.0.0 --port={}; echo exitcode: $?\n".format(
        STARTNBCMD, port
      )
    )
  else:
    pqsub.stdin.write(
      "{}{}{} --no-browser --ip=0.0.0.0 --port={}; echo exitcode: $?\n".format(
        APPTAINER_ACTIVATE, ptf_part, STARTNBCMD, port
      )
    )  
  pqsub.stdin.flush()
  
  # wait for the Jupyter notebook URL
  if (STARTNBCMD == 'jupyter notebook'):
    readwhile_threaded(pqsub.stdout, gethttpcommand)
  else:
    readwhile_threaded(pqsub.stdout, lambda line: line.find('http://%s:%s/'%(hostname,port))>0)
  
  if (httpcommand == ''):
    print("\nJupyter could not be started on", hostname, "(see the output above). Exiting...")
    try:
      pqsub.kill()
    except Exception:
      pass
    quit()
  print("httpcommand=",httpcommand)  
  token = httpcommand.split("=")[-1]
  newhttpcommand = 'http://localhost:%s/?token=%s'%(port,token)
  print("httpcommand=",newhttpcommand)

  # ------------------------------
  # Start SSH port forwarding
  # ------------------------------
  # The tunnel listens only on localhost of this machine
  forward = '%s:%s:%s' % (port, hostname, port)
  forwarded = False
  if cpath:
    print('\nAdding the SSH tunnel to the existing connection: -L', forward)
    forwarded = ssh_control(cpath, 'forward', ['-L', forward])
    if not forwarded:
      print('Could not add the tunnel to the existing connection: opening a second one...')

  if not forwarded:
    tunnel = [
      'ssh',
      '-o', 'ServerAliveCountMax=5',
      '-o', 'IPQoS=throughput',
      '-o', 'ServerAliveInterval=30',
      '-L', forward,
      '-t', '%s@%s' % (username, H2_HOST)
    ]
    if sshlogs == 'yes':
      tunnel = tunnel[:-2] + ['-vvv', '-E', 'mytunnel.log'] + tunnel[-2:]

    print('\nAbout to open SSH tunnel:')
    print(' You may be prompted for your Hoffman2 password or your SSH key passphrase')
    print(' a second time when the tunnel is created. This is normal if SSH key-based')
    print(' authentication or an SSH agent is not in use for the tunnel step.')
    print(' To avoid a second prompt, consider:')
    print('  - loading your key into an SSH agent (ssh-agent / Pageant),')
    print('  - or using an SSH key without a passphrase (less secure).')
    print('')
    print('Proceeding to open the SSH tunnel...')
    print(' '.join(tunnel))

    # Open tunnel pipes. On Python 2 we use binary pipes and let readwhile decode;
    # on Python 3 we use text mode with explicit UTF-8 decoding and errors='replace'.
    if is_py2:
      if opsys == 'Windows':
        ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE,creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
      else:
        ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE)
    else:
      if opsys == 'Windows':
        ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE,text=True,encoding='utf-8',errors='replace',creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
      else:
        ptunnel = Popen(tunnel,stdout=PIPE,stdin=PIPE,text=True,encoding='utf-8',errors='replace')

    ptunnel.stdin.write('echo TUNNEL\n')
    ptunnel.stdin.flush()

    # Wait until we see "TUNNEL" in stdout
    readwhile_threaded(ptunnel.stdout,lambda line: line.startswith('TUNNEL'))
  
  if ('CYGWIN' in opsys) or ('Cygwin' in opsys) or ('cygwin' in opsys):
    if (STARTNBCMD == 'jupyter notebook'):
      call(["cygstart",newhttpcommand])
    else:
      call(["cygstart",'http://localhost:%s'%(port)])
  else:
    if (STARTNBCMD == 'jupyter notebook'):
      webbrowser.open(newhttpcommand)
    else:
      webbrowser.open('http://localhost:%s'%(port))
  
  print("Succesfully opened notebook!")
  print("Kill this process to end your notebook connection.")
  print("")
  print("How to terminate:")
  print(" - On Unix/macOS: press Ctrl-C in this terminal to request shutdown.")
  print(" - On Windows PowerShell: if Ctrl-C does not stop the script, press Ctrl+Break (Pause/Break)")
  print("   or the laptop equivalent (often Ctrl+Fn+Break or Ctrl+Fn+B). The distributable will also")
  print("   attempt to send a CTRL_BREAK_EVENT to the SSH/tunnel child processes to help them exit.")
  print("")
  print("If child processes do not terminate, the script will fall back to killing them. If that fails,")
  print("you can close this terminal window as a last resort.")
  
  # set up Control C handler 
  #signal.signal(signal.SIGINT, handler)
  
  # Start a background countdown so the main thread remains responsive to Ctrl-C.
  countdown = Countdown(timeinhours * 3600, interval=10)
  countdown.start()
  try:
    while countdown.is_alive():
      countdown.join(timeout=1)
  except KeyboardInterrupt:
    print('\nCtrl-C pressed. Cleaning up...')
  finally:
    countdown.stop()
    countdown.join(timeout=2)
    
    # Clean up child processes
    if cpath:
      ssh_control(cpath, 'exit')
    try:
      pqsub.kill()
    except Exception:
      pass
    try:
      ptunnel.kill()
    except Exception:
      pass
    
    print("Elapsted runtime. Terminating...")
    print("Succesfully cleaned up connections.")
### function main def ends here


if __name__ == "__main__":
   main(sys.argv[1:])