Replace shell commands with subprocess.run

Signed-off-by: Bruno Alvisio <[email protected]>

Author: balvisio

ci/scripts/pytest_runner.sh

@@ -57,7 +57,7 @@ SKIP_SLOW=false
 ONLY_SLOW=false
 ALLOW_NO_TESTS=false
 # TODO(@cspades): Ignore this Evo2 notebook test, which has a tendency to leave a 32GB orphaned process in GPU.
-declare -a IGNORE_FILES=("sub-packages/bionemo-evo2/examples/fine-tuning-tutorial.ipynb")
+declare -a IGNORE_FILES=()
 error=false
 
 # Parse command line arguments

ci/scripts/run_pytest_notebooks.sh

@@ -19,4 +19,4 @@
 # Enable strict mode with better error handling
 set -euox pipefail
 
-pytest -v --nbval-lax -p no:python docs/ sub-packages/
+pytest -v --nbval-lax -x -p no:python docs/ sub-packages/

pyproject.toml

@@ -114,12 +114,12 @@ known-first-party = ["bionemo", "infra_bionemo"]
 convention = "google"
 
 [tool.pytest.ini_options]
+xfail_strict = true
 norecursedirs = ["3rdparty"]
 addopts = [
     "--durations-min=30.0",
     "--durations=0",
     "--ignore=3rdparty",
-    "--ignore-glob=sub-packages/bionemo-evo2/examples/fine-tuning-tutorial.ipynb",
     "--ignore-glob=sub-packages/bionemo-moco/examples/discrete_data_interpolant_tutorial.ipynb"
 ]
 markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]

