address comments

Signed-off-by: Meng Xin <mxin@nvidia.com>

Author: mxinO

.claude/skills/common/remote_exec.sh

@@ -35,7 +35,8 @@
 #   REMOTE_WORKSPACE, REMOTE_GPU_TYPE, REMOTE_ENV_TYPE,
 #   REMOTE_CONTAINER_IMAGE, REMOTE_SLURM_ACCOUNT, REMOTE_SLURM_PARTITION
 
-set -euo pipefail
+# NOTE: This file is designed to be sourced. It does NOT set shell options
+# (set -euo pipefail) to avoid mutating the caller's environment.
 
 # ── Helpers ──────────────────────────────────────────────────────────────────
 
@@ -155,8 +156,7 @@ remote_start_session() {
     local rc=$?
     if (( rc == 0 )); then
         echo "SSH session established. All commands will reuse this connection."
-        # Register cleanup trap
-        trap 'remote_stop_session 2>/dev/null' EXIT
+        echo "Call 'remote_enable_cleanup_trap' to auto-close on exit, or 'remote_stop_session' manually."
     else
         echo "WARNING: Failed to start persistent SSH session (rc=$rc). Commands will use individual connections." >&2
     fi
@@ -173,6 +173,18 @@ remote_stop_session() {
     fi
 }
 
+remote_enable_cleanup_trap() {
+    # Opt-in: register an EXIT trap to auto-close the SSH session.
+    # Chains with any existing EXIT trap to avoid breaking the caller.
+    local existing_trap
+    existing_trap=$(trap -p EXIT | sed "s/^trap -- '//;s/' EXIT$//")
+    if [[ -n "$existing_trap" ]]; then
+        trap "${existing_trap}; remote_stop_session 2>/dev/null" EXIT
+    else
+        trap 'remote_stop_session 2>/dev/null' EXIT
+    fi
+}
+
 # ── Core Functions ───────────────────────────────────────────────────────────
 
 remote_load_cluster() {
@@ -278,8 +290,12 @@ remote_detect_env() {
         # Check Docker
         if command -v docker &>/dev/null; then
             echo 'HAS_DOCKER=yes';
-            # Check if docker can access GPUs
-            docker run --rm --gpus all nvidia/cuda:12.4.0-base-ubuntu22.04 nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null && echo 'DOCKER_GPU=yes' || echo 'DOCKER_GPU=no';
+            # Check if nvidia-container-cli is available (GPU support without pulling an image)
+            if command -v nvidia-container-cli &>/dev/null || docker info 2>/dev/null | grep -qi nvidia; then
+                echo 'DOCKER_GPU=yes';
+            else
+                echo 'DOCKER_GPU=no';
+            fi;
         else
             echo 'HAS_DOCKER=no';
         fi;
@@ -297,10 +313,13 @@ remote_detect_env() {
 
     if echo "$info" | grep -q "HAS_SLURM=yes"; then
         REMOTE_ENV_TYPE="slurm"
-    elif echo "$info" | grep -q "HAS_DOCKER=yes"; then
+    elif echo "$info" | grep -q "DOCKER_GPU=yes"; then
         REMOTE_ENV_TYPE="docker"
     elif echo "$info" | grep -q "HAS_BARE_GPU=yes"; then
         REMOTE_ENV_TYPE="bare"
+    elif echo "$info" | grep -q "HAS_DOCKER=yes"; then
+        # Docker available but no GPU support — fall back to bare
+        REMOTE_ENV_TYPE="bare"
     else
         REMOTE_ENV_TYPE="unknown"
     fi
@@ -316,7 +335,7 @@ remote_run() {
     # Retries up to 3 times on SSH connection failures.
     local cmd="$1"
     local ws="${REMOTE_WORKSPACE:-\$HOME}"
-    local full_cmd="cd $ws && $cmd"
+    local full_cmd="cd \"$ws\" && $cmd"
     local encoded
     encoded=$(printf '%s' "$full_cmd" | base64 -w0)
 

.claude/skills/ptq/references/remote-execution.md

@@ -140,7 +140,7 @@ remote_sync_from <remote_output_subdir> /local/output/
 | `unix_listener: cannot bind to path ... Read-only file system` | SSH ControlMaster socket in non-writable `/tmp` | `remote_exec.sh` auto-finds writable dir; ensure `TMPDIR` or `/tmp/claude-*` exists |
 | `cd: /home/user/~/path: No such file or directory` | `~` not expanding on remote | Use absolute paths in `workspace` config, not `~/...` |
 | Login nodes resolve home dirs differently | Symlinked home dirs vary by node | Use absolute lustre/NFS paths (e.g., `/lustre/fs1/...`) in job scripts |
-| `#!` becomes `#\!` in scripts | Shell environment mangles shebang | Fix with `sed -i 's\|^#\\\\!\|#!\| script.sh'` after writing |
+| `#!` becomes `#\!` in scripts | Shell environment mangles shebang | Fix with `sed -i 's\|^#\\\\!\|#!\|' script.sh` after writing |
 
 ## Reference Files
 

.claude/skills/ptq/references/unsupported-models.md

@@ -18,12 +18,12 @@ print(type(cfg).__name__)
 
   ```bash
   python -c "
-  import inspect
-  from transformers import AutoConfig, AutoModel
+  import importlib, inspect
+  from transformers import AutoConfig
   cfg = AutoConfig.from_pretrained('<ckpt_path>')
-  cls = AutoModel._model_type_to_module_name.get(cfg.model_type)
-  import transformers; mod = getattr(transformers, cls, None)
-  print(inspect.getfile(mod) if mod else 'not found')
+  mod_name = 'transformers.models.' + cfg.model_type.replace('-', '_')
+  mod = importlib.import_module(mod_name + '.modeling_' + cfg.model_type.replace('-', '_'))
+  print(inspect.getfile(mod))
   "
   ```