sub-packages/bionemo-core/src/bionemo/core/utils/subprocess_utils.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import subprocess
+from typing import Any, Dict, List
+
+
+logger = logging.getLogger(__name__)
+
+
+def run_subprocess_safely(command: List[str], timeout: int = 2000) -> Dict[str, Any]:
+    """Run a subprocess and raise an error if it fails.
+
+    Args:
+        command: The command to run.
+        timeout: The timeout for the command.
+
+    Returns:
+        The result of the subprocess.
+    """
+    try:
+        result = subprocess.run(command, capture_output=True, timeout=timeout, check=True, text=True)
+        return {"stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode}
+    except subprocess.TimeoutExpired as e:
+        logger.error(f"Command timed out. Command: {command}\nstdout:\n{e.stdout}\nstderr:\n{e.stderr}")
+        return {"error": "timeout", "stdout": e.stdout, "stderr": e.stderr, "returncode": None}
+
+    except subprocess.CalledProcessError as e:
+        logger.error(
+            f"Command failed. Command: {command}\nreturncode: {e.returncode}\nstdout:\n{e.stdout}\nstderr:\n{e.stderr}"
+        )
+        return {"error": "non-zero exit", "stdout": e.stdout, "stderr": e.stderr, "returncode": e.returncode}
+
+    except FileNotFoundError as e:
+        logger.error(f"Command not found. Command: {command}\nstderr:\n{str(e)}")
+        return {"error": "not found", "stdout": "", "stderr": str(e), "returncode": None}
+
+    except Exception as e:
+        # catch-all for other unexpected errors
+        return {"error": "other", "message": str(e), "stdout": "", "stderr": "", "returncode": None}

sub-packages/bionemo-evo2/examples/fine-tuning-tutorial.ipynb

@@ -96,6 +96,8 @@
     "%%capture\n",
     "import os\n",
     "\n",
+    "from bionemo.core.utils.subprocess_utils import run_subprocess_safely\n",
+    "\n",
     "\n",
     "concat_path = \"chr20_21_22.fa\"\n",
     "if not os.path.exists(concat_path):\n",
@@ -270,38 +272,78 @@
     "#  the big difference is the --ckpt-dir argument which points to a pre-existing checkpoint from some other training run.\n",
     "\n",
     "if FAST_CI_MODE:\n",
-    "    model_subset_option = (\n",
-    "        \"--num-layers 4 --hybrid-override-pattern SDH* --activation-checkpoint-recompute-num-layers 2\"\n",
-    "    )\n",
+    "    model_subset_option = [\n",
+    "        \"--num-layers\",\n",
+    "        \"4\",\n",
+    "        \"--hybrid-override-pattern\",\n",
+    "        \"SDH*\",\n",
+    "        \"--activation-checkpoint-recompute-num-layers\",\n",
+    "        \"2\",\n",
+    "    ]\n",
     "else:\n",
     "    # By default do 5 layers of activation checkpointing\n",
-    "    model_subset_option = \"--activation-checkpoint-recompute-num-layers 5\"\n",
-    "train_cmd = f\"\"\"train_evo2 \\\n",
-    "    -d training_data_config.yaml \\\n",
-    "    --dataset-dir ./preprocessed_data \\\n",
-    "    --result-dir pretraining_demo \\\n",
-    "    --experiment-name evo2 \\\n",
-    "    --model-size 1b \\\n",
-    "    --devices 1 \\\n",
-    "    --num-nodes 1 \\\n",
-    "    --seq-length 8192 \\\n",
-    "    --micro-batch-size 2 \\\n",
-    "    --lr 0.000015 \\\n",
-    "    --min-lr 0.0000149 \\\n",
-    "    --warmup-steps {warmup_steps} \\\n",
-    "    --grad-acc-batches 4 \\\n",
-    "    --max-steps {MAX_STEPS} \\\n",
-    "    --ckpt-dir nemo2_evo2_1b_8k \\\n",
-    "    --clip-grad 250 \\\n",
-    "    --wd 0.001 \\\n",
-    "    --attention-dropout 0.01 \\\n",
-    "    --hidden-dropout 0.01 \\\n",
-    "    --val-check-interval {val_check_interval} \\\n",
-    "    {model_subset_option} \\\n",
-    "    --create-tensorboard-logger \\\n",
-    "    --ckpt-async-save\"\"\"\n",
-    "\n",
-    "!{train_cmd}"
+    "    model_subset_option = [\"--activation-checkpoint-recompute-num-layers\", \"5\"]\n",
+    "\n",
+    "train_cmd = [\n",
+    "    \"train_evo2\",\n",
+    "    \"-d\",\n",
+    "    \"training_data_config.yaml\",\n",
+    "    \"--dataset-dir\",\n",
+    "    \"./preprocessed_data\",\n",
+    "    \"--result-dir\",\n",
+    "    \"pretraining_demo\",\n",
+    "    \"--experiment-name\",\n",
+    "    \"evo2\",\n",
+    "    \"--model-size\",\n",
+    "    \"1b\",\n",
+    "    \"--devices\",\n",
+    "    \"1\",\n",
+    "    \"--num-nodes\",\n",
+    "    \"1\",\n",
+    "    \"--seq-length\",\n",
+    "    \"8192\",\n",
+    "    \"--micro-batch-size\",\n",
+    "    \"2\",\n",
+    "    \"--lr\",\n",
+    "    \"0.000015\",\n",
+    "    \"--min-lr\",\n",
+    "    \"0.0000149\",\n",
+    "    \"--warmup-steps\",\n",
+    "    str(warmup_steps),\n",
+    "    \"--grad-acc-batches\",\n",
+    "    \"4\",\n",
+    "    \"--max-steps\",\n",
+    "    str(MAX_STEPS),\n",
+    "    \"--ckpt-dir\",\n",
+    "    \"nemo2_evo2_1b_8k\",\n",
+    "    \"--clip-grad\",\n",
+    "    \"250\",\n",
+    "    \"--wd\",\n",
+    "    \"0.001\",\n",
+    "    \"--attention-dropout\",\n",
+    "    \"0.01\",\n",
+    "    \"--hidden-dropout\",\n",
+    "    \"0.01\",\n",
+    "    \"--val-check-interval\",\n",
+    "    str(val_check_interval),\n",
+    "    \"--create-tensorboard-logger\",\n",
+    "    \"--ckpt-async-save\",\n",
+    "]\n",
+    "\n",
+    "train_cmd.extend(model_subset_option)\n",
+    "\n",
+    "print(f\"Running command: {train_cmd}\")\n",
+    "\n",
+    "result = run_subprocess_safely(train_cmd)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert result[\"returncode\"] == 0, result"
    ]
   },
   {

sub-packages/bionemo-evo2/examples/zeroshot_brca1.ipynb

@@ -50,7 +50,9 @@
     "import seaborn as sns\n",
     "import torch\n",
     "from Bio import SeqIO\n",
-    "from sklearn.metrics import auc, roc_auc_score, roc_curve"
+    "from sklearn.metrics import auc, roc_auc_score, roc_curve\n",
+    "\n",
+    "from bionemo.core.utils.subprocess_utils import run_subprocess_safely"
    ]
   },
   {
@@ -662,24 +664,11 @@
     "# The Evo2 config has 'use_fp8_input_projections: True' by default\n",
     "\n",
     "if FAST_CI_MODE:\n",
-    "    model_subset_option = \"--num-layers 4 --hybrid-override-pattern SDH*\"\n",
+    "    model_subset_options = [\"--num-layers 4\", \"--hybrid-override-pattern SDH*\"]\n",
     "else:\n",
-    "    model_subset_option = \"\"\n",
-    "\n",
-    "fp8_option = \"--fp8\" if fp8_supported else \"\"\n",
+    "    model_subset_options = []\n",
     "\n",
-    "# Update predict commands to run on the full dataset\n",
-    "predict_ref_command = (\n",
-    "    f\"predict_evo2 --fasta {ref_fasta_path} --ckpt-dir {checkpoint_path} \"\n",
-    "    f\"--output-dir {predict_ref_dir} --model-size {MODEL_SIZE} --tensor-parallel-size 1  {model_subset_option} \"\n",
-    "    f\"--pipeline-model-parallel-size 1 --context-parallel-size 1 --output-log-prob-seqs {fp8_option}\"\n",
-    ")\n",
-    "\n",
-    "predict_var_command = (\n",
-    "    f\"predict_evo2 --fasta {var_fasta_path} --ckpt-dir {checkpoint_path} \"\n",
-    "    f\"--output-dir {predict_var_dir} --model-size {MODEL_SIZE} --tensor-parallel-size 1 {model_subset_option} \"\n",
-    "    f\"--pipeline-model-parallel-size 1 --context-parallel-size 1 --output-log-prob-seqs {fp8_option}\"\n",
-    ")"
+    "fp8_option = \"--fp8\" if fp8_supported else \"\""
    ]
   },
   {
@@ -696,8 +685,46 @@
    "outputs": [],
    "source": [
     "%%capture\n",
+    "# Update predict commands to run on the full dataset\n",
+    "# Update predict commands to run on the full dataset\n",
+    "predict_ref_command = [\n",
+    "    \"predict_evo2\",\n",
+    "    \"--fasta\",\n",
+    "    ref_fasta_path,\n",
+    "    \"--ckpt-dir\",\n",
+    "    checkpoint_path,\n",
+    "    \"--output-dir\",\n",
+    "    predict_ref_dir,\n",
+    "    \"--model-size\",\n",
+    "    MODEL_SIZE,\n",
+    "    \"--tensor-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--pipeline-model-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--context-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--output-log-prob-seqs\",\n",
+    "]\n",
+    "\n",
+    "# Optional flags\n",
+    "if model_subset_options:\n",
+    "    predict_ref_command.extend(model_subset_options)\n",
+    "\n",
+    "if fp8_option:\n",
+    "    predict_ref_command.append(fp8_option)\n",
+    "\n",
     "print(f\"Running command: {predict_ref_command}\")\n",
-    "!{predict_ref_command}"
+    "\n",
+    "result = run_subprocess_safely(predict_ref_command)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert result[\"returncode\"] == 0, result"
    ]
   },
   {
@@ -714,8 +741,43 @@
    "outputs": [],
    "source": [
     "%%capture\n",
+    "predict_var_command = [\n",
+    "    \"predict_evo2\",\n",
+    "    \"--fasta\",\n",
+    "    var_fasta_path,\n",
+    "    \"--ckpt-dir\",\n",
+    "    checkpoint_path,\n",
+    "    \"--output-dir\",\n",
+    "    predict_var_dir,\n",
+    "    \"--model-size\",\n",
+    "    MODEL_SIZE,\n",
+    "    \"--tensor-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--pipeline-model-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--context-parallel-size\",\n",
+    "    \"1\",\n",
+    "    \"--output-log-prob-seqs\",\n",
+    "]\n",
+    "\n",
+    "if model_subset_options:\n",
+    "    predict_var_command.extend(model_subset_options)  # make sure this is already a single flag or \"--key value\"\n",
+    "\n",
+    "if fp8_option:\n",
+    "    predict_var_command.append(fp8_option)\n",
+    "\n",
     "print(f\"Running command: {predict_var_command}\")\n",
-    "!{predict_var_command}"
+    "\n",
+    "result = run_subprocess_safely(predict_var_command)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert result[\"returncode\"] == 0, result"
    ]
   },
   {