Repository: NVIDIA/NeMo-Skills
Pull request #455: [CodeGen] Refactor pipeline scripts for additional features
Author: login: titu1994
Labels: name: run GPU tests
Created: 2025-05-09T05:19:25Z | Merged: 2025-05-23T03:32:55Z

# Changelog
Primarily this PR adds two major things:
- support for dynamic mounting via str cli arg
- full featured `run_cmd`

Some of the scripts have an optional new argument `check_mounted_paths` which can perform cluster tunnel level checks to assert that some folders and paths exist. 

While codegen had far broader checks, going into individual input file level checks to assert that script cannot run until at least all filepaths are asserted to exist on cluster, we don't need to be so aggressive with checks in skills so I've kept it to minimal items and optional. I.e. auto mounting (mounting of absolute paths passed by user automatically) will only be done if this flag is enabled and only for a very specific set of inputs (like output dir, log dir and models sometimes).

**NOTE**: I still have to test each command. @Kipok  if you can share generic commands for them that would be good, or if its easier for you to run some tests yourself, let me know.

-----

# Enhancements to mounted path handling:
[GPT generated changelog]

* **New parameters for mounted paths:**
  - Added `mount_paths` and `check_mounted_paths` options to multiple pipeline functions (`check_contamination`, `convert`, `eval`, `generate`, and `run_cmd`) to allow specifying and validating paths to be mounted on remote machines. [[1]](diffhunk://#diff-3358d2230e92d60a6e4a72a94db50693501340b661d9dd7a2368fc310802804fR84) [[2]](diffhunk://#diff-3358d2230e92d60a6e4a72a94db50693501340b661d9dd7a2368fc310802804fR114) [[3]](diffhunk://#diff-7f42d537894a57a2a9a9b3681b793b6f988d6c80569fcd625dd6ad8c1867d138R184) [[4]](diffhunk://#diff-7f42d537894a57a2a9a9b3681b793b6f988d6c80569fcd625dd6ad8c1867d138R207) [[5]](diffhunk://#diff-58739ee434fd929a0c16a140b270be222c2367f0a7ed7f7fac6b492aa390408aR161) [[6]](diffhunk://#diff-58739ee434fd929a0c16a140b270be222c2367f0a7ed7f7fac6b492aa390408aR186-R193) [[7]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0R438) [[8]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0R498) [[9]](diffhunk://#diff-cae186e947820fc321b89b73412e2530f754ffbc31cf31e2581007f422d317a4R67-R83)
  - Introduced logic to resolve, validate, and handle mounted paths using utility functions like `resolve_mount_paths`, `is_mounted_filepath`, `add_mount_path`, and `check_remote_mount_directories`. [[1]](diffhunk://#diff-3358d2230e92d60a6e4a72a94db50693501340b661d9dd7a2368fc310802804fL124-R158) [[2]](diffhunk://#diff-7f42d537894a57a2a9a9b3681b793b6f988d6c80569fcd625dd6ad8c1867d138R236-R257) [[3]](diffhunk://#diff-58739ee434fd929a0c16a140b270be222c2367f0a7ed7f7fac6b492aa390408aL197-R231) [[4]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0R535-R559)

* **Cluster configuration updates:**
  - Updated cluster configuration handling to integrate mounted path resolution and ensure proper directory creation on remote machines when required. [[1]](diffhunk://#diff-3358d2230e92d60a6e4a72a94db50693501340b661d9dd7a2368fc310802804fL124-R158) [[2]](diffhunk://#diff-7f42d537894a57a2a9a9b3681b793b6f988d6c80569fcd625dd6ad8c1867d138R236-R257) [[3]](diffhunk://#diff-58739ee434fd929a0c16a140b270be222c2367f0a7ed7f7fac6b492aa390408aL197-R231) [[4]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0R535-R559)

### Code refactoring and cleanup:

* **Removed redundant methods:**
  - Eliminated the `wrap_cmd` function from `generate.py` and replaced it with a more comprehensive `get_cmd` function in `run_cmd.py`, which now handles command construction with support for preprocessing and postprocessing steps. [[1]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0L344-L355) [[2]](diffhunk://#diff-cae186e947820fc321b89b73412e2530f754ffbc31cf31e2581007f422d317a4L16-R45)

* **Simplified CLI utilities:**
  - Moved the `wrap_arguments` function from `cli.py` to `utils`, streamlining the CLI setup and maintaining consistency across modules. [[1]](diffhunk://#diff-2b5023e075231063a0b2fd8e4da68ddfbdb8ef9a0a6d19d7694356824d0e32f2R18-R21) [[2]](diffhunk://#diff-2b5023e075231063a0b2fd8e4da68ddfbdb8ef9a0a6d19d7694356824d0e32f2L36-L48)

* **Improved modularity:**
  - Consolidated utility imports across pipeline modules, ensuring all required functions for mounted path handling are consistently available. [[1]](diffhunk://#diff-3358d2230e92d60a6e4a72a94db50693501340b661d9dd7a2368fc310802804fR31-R36) [[2]](diffhunk://#diff-7f42d537894a57a2a9a9b3681b793b6f988d6c80569fcd625dd6ad8c1867d138L23-R35) [[3]](diffhunk://#diff-58739ee434fd929a0c16a140b270be222c2367f0a7ed7f7fac6b492aa390408aR33-R49) [[4]](diffhunk://#diff-2bf0a3d00b2c8bc23a7ee220ddce41ecd8d4428efc4404c6b8e6857ec91c6ca0L29-R43)

### Minor adjustments:

* Added a new `ExtraDatasetType` enum in `eval.py` for potential future enhancements related to handling extra datasets.


Diff:
diff --git a/docs/pipelines/index.md b/docs/pipelines/index.md
index 2c6a901737..12fb50696b 100644
--- a/docs/pipelines/index.md
+++ b/docs/pipelines/index.md
@@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
+    - Check if the paths specified in the script are mounted correctly. This is useful if you want to make 
+    sure that the paths that are mounted are available  on remote machine before running the job. 
+    E.g. use `--check_mounted_paths` to check if `/my/remote/workspace` folder from the host machine 
+    is a folder that exists and can be mounted.
+    - In many cases, if the directory does not exist, we will create it for you. This is useful for 
+    output and log directories.
+    - If paths are provided but not mounted, often times we will dynamically mount them for you.
 - **--partition**: Can be used to run in a specific slurm partition (e.g. commonly used
   to launch interactive jobs).
 - **--not_exclusive**: Can be used if you want to request a part o the slurm node. By default
diff --git a/docs/pipelines/run-cmd.md b/docs/pipelines/run-cmd.md
index fbfe5a820d..30a7fff1c2 100644
--- a/docs/pipelines/run-cmd.md
+++ b/docs/pipelines/run-cmd.md
@@ -14,4 +14,57 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command as below. The --with_sandbox argument starts the 
+code execution server that can be used to run arbitrary code in a sandboxed environment and is added here just as a 
+demonstration. While the current example does not use it, this can be useful to execute code or to run code that 
+requires a specific environment in a container. 
+
+**Note**: While the container is a little more secure than running code directly on the host, it is still not a fully 
+secure sandbox and should not be used to run untrusted code.
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=Qwen/Qwen3-1.7B \
+    --server_type=vllm \
+    --server_gpus=1 \
+    --with_sandbox \
+    cd /nemo_run/code/ && python inference.py
+```
+
+This will launch the LLM inference server, the
\ No newline at end of file
diff --git a/nemo_skills/pipeline/check_contamination.py b/nemo_skills/pipeline/check_contamination.py
index 196dc4a811..78f140fa94 100644
--- a/nemo_skills/pipeline/check_contamination.py
+++ b/nemo_skills/pipeline/check_contamination.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from enum import Enum
 from typing import List
 
@@ -20,12 +21,18 @@
 from nemo_skills.pipeline.app import app, typer_unpacker
 from nemo_skills.pipeline.generate import wrap_cmd
 from nemo_skills.pipeline.utils import (
+    add_mount_path,
     add_task,
     check_if_mounted,
+    check_mounts,
+    create_remote_directory,
     get_cluster_config,
     get_exp,
     get_free_port,
     get_generation_command,
+    get_mounted_path,
+    is_mounted_filepath,
+    resolve_mount_paths,
     run_exp,
 )
 from nemo_skills.utils import setup_logging
@@ -79,6 +86,7 @@ def check_contamination(
         None, help="Can specify if need interactive jobs or a specific non-default partition"
     ),
     time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     run_after: List[str] = typer.Option(
         None, help="Can specify a list of expnames that need to be completed before this one starts"
     ),
@@ -108,6 +116,7 @@ def check_contamination(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
 ):
     """Check contamination between train/test via an LLM call.
 
@@ -126,10 +135,14 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
-    check_if_mounted(cluster_config, output_file)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    input_file, output_file, log_dir = check_mounts(
+        cluster_config,
+        log_dir=log_dir,
+        mount_map={input_file: "/mounted_data/input", output_file: "/mounted_data/output"},
+        check_mounted_paths=check_mounted_paths,
+    )
 
     if server_address is None:  # we need to host the model
         assert server_gpus is not None, "Need to specify server_gpus if hosting the model"
diff --git a/nemo_skills/pipeline/cli.py b/nemo_skills/pipeline/cli.py
index 7b986a11f5..4a2067bb22 100644
--- a/nemo_skills/pipeline/cli.py
+++ b/nemo_skills/pipeline/cli.py
@@ -15,7 +15,10 @@
 
 import typer
 
+# isort: off
 from nemo_skills.pipeline.app import app
+from nemo_skills.pipeline.utils import wrap_arguments
+# isort: on
 
 # need the imports to make sure the commands are registered
 from nemo_skills.pipeline.check_contamination import check_contamination
@@ -33,19 +36,6 @@
 
 typer.main.get_command_name = lambda name: name
 
-
-def wrap_arguments(arguments: str):
-    """Returns a mock context object to allow using the cli entrypoints as functions."""
-
-    class MockContext:
-        def __init__(self, args):
-            self.args = args
-            self.obj = None
-
-    # first one is the cli name
-    return MockContext(args=arguments.split(" "))
-
-
 if __name__ == "__main__":
     # workaround for https://github.com/fastapi/typer/issues/341
     app()
diff --git a/nemo_skills/pipeline/convert.py b/nemo_skills/pipeline/convert.py
index 2005c7a3d0..ff72f1ec4f 100644
--- a/nemo_skills/pipeline/convert.py
+++ b/nemo_skills/pipeline/convert.py
@@ -20,7 +20,19 @@
 import typer
 
 from nemo_skills.pipeline.app import app, typer_unpacker
-from nemo_skills.pipeline.utils import add_task, check_if_mounted, get_cluster_config, get_exp, run_exp
+from nemo_skills.pipeline.utils import (
+    add_mount_path,
+    add_task,
+    check_if_mounted,
+    check_mounts,
+    create_remote_directory,
+    get_cluster_config,
+    get_exp,
+    get_mounted_path,
+    is_mounted_filepath,
+    resolve_mount_paths,
+    run_exp,
+)
 from nemo_skills.utils import get_logger_name, setup_logging
 
 LOG = logging.getLogger(get_logger_name(__file__))
@@ -234,6 +246,7 @@ def convert(
         None, help="Can specify if need interactive jobs or a specific non-default partition"
     ),
     time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     run_after: List[str] = typer.Option(
         None, help="Can specify a list of expnames that need to be completed before this one starts"
     ),
@@ -256,6 +269,7 @@ def convert(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
 ):
     """Convert a checkpoint from one format to another.
 
@@ -300,12 +314,18 @@ def convert(
             # TODO: that's probably not true, but need to figure out how it's passed
             raise ValueError("Conversion to Megatron is only supported for bf16 models")
 
+    # Prepare cluster config and mount paths
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_model)
-    check_if_mounted(cluster_config, output_model)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
-    else:
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    input_model, output_model, log_dir = check_mounts(
+        cluster_config,
+        log_dir=log_dir,
+        mount_map={input_model: '/input_model', output_model: '/output_model'},
+        check_mounted_paths=check_mounted_paths,
+    )
+
+    if log_dir is None:
         log_dir = str(Path(output_model) / "conversion-logs")
 
     conversion_cmd_map = {
diff --git a/nemo_skills/pipeline/eval.py b/nemo_skills/pipeline/eval.py
index e9f40ca6a8..150d2b5d34 100644
--- a/nemo_skills/pipeline/eval.py
+++ b/nemo_skills/pipeline/eval.py
@@ -22,13 +22,19 @@
 from nemo_skills.dataset.utils import get_dataset_module
 from nemo_skills.pipeline.app import app, typer_unpacker
 from nemo_skills.pipeline.utils import (
+    add_mount_path,
     add_task,
     check_if_mounted,
+    check_mounts,
+    create_remote_directory,
     get_cluster_config,
     get_exp,
     get_free_port,
     get_generation_command,
+    get_mounted_path,
     get_server_command,
+    is_mounted_filepath,
+    resolve_mount_paths,
     run_exp,
 )
 from nemo_skills.utils import compute_chunk_ids, get_chunked_filename, get_logger_name, setup_logging
@@ -36,6 +42,11 @@
 LOG = logging.getLogger(get_logger_name(__file__))
 
 
+class ExtraDatasetType(str, Enum):
+    copy = "copy"
+    mount = "mount"
+
+
 def get_greedy_cmd(
     benchmark,
     split,
@@ -153,6 +164,7 @@ def eval(
     ),
     partition: str = typer.Option(None, help="Cluster partition to use"),
     time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     extra_eval_args: str = typer.Option("", help="Additional arguments for evaluation"),
     add_greedy: bool = typer.Option(
         False,
@@ -180,12 +192,14 @@ def eval(
         help="Path to a custom dataset folder that will be searched in addition to the main one. "
         "Can also specify through NEMO_SKILLS_EXTRA_DATASETS.",
     ),
+    # extra_datasets_type: ExtraDatasetType = typer.Option(ExtraDatasetType.copy, help="How to handle extra datasets"),
     exclusive: bool = typer.Option(
         True,
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
     with_sandbox: bool = typer.Option(False, help="If True, will start a sandbox container alongside this job"),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
 ):
     """Evaluate a model on specified benchmarks.
 
@@ -203,12 +217,18 @@ def eval(
         pass
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, output_dir)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
-    else:
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    if log_dir is None:
         log_dir = f"{output_dir}/eval-logs"
 
+    output_dir, log_dir = check_mounts(
+        cluster_config,
+        log_dir=log_dir,
+        mount_map={output_dir: None},
+        check_mounted_paths=check_mounted_paths,
+    )
+
     if num_chunks:
         chunk_ids = compute_chunk_ids(chunk_ids, num_chunks)
     should_chunk_dataset = num_chunks is not None and chunk_ids is not None
@@ -246,6 +266,11 @@ def eval(
     benchmarks = {k: int(v) for k, v in [b.split(":") for b in benchmarks.split(",")]}
 
     extra_datasets = extra_datasets or os.environ.get("NEMO_SKILLS_EXTRA_DATASETS")
+    # TODO(@titu1994): add support for extra_datasets_type in future pr
+    # try:
+    #     extra_datasets_type = extra_datasets_type.value
+    # except AttributeError:
+    #     pass
 
     # Check which benchmarks require sandbox
     benchmark_requires_sandbox = {}
diff --git a/nemo_skills/pipeline/generate.py b/nemo_skills/pipeline/generate.py
index 07c9339cbf..f7e2731e97 100644
--- a/nemo_skills/pipeline/generate.py
+++ b/nemo_skills/pipeline/generate.py
@@ -25,17 +25,23 @@
 from nemo_skills.inference.generate import GenerationTask
 from nemo_skills.pipeline.app import app, typer_unpacker
 from nemo_skills.pipeline.utils import (
+    add_mount_path,
     add_task,
     check_if_mounted,
+    check_mounts,
+    create_remote_directory,
     get_cluster_config,
     get_exp,
     get_free_port,
     get_generation_command,
+    get_mounted_path,
     get_reward_server_command,
     get_server_command,
     get_tunnel,
     get_unmounted_path,
+    resolve_mount_paths,
     run_exp,
+    wrap_cmd,
 )
 from nemo_skills.utils import compute_chunk_ids, get_chunked_filename, get_logger_name, setup_logging, str_ids_to_list
 
@@ -347,28 +353,6 @@ def get_genselect_cmd(
     return cmd, postprocess_cmd
 
 
-def wrap_cmd(cmd, preprocess_cmd, postprocess_cmd, random_seed=None, wandb_parameters=None):
-    if preprocess_cmd:
-        if random_seed is not None:
-            preprocess_cmd = preprocess_cmd.format(random_seed=random_seed)
-        cmd = f" {preprocess_cmd} && {cmd} "
-    if postprocess_cmd:
-        if random_seed is not None:
-            postprocess_cmd = postprocess_cmd.format(random_seed=random_seed)
-        cmd = f" {cmd} && {postprocess_cmd} "
-    if wandb_parameters:
-        log_wandb_cmd = (
-            f"python -m nemo_skills.inference.log_samples_wandb "
-            f"    {wandb_parameters['samples_file']} "
-            f"    --name={wandb_parameters['name']} "
-            f"    --project={wandb_parameters['project']} "
-        )
-        if wandb_parameters['group'] is not None:
-            log_wandb_cmd += f" --group={wandb_parameters['group']} "
-        cmd = f"{cmd} && {log_wandb_cmd} "
-    return cmd
-
-
 class GenerationType(str, Enum):
     generate = "generate"
     reward = "reward"
@@ -464,6 +448,7 @@ def generate(
         "If not specified, will use the default entrypoint for the server type.",
     ),
     dependent_jobs: int = typer.Option(0, help="Specify this to launch that number of dependent jobs"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     num_random_seeds: int = typer.Option(
         None, help="Specify if want to run many generations with high temperature for the same input"
     ),
@@ -521,6 +506,7 @@ def generate(
         False, help="If True, will re-run jobs even if a corresponding '.done' file already exists"
     ),
     with_sandbox: bool = typer.Option(False, help="If True, will start a sandbox container alongside this job"),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
     log_samples: bool = typer.Option(
         False,
         help="If True, will log random samples from the output files to wandb. "
@@ -581,13 +567,21 @@ def generate(
         chunk_ids = compute_chunk_ids(chunk_ids, num_chunks)
     if chunk_ids is None:
         chunk_ids = [None]
+
+    # Prepare cluster config and mount paths
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, output_dir)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
-    else:
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths, create_remote_dir=check_mounted_paths)
+
+    if not log_dir:
         log_dir = f"{output_dir}/generation-logs"
 
+    output_dir, log_dir = check_mounts(
+        cluster_config,
+        log_dir=log_dir,
+        mount_map={output_dir: None},
+        check_mounted_paths=check_mounted_paths,
+    )
+
     get_server_command = server_command_factories[generation_type]
     get_cmd = client_command_factories[generation_type]
     cmd_script = client_command_scripts[generation_type]
diff --git a/nemo_skills/pipeline/run_cmd.py b/nemo_skills/pipeline/run_cmd.py
index 248369483a..d65e99ee3f 100644
--- a/nemo_skills/pipeline/run_cmd.py
+++ b/nemo_skills/pipeline/run_cmd.py
@@ -13,21 +13,27 @@
 # limitations under the License.
 
 import logging
-from typing import List
+from typing import List, Optional
 
 import typer
 
+from nemo_skills import utils
+from nemo_skills.pipeline import utils as pipeline_utils
 from nemo_skills.pipeline.app import app, typer_unpacker
 from nemo_skills.pipeline.generate import wrap_cmd
-from nemo_skills.pipeline.utils import add_task, check_if_mounted, get_cluster_config, get_exp, run_exp
+from nemo_skills.pipeline.utils import add_task, check_if_mounted, check_mounts, get_cluster_config, get_exp, run_exp
 from nemo_skills.utils import get_logger_name, setup_logging
 
 LOG = logging.getLogger(get_logger_name(__file__))
 
 
-def get_cmd(extra_arguments):
-    cmd = f"{extra_arguments} "
-    cmd = f"export HYDRA_FULL_ERROR=1 && export PYTHONPATH=$PYTHONPATH:/nemo_run/code && cd /nemo_run/code && {cmd}"
+def get_cmd(command):
+    cmd = (
+        f"export HYDRA_FULL_ERROR=1 && "
+        f"export PYTHONPATH=$PYTHONPATH:/nemo_run/code && "
+        f"cd /nemo_run/code && "
+        f"{command.strip()} "
+    )
     return cmd
 
 
@@ -40,6 +46,9 @@ def run_cmd(
         help="One of the configs inside config_dir or NEMO_SKILLS_CONFIG_DIR or ./cluster_configs. "
         "Can also use NEMO_SKILLS_CONFIG instead of specifying as argument.",
     ),
+    command: str = typer.Option(
+        None, help="Command to run in the container. Can also be specified as extra arguments."
+    ),
     container: str = typer.Option("nemo-skills", help="Container to use for the run"),
     expname: str = typer.Option("script", help="Nemo run experiment name"),
     partition: str = typer.Option(
@@ -48,6 +57,14 @@ def run_cmd(
     time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
     num_gpus: int | None = typer.Option(None, help="Number of GPUs to use"),
     num_nodes: int = typer.Option(1, help="Number of nodes to use"),
+    model: str = typer.Option(None, help="Path to the model to evaluate"),
+    server_address: str = typer.Option(None, help="Address of the server hosting the model"),
+    server_type: str = typer.Option(None, help="Type of server to use"),
+    server_gpus: int = typer.Option(None, help="Number of GPUs to use if hosting the model"),
+    server_nodes: int = typer.Option(1, help="Number of nodes to use if hosting the model"),
+    server_args: str = typer.Option("", help="Additional arguments for the server"),
+    dependent_jobs: int = typer.Option(0, help="Specify this to launch that number of dependent jobs"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     run_after: List[str] = typer.Option(
         None, help="Can specify a list of expnames that need to be completed before this one starts"
     ),
@@ -66,6 +83,7 @@ def run_cmd(
     preprocess_cmd: str = typer.Option(None, help="Command to run before job"),
     postprocess_cmd: str = typer.Option(None, help="Command to run after job"),
     config_dir: str = typer.Option(None, help="Can customize where we search for cluster configs"),
+    with_sandbox: bool = typer.Option(False, help="If True, will start a sandbox container alongside this job"),
     log_dir: str = typer.Option(
         None,
         help="Can specify a custom location for slurm logs. "
@@ -76,32 +94,80 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
 ):
     """Run a pre-defined module or script in the NeMo-Skills container."""
     setup_logging(disable_hydra_logs=False, use_rich=True)
     extra_arguments = f'{" ".join(ctx.args)}'
 
-    cluster_config = get_cluster_config(cluster, config_dir)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+    # Assert that either command or extra_arguments is provided, not both
+    if command and extra_arguments:
+        raise ValueError("Please provide either a command or extra arguments, not both.")
+    elif not command and not extra_arguments:
+        raise ValueError("Please provide either a command or extra arguments.")
+
+    command = command or extra_arguments  # From here on, `command` will be used as the command to run
+    extra_arguments = ""  # Reset extra_arguments to avoid confusion
+
+    # Setup cluster config
+    cluster_config = pipeline_utils.get_cluster_config(cluster, config_dir)
+    cluster_config = pipeline_utils.resolve_mount_paths(
+        cluster_config, mount_paths, create_remote_dir=check_mounted_paths
+    )
+
+    log_dir = check_mounts(cluster_config, log_dir, check_mounted_paths=check_mounted_paths)
 
     with get_exp(expname, cluster_config) as exp:
-        add_task(
-            exp,
-            cmd=wrap_cmd(get_cmd(extra_arguments=extra_arguments), preprocess_cmd, postprocess_cmd),
-            task_name=expname,
-            log_dir=log_dir,
-            container=cluster_config["containers"][container],
-            cluster_config=cluster_config,
-            partition=partition,
-            time_min=time_min,
-            run_after=run_after,
-            reuse_code=reuse_code,
-            reuse_code_exp=reuse_code_exp,
-            num_gpus=num_gpus,
-            num_nodes=num_nodes,
-            slurm_kwargs={"exclusive": exclusive} if exclusive else None,
-        )
+        # Setup server config if model is provided
+        server_port = None if get_random_port else 5000
+        if model is not None:
+            server_config, extra_arguments, server_address, server_port = pipeline_utils.configure_client(
+                model=model,
+                server_type=server_type,
+                server_address=server_address,
+                server_port=server_port,
+                server_gpus=server_gpus,
+                server_nodes=server_nodes,
+                server_args=server_args,
+                extra_arguments=extra_arguments,  # this is empty string by design
+                get_random_port=get_random_port,
+            )
+        else:
+            server_config = None
+
+        # Prepare command
+        cmd = get_cmd(command=command)
+        cmd = wrap_cmd(cmd, preprocess_cmd, postprocess_cmd)
+
+        # Wrap command with generation command if model is provided
+        if model is not None and server_config is not None:
+            cmd = pipeline_utils.get_generation_command(server_address, cmd)
+
+        prev_tasks = None
+        for _ in range(dependent_jobs + 1):
+            # Add the task to the experiment
+            new_task = add_task(
+                exp,
+                cmd=cmd,
+                task_name=expname,
+                log_dir=log_dir,
+                container=cluster_config["containers"][container],
+                cluster_config=cluster_config,
+                partition=partition,
+                time_min=time_min,
+                server_config=server_config,
+                with_sandbox=with_sandbox,
+                sandbox_port=server_port,
+                run_after=run_after,
+                reuse_code=reuse_code,
+                reuse_code_exp=reuse_code_exp,
+                task_dependencies=prev_tasks,
+                num_gpus=num_gpus,
+                num_nodes=num_nodes,
+                slurm_kwargs={"exclusive": exclusive} if exclusive else None,
+            )
+            prev_tasks = [new_task]
         run_exp(exp, cluster_config)
 
     return exp
diff --git a/nemo_skills/pipeline/start_server.py b/nemo_skills/pipeline/start_server.py
index bd42922061..5cf7d318e0 100644
--- a/nemo_skills/pipeline/start_server.py
+++ b/nemo_skills/pipeline/start_server.py
@@ -16,7 +16,19 @@
 import typer
 
 from nemo_skills.pipeline.app import app, typer_unpacker
-from nemo_skills.pipeline.utils import add_task, check_if_mounted, get_cluster_config, get_exp, get_free_port
+from nemo_skills.pipeline.utils import (
+    add_mount_path,
+    add_task,
+    check_if_mounted,
+    check_mounts,
+    create_remote_directory,
+    get_cluster_config,
+    get_exp,
+    get_free_port,
+    get_mounted_path,
+    is_mounted_filepath,
+    resolve_mount_paths,
+)
 from nemo_skills.utils import setup_logging
 
 
@@ -48,6 +60,7 @@ def start_server(
     ),
     partition: str = typer.Option(None, help="Cluster partition to use"),
     time_min: str = typer.Option(None, help="If specified, will use as a time-min slurm parameter"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     with_sandbox: bool = typer.Option(
         False, help="Starts a sandbox (set this flag if model supports calling Python interpreter)"
     ),
@@ -62,20 +75,21 @@ def start_server(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
     get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
 ):
     """Self-host a model server."""
     setup_logging(disable_hydra_logs=False, use_rich=True)
 
     cluster_config = get_cluster_config(cluster, config_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
 
     try:
         server_type = server_type.value
     except AttributeError:
         pass
 
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+    log_dir = check_mounts(cluster_config, log_dir, check_mounted_paths=check_mounted_paths)
 
     server_config = {
         "model_path": model,
diff --git a/nemo_skills/pipeline/summarize_results.py b/nemo_skills/pipeline/summarize_results.py
index 843b3228c5..1a0fba6f6e 100644
--- a/nemo_skills/pipeline/summarize_results.py
+++ b/nemo_skills/pipeline/summarize_results.py
@@ -33,6 +33,7 @@
     get_cluster_config,
     get_tunnel,
     get_unmounted_path,
+    resolve_mount_paths,
 )
 from nemo_skills.utils import setup_logging
 
@@ -59,6 +60,7 @@ def summarize_results(
     ),
     remote_tar_dir: str = typer.Option(None, help="Directory where remote tar files are created on clusters"),
     debug: bool = typer.Option(False, help="Print debug information"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     max_samples: int = typer.Option(-1, help="Limit metric computation only to first `max_samples`"),
     extra_datasets: str = typer.Option(
         None,
@@ -89,6 +91,7 @@ def summarize_results(
     upload_path = None
     if cluster is not None:
         cluster_config = get_cluster_config(cluster, config_dir)
+        cluster_config = resolve_mount_paths(cluster_config, mount_paths)
         check_if_mounted(cluster_config, results_dir)
         if cluster_config.get("executor", "") == "local":
             results_dir = get_unmounted_path(cluster_config, results_dir)
diff --git a/nemo_skills/pipeline/train.py b/nemo_skills/pipeline/train.py
index 3379fb4e6a..5ee992ffc0 100755
--- a/nemo_skills/pipeline/train.py
+++ b/nemo_skills/pipeline/train.py
@@ -21,7 +21,18 @@
 import typer
 
 from nemo_skills.pipeline.app import app, typer_unpacker
-from nemo_skills.pipeline.utils import add_task, check_if_mounted, get_cluster_config, get_exp, get_timeout, run_exp
+from nemo_skills.pipeline.utils import (
+    add_task,
+    check_if_mounted,
+    check_mounts,
+    get_cluster_config,
+    get_exp,
+    get_free_port,
+    get_mounted_path,
+    get_timeout,
+    resolve_mount_paths,
+    run_exp,
+)
 from nemo_skills.utils import get_logger_name, setup_logging
 
 LOG = logging.getLogger(get_logger_name(__file__))
@@ -274,6 +285,7 @@ def train(
     server_gpus: int = typer.Option(None, help="Number of GPUs to use if hosting the model"),
     server_nodes: int = typer.Option(1, help="Number of nodes required for hosting LLM server"),
     server_args: str = typer.Option("", help="Any extra arguments to pass to the server"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
     run_after: List[str] = typer.Option(
         None, help="Can specify a list of expnames that need to be completed before this one starts"
     ),
@@ -296,6 +308,7 @@ def train(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
 ):
     """Train (SFT or DPO) an LLM model.
 
@@ -316,24 +329,29 @@ def train(
         pass
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, output_dir)
-    check_if_mounted(cluster_config, nemo_model)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
-    else:
-        log_dir = output_dir
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    if log_dir is None:
+        log_dir = f"{output_dir}"
+
+    nemo_model, output_dir, log_dir = check_mounts(
+        cluster_config,
+        log_dir=log_dir,
+        mount_map={nemo_model: None, output_dir: None},
+        check_mounted_paths=check_mounted_paths,
+    )
 
     if num_training_jobs > 0:
         if training_data is None:
             raise ValueError("training_data is required when num_training_jobs > 0")
-        check_if_mounted(cluster_config, training_data)
+        training_data = get_mounted_path(cluster_config, training_data)
 
     if not final_nemo_path:
         final_nemo_path = f"{output_dir}/model-averaged-nemo"
-    check_if_mounted(cluster_config, final_nemo_path)
+    final_nemo_path = get_mounted_path(cluster_config, final_nemo_path)
 
     if validation_data:
-        check_if_mounted(cluster_config, validation_data)
+        validation_data = get_mounted_path(cluster_config, validation_data)
 
     if " " in str(average_steps):
         raise ValueError("average steps should be separated with commas")
diff --git a/nemo_skills/pipeline/utils.py b/nemo_skills/pipeline/utils.py
index c23b4ebe5f..7edf4b72b6 100644
--- a/nemo_skills/pipeline/utils.py
+++ b/nemo_skills/pipeline/utils.py
@@ -24,7 +24,7 @@
 from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
-from typing import Optional
+from typing import Dict, Optional
 
 import nemo_run as run
 import yaml
@@ -43,6 +43,18 @@
 LOG = logging.getLogger(get_logger_name(__file__))
 
 
+def wrap_arguments(arguments: str):
+    """Returns a mock context object to allow using the cli entrypoints as functions."""
+
+    class MockContext:
+        def __init__(self, args):
+            self.args = args
+            self.obj = None
+
+    # first one is the cli name
+    return MockContext(args=arguments.split(" "))
+
+
 # TODO: this file is way too big - we need to split it into pieces
 
 # keeping a global variable for first submitted experiment (per cluster) and reusing it by default
@@ -131,6 +143,84 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
+                if is_file:
+                    path_dir, _ = os.path.split(path)
+                else:
+                    path_dir = path
+                remote_dir_list.append(path_dir)
+
+                if default_mount is not None:
+                    # Check that path is not empty and is an absolute path
+                    assert (
+                        default_mount[0] == '/'
+                    ), f"Default mount path should be absolute path, given {default_mount}"
+
+                    # Add mount path to the cluster config
+                    add_mount_path(path_dir, default_mount, cluster_config)
+
+    else:
+        # Just check if the paths are mounted
+        for path in mount_map.keys():
+            check_if_mounted(cluster_config, path)
+
+    # check if the paths are mounted, get them if they arent but have mount sources
+    # will error out if there are no corresponding mount sources
+    new_paths = [get_mounted_path(cluster_config, path) for path in mount_map.keys()]
+
+    if log_dir:
+        if check_mounted_paths:
+            # Create log dir in some location that will be mounted
+            remote_dir_list.append(log_dir)
+        else:
+            check_if_mounted(cluster_config, log_dir)
+        log_dir = get_mounted_path(cluster_config, log_dir)
+
+    if check_mounted_paths:
+        # Create missing remote directories
+        if remote_dir_list:
+            create_remote_directory(remote_dir_list, cluster_config)
+
+        # Check that the file or dir exists at the remote location
+        check_remote_mount_directories(list(mount_map.keys()) + [log_dir], cluster_config)
+
+    if new_paths:
+        return *new_paths, log_dir
+    return log_dir
+
+
 def get_mounted_path(cluster_config: dict, path: str):
     """
     Resolve the mounted filepath using the cluster config to merge the mount destination path to the filepath.
@@ -153,7 +243,7 @@ def get_mounted_path(cluster_config: dict, path: str):
 
     # Find which mount path matches the filepaths prefix
     mount_path = None
-    for mount in cluster_config['mounts']:
+    for mount in get_mounts_from_config(cluster_config) + ['/nemo_run/code:/nemo_run/code']:
         mount_source, mount_dest = mount.split(':')
         if path.startswith(mount_source):
             mount_path = mount
@@ -200,7 +290,7 @@ def get_unmounted_path(cluster_config: dict, path: str):
 
     # Find which mount path matches the filepaths prefix
     mount_path = None
-    for mount in get_mounts_from_config(cluster_config):
+    for mount in get_mounts_from_config(cluster_config) + ['/nemo_run/code:/nemo_run/code']:
         mount_source, mount_dest = mount.split(':')
         if path.startswith(mount_dest):
             mount_path = mount
@@ -213,7 +303,7 @@ def get_unmounted_path(cluster_config: dict, path: str):
     if mount_path is None:
         raise ValueError(
             f"Could not find a mount path for the file path `{path}`. Check cluster config. Below paths are mounted: \n"
-            f"{cluster_config['mounts']}"
+            f"{cluster_config['mounts'] + ['/nemo_run/code:/nemo_run/code']}"
         )
 
     # replace the mount destination inside the filepath with the mount source
@@ -232,7 +322,7 @@ def add_mount_path(mount_source: str, mount_dest: str, cluster_config):
         raise ValueError("Cluster config is not provided.")
 
     if 'mounts' in cluster_config:
-        original_mounts = get_mounts_from_config(cluster_config)
+        original_mounts = get_mounts_from_config(cluster_config) + ['/nemo_run/code:/nemo_run/code']
         added_mount = False
         for mount_path in original_mounts:
             source, destination = mount_path.split(':')
@@ -257,6 +347,12 @@ def create_remote_directory(directory: str | list, cluster_config: dict):
     if isinstance(directory, str):
         directory = [directory]
 
+    # Get unmounted path of all directories
+    directory = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
+        for dir_path in directory
+    ]
+
     if cluster_config.get('executor') != 'slurm':
         tunnel = run.LocalTunnel(job_dir=directory[0])
         for dir_path in directory:
@@ -334,6 +430,13 @@ def check_remote_mount_directories(directories: list, cluster_config: dict, exit
         raise ValueError("Cluster config is not provided.")
     if isinstance(directories, str):
         directories = [directories]
+
+    # Get unmounted path of all directories
+    directories = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
+        for dir_path in directories
+    ]
+
     if cluster_config.get('executor') != 'slurm':
         tunnel = run.LocalTunnel(job_dir=None)
         all_dirs_exist = True
@@ -689,6 +792,68 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+# TODO: Unify the signature of generate.py to use this
+def configure_client(
+    *,  # Force keyword arguments
+    model: str,
+    server_type: str,
+    server_gpus: int,
+    server_nodes: int,
+    server_address: str,
+    server_port: Optional[int],
+    server_args: str,
+    extra_arguments: str,
+    get_random_port: bool,
+):
+    """
+    Utility function to configure a client for the model inference server.
+
+    Args:
+        model: Mounted Path to the model to evaluate.
+        server_type: String name of the server type.
+        server_address: URL of the server hosting the model.
+        server_gpus: Number of GPUs to use for the server.
+        server_nodes: Number of nodes to use for the server.
+        server_port: Port number for the server.
+        server_args: Additional arguments for the server.
+        extra_arguments: Extra arguments to pass to the command.
+        get_random_port: Whether to get a random port for the server.
+
+    Returns:
+        A tuple containing:
+            - server_config: Configuration for the server.
+            - extra_arguments: Updated extra arguments for the command.
+            - server_address: Address of the server.
+            - server_port: Port number for the server.
+    """
+    if server_address is None:  # we need to host the model
+        if server_port is None:  # if not specified, we will use a random port or 5000 depending get_random_port
+            server_port = get_free_port(strategy="random") if not get_random_port else 5000
+        assert server_gpus is not None, "Need to specify server_gpus if hosting the model"
+        server_address = f"localhost:{server_port}"
+
+        server_config = {
+            "model_path": model,
+            "server_type": server_type,
+            "num_gpus": server_gpus,
+            "num_nodes": server_nodes,
+            "server_args": server_args,
+            "server_port": server_port,
+        }
+        extra_arguments = (
+            f"{extra_arguments} ++server.server_type={server_type} "
+            f"++server.host=localhost ++server.port={server_port} "
+        )
+    else:  # model is hosted elsewhere
+        server_config = None
+        extra_arguments = (
+            f"{extra_arguments} ++server.server_type={server_type} "
+            f"++server.base_url={server_address} ++server.model={model} "
+        )
+        server_port = None
+    return server_config, extra_arguments, server_address, server_port
+
+
 @dataclass(kw_only=True)
 class CustomJobDetails(SlurmJobDetails):
     # we have 1 srun per sub-task (e.g. server/sandbox/main), but only a single sbatch
@@ -782,6 +947,10 @@ def get_cluster_config(cluster=None, config_dir=None):
 
     cluster_config = read_config(config_file)
 
+    # resolve ssh tunnel config
+    if "ssh_tunnel" in cluster_config:
+        cluster_config = update_ssh_tunnel_config(cluster_config)
+
     if cluster_config['executor'] == 'slurm' and "ssh_tunnel" not in cluster_config:
         if "job_dir" not in cluster_config:
             raise ValueError("job_dir must be provided in the cluster config if ssh_tunnel is not provided.")
@@ -790,6 +959,49 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    resolve_map = [
+        dict(key='user', default_env_key='USER'),
+        dict(key='job_dir', default_env_key=None),
+        dict(key='identity', default_env_key=None),
+    ]
+
+    for item in resolve_map:
+        key = item['key']
+        default_env_key = item['default_env_key']
+
+        if key in cluster_config['ssh_tunnel']:
+            # Resolve `user` from env if not provided
+            if cluster_config['ssh_tunnel'][key] is None and default_env_key is not None:
+                cluster_config['ssh_tunnel'][key] = os.environ[default_env_key]
+                LOG.info(f"Resolved `{key}` to `{cluster_config['ssh_tunnel'][key]}`")
+
+            elif isinstance(cluster_config['ssh_tunnel'][key], str) and '$' in cluster_config['ssh_tunnel'][key]:
+                cluster_config['ssh_tunnel'][key] = os.path.expandvars(cluster_config['ssh_tunnel'][key])
+                LOG.info(f"Resolved `{key}` to `{cluster_config['ssh_tunnel'][key]}`")
+
+    if "$" in cluster_config['ssh_tunnel']['identity']:
+        raise ValueError(
+            "SSH identity cannot be resolved from environment variables. "
+            "Please provide a valid path to the identity file."
+        )
+
+    return cluster_config
+
+
 @lru_cache
 def _get_tunnel_cached(
     job_dir: str,
@@ -799,23 +1011,6 @@ def _get_tunnel_cached(
     shell: str | None = None,
     pre_command: str | None = None,
 ):
-    # If the str provided is an env variable, resolve it
-    if "$" in job_dir:
-        job_dir = os.path.expandvars(job_dir)
-        LOG.info(f"Resolved `job_dir` from env var to `{job_dir}`")
-
-    if "$" in host:
-        host = os.path.expandvars(host)
-        LOG.info(f"Resolved `host` from env var to `{host}`")
-
-    if "$" in user:
-        user = os.path.expandvars(user)
-        LOG.info(f"Resolved `user` from env var to `{user}`")
-
-    if "$" in identity:
-        identity = os.path.expandvars(identity)
-        LOG.info(f"Resolved `identity` from env var to `{identity}`")
-
     return run.SSHTunnel(
         host=host,
         user=user,
@@ -1121,6 +1316,28 @@ def get_env_variables(cluster_config):
     return env_vars
 
 
+def wrap_cmd(cmd, preprocess_cmd, postprocess_cmd, random_seed=None, wandb_parameters=None):
+    if preprocess_cmd:
+        if random_seed is not None:
+            preprocess_cmd = preprocess_cmd.format(random_seed=random_seed)
+        cmd = f" {preprocess_cmd} && {cmd} "
+    if postprocess_cmd:
+        if random_seed is not None:
+            postprocess_cmd = postprocess_cmd.format(random_seed=random_seed)
+        cmd = f" {cmd} && {postprocess_cmd} "
+    if wandb_parameters:
+        log_wandb_cmd = (
+            f"python -m nemo_skills.inference.log_samples_wandb "
+            f"    {wandb_parameters['samples_file']} "
+            f"    --name={wandb_parameters['name']} "
+            f"    --project={wandb_parameters['project']} "
+        )
+        if wandb_parameters['group'] is not None:
+            log_wandb_cmd += f" --group={wandb_parameters['group']} "
+        cmd = f"{cmd} && {log_wandb_cmd} "
+    return cmd
+
+
 def get_mounts_from_config(cluster_config: dict):
     """
     Determines if there are mount paths that are being passed via environment variables.
diff --git a/nemo_skills/utils.py b/nemo_skills/utils.py
index c2c6850901..de733ab05c 100644
--- a/nemo_skills/utils.py
+++ b/nemo_skills/utils.py
@@ -23,7 +23,7 @@
 import typing
 from dataclasses import MISSING, dataclass, fields, is_dataclass
 from pathlib import Path
-from typing import Any, List, Optional
+from typing import Any, Callable, List, Optional, Union
 
 import fire
 from fire import decorators as fire_decorators
@@ -508,3 +508,44 @@ def resolve_python_module_from_file(py_filepath: str, root_module: str = 'nemo_s
     striped_module = str(stripped_module_path).replace("/", ".")
     striped_module = os.path.splitext(striped_module)[0]
     return striped_module
+
+
+def maybe_get_env(value: Union[Any, List[Any]], env_name, default=None, cast: Callable[[Any], Any] = None):
+    """
+    Retrieve the value of an environment variable, iff the value is originally None.
+
+    If the provided value is None, this function attempts to retrieve the value
+    from the environment variables specified by `env_name`. If none of the
+    environment variables are set, it returns the provided default value.
+
+    Args:
+        value: The initial value to check or a list of values to check.
+        env_name (str or list of str): The name(s) of the environment variable(s) to check.
+        default: The default value to return if the environment variable is not set.
+        cast: A function to cast the environment variable to a specific type. If None, no casting is performed.
+
+    Returns:
+        The value of the environment variable if set, otherwise the default value.
+    """
+    if value is None:
+        # Convert env_name to a list if it's a string
+        if isinstance(env_name, str):
+            env_name = [env_name]
+
+        found_value = False
+        # Iterate over the list of environment variable names
+        for env in env_name:
+            # Try to get the value from the environment
+            value = os.environ.get(env, None)
+            if value is not None:
+                # Cast the value if a casting function is provided
+                if cast is not None:
+                    value = cast(value)
+
+                found_value = True
+                break
+
+        # If no environment variable is found, use the default value
+        if not found_value:
+            value = default
+    return value


Reviews:
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T05:20:39Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:23:21Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T04:50:21Z
user: login: Kipok
body: Some comments + tests currently fail, so would need to fix those
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T22:53:09Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T03:20:57Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T00:42:41Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T02:26:27Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T22:43:20Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:06:41Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:10:57Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:11:37Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:12:52Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:12:58Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:13:37Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:16:26Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:17:01Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:18:18Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:18:37Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:20:39Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:21:01Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:29:01Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:35:27Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:37:42Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T23:43:09Z
user: login: Kipok
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T00:43:33Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T02:33:11Z
user: login: Kipok
body: Thanks! Please wait for cpu tests to finish and feel free to merge after that
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-23T03:06:55Z
user: login: titu1994
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T21:45:49Z

Review comments:
diff_hunk: @@ -76,32 +110,87 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
path: nemo_skills/pipeline/run_cmd.py
user: login: titu1994
body: Btw, id like this flag to be true in general to offer users some safety. But that can be discussed. 
created_at: 2025-05-09T05:20:38Z
author_association: COLLABORATOR
diff_hunk: @@ -76,32 +110,87 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: The problem with having it True is that when you schedule large pipelines, it's often the case that subsequent steps take as input files generated by previous steps. So at scheduling time, they don't exist and thus will cause failure. But they will be available by the time the job runs because of slurm dependencies
created_at: 2025-05-09T20:23:21Z
author_association: COLLABORATOR
diff_hunk: @@ -76,32 +110,87 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
path: nemo_skills/pipeline/run_cmd.py
user: login: titu1994
body: I see, that's a fair point to keep it disabled by default. 
created_at: 2025-05-10T04:50:20Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
-    check_if_mounted(cluster_config, output_file)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    if check_mounted_paths:
+        if not is_mounted_filepath(cluster_config, input_file):
+            input_dir, input_filename = os.path.split(input_file)
+            add_mount_path(input_dir, "/mounted_data/input", cluster_config)
+
+        if not is_mounted_filepath(cluster_config, output_file):
+            output_dir, output_filename = os.path.split(output_file)
+            add_mount_path(output_dir, "/mounted_data/output", cluster_config)
+
+    input_file = get_mounted_path(cluster_config, input_file)
+    output_file = get_mounted_path(cluster_config, output_file)
+
     if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+        if check_mounted_paths: create_remote_directory(log_dir, cluster_config)
path: nemo_skills/pipeline/check_contamination.py
user: login: Kipok
body: ```suggestion
        if check_mounted_paths: 
            create_remote_directory(log_dir, cluster_config)
```
created_at: 2025-05-12T22:17:32Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
-    check_if_mounted(cluster_config, output_file)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    if check_mounted_paths:
+        if not is_mounted_filepath(cluster_config, input_file):
+            input_dir, input_filename = os.path.split(input_file)
+            add_mount_path(input_dir, "/mounted_data/input", cluster_config)
+
+        if not is_mounted_filepath(cluster_config, output_file):
+            output_dir, output_filename = os.path.split(output_file)
+            add_mount_path(output_dir, "/mounted_data/output", cluster_config)
+
+    input_file = get_mounted_path(cluster_config, input_file)
+    output_file = get_mounted_path(cluster_config, output_file)
+
     if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+        if check_mounted_paths: create_remote_directory(log_dir, cluster_config)
+        log_dir = get_mounted_path(cluster_config, log_dir)
+
+    if check_mounted_paths:
+        # Final check for existance of mounted paths
path: nemo_skills/pipeline/check_contamination.py
user: login: Kipok
body: ```suggestion
        # Final check for existence of mounted paths
```
created_at: 2025-05-12T22:18:32Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
path: nemo_skills/pipeline/check_contamination.py
user: login: Kipok
body: Can we keep this logic in the "else" part of `if check_mounted_paths`?
created_at: 2025-05-12T22:19:37Z
author_association: COLLABORATOR
diff_hunk: @@ -255,6 +267,12 @@ def create_remote_directory(directory: str | list, cluster_config: dict):
     if isinstance(directory, str):
         directory = [directory]
 
+    # Get unmounted path of all directories
+    directory = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: isn't the if part redundant? I think the `get_unmounted_path` already returns original if it's not mounted
created_at: 2025-05-12T22:21:29Z
author_association: COLLABORATOR
diff_hunk: @@ -332,6 +350,13 @@ def check_remote_mount_directories(directories: list, cluster_config: dict, exit
         raise ValueError("Cluster config is not provided.")
     if isinstance(directories, str):
         directories = [directories]
+
+    # Get unmounted path of all directories
+    directories = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: isn't the if part redundant? I think the get_unmounted_path already returns original if it's not mounted
created_at: 2025-05-12T22:21:46Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
path: nemo_skills/pipeline/check_contamination.py
user: login: Kipok
body: This applies to all pipeline scripts
created_at: 2025-05-12T22:22:46Z
author_association: COLLABORATOR
diff_hunk: @@ -219,14 +233,28 @@ def convert(
     if convert_to in ["hf", "nemo"] and model_type == "deepseek_v3":
         raise ValueError("Conversion to HF/Nemo is not yet supported for DeepSeek v3 models")
 
+    # Prepare cluster config and mount paths
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_model)
-    check_if_mounted(cluster_config, output_model)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
+
+    if check_mounted_paths:
+        if not is_mounted_filepath(cluster_config, input_model): add_mount_path(input_model, "/input_model", cluster_config)
path: nemo_skills/pipeline/convert.py
user: login: Kipok
body: best not to use single line ifs, harder to read and also our current black setup in precommit will reformat this back to 2 lines anyway. And in general, best to setup precommit to ensure your code is automatically formatted
created_at: 2025-05-12T22:23:54Z
author_association: COLLABORATOR
diff_hunk: @@ -197,12 +211,25 @@ def eval(
         pass
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, output_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
path: nemo_skills/pipeline/eval.py
user: login: Kipok
body: Feels like we can have a utility that does all of this and takes all boolean arguments as well as all files to check and log_dir and then just call it as one line in all pipelines to avoid duplication
created_at: 2025-05-12T22:25:51Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+    if 'job_dir' in cluster_config['ssh_tunnel']:
+        # Resolve `job_dir` from env if not provided
+        cluster_config['ssh_tunnel']['job_dir'] = os.path.expandvars(cluster_config['ssh_tunnel']['job_dir'])
+        LOG.info(f"Resolved `job_dir` to `{cluster_config['ssh_tunnel']['job_dir']}`")
+
+    if 'identity' in cluster_config['ssh_tunnel']:
+        # Resolve `identity` from env if not provided
+        cluster_config['ssh_tunnel']['identity'] = os.path.expandvars(cluster_config['ssh_tunnel']['identity'])
+        LOG.info(f"Resolved `identity` to `{cluster_config['ssh_tunnel']['identity']}`")
+
+        if "$" in cluster_config['ssh_tunnel']['identity']:
+            raise ValueError(
+                f"Identity file path could not be loaded from env variable `{cluster_config['ssh_tunnel']['identity']}`"
+            )
+
+    # Add mount container list if not present
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: feels out of place in this function as from the name it only supposed to update ssh section
created_at: 2025-05-12T22:31:54Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: can we just have some recursive logic to do this for all variables, not just ssh tunnel? That would make things more consistent and maybe there is already some library we can reuse that does that. But if it's hard, ok to keep it like this
created_at: 2025-05-12T22:33:25Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: and if not, then at least this function can be much shorter by having a for loop over all variables and same logic inside the loop, instead of current copy-paste
created_at: 2025-05-12T22:34:01Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+    if 'job_dir' in cluster_config['ssh_tunnel']:
+        # Resolve `job_dir` from env if not provided
+        cluster_config['ssh_tunnel']['job_dir'] = os.path.expandvars(cluster_config['ssh_tunnel']['job_dir'])
+        LOG.info(f"Resolved `job_dir` to `{cluster_config['ssh_tunnel']['job_dir']}`")
+
+    if 'identity' in cluster_config['ssh_tunnel']:
+        # Resolve `identity` from env if not provided
+        cluster_config['ssh_tunnel']['identity'] = os.path.expandvars(cluster_config['ssh_tunnel']['identity'])
+        LOG.info(f"Resolved `identity` to `{cluster_config['ssh_tunnel']['identity']}`")
+
+        if "$" in cluster_config['ssh_tunnel']['identity']:
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: why is special logic needed here?
created_at: 2025-05-12T22:34:27Z
author_association: COLLABORATOR
diff_hunk: @@ -500,3 +500,44 @@ def resolve_python_module_from_file(py_filepath: str, root_module: str = 'nemo_s
     striped_module = str(stripped_module_path).replace("/", ".")
     striped_module = os.path.splitext(striped_module)[0]
     return striped_module
+
+
+def maybe_get_env(value: Union[int, List[int]], env_name, default=None, cast: Callable[[Any], Any] = None):
path: nemo_skills/utils.py
user: login: Kipok
body: type of value is not correct? Best not to provide type, rather than have incorrect one
created_at: 2025-05-12T22:40:31Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: Probably best to keep this in utils as well (we really should split that file though) and then make use of it inside generate.py, eval.py, start_server.py 
created_at: 2025-05-12T22:46:46Z
author_association: COLLABORATOR
diff_hunk: @@ -13,21 +13,36 @@
 # limitations under the License.
 
 import logging
-from typing import List
+from typing import List, Optional
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: Will discuss changes to this file offline
created_at: 2025-05-12T22:52:35Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
path: nemo_skills/pipeline/check_contamination.py
user: login: titu1994
body: Yes will do
created_at: 2025-05-13T02:42:36Z
author_association: COLLABORATOR
diff_hunk: @@ -197,12 +211,25 @@ def eval(
         pass
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, output_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
path: nemo_skills/pipeline/eval.py
user: login: titu1994
body: The issue with that is not all scripts have all args. For example, log dir is sometimes defaulted to output dir, vs sometimes explicitly required in run_cmd vs sometimes ignored entirely (in start server). 

Keeping resolve mount paths as a separate arg allows users in scripts to dynamically mount things as needed, which is convenient without modifying cluster config again and again. Especially for long scripts which start with one shared cluster config but dynamically mount different things for each call to cli functions (say different set of mounts, mount previous runs output folder etc). 

Im open to ideas of how to unify it though, but dunno if it needs to be part of this pr.
created_at: 2025-05-13T02:46:53Z
author_association: COLLABORATOR
diff_hunk: @@ -255,6 +267,12 @@ def create_remote_directory(directory: str | list, cluster_config: dict):
     if isinstance(directory, str):
         directory = [directory]
 
+    # Get unmounted path of all directories
+    directory = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: So the thing is get_unmounted_path will raise an error if the file was not in the mounted list at all. This is to raise awareness that a path was passed to this function that has not been registered for mounting. The two cases to use this function are
- Pass a mounted path, rethrn its unmounted path after resolution
- Pass an unmounted path (which has a corresponding registered mount), return as is 
It should error out if an unmounted path is provided which has no corresponding mount path - that means that in the future the mount path wont be available on cluster.

created_at: 2025-05-13T02:55:06Z
author_association: COLLABORATOR
diff_hunk: @@ -332,6 +350,13 @@ def check_remote_mount_directories(directories: list, cluster_config: dict, exit
         raise ValueError("Cluster config is not provided.")
     if isinstance(directories, str):
         directories = [directories]
+
+    # Get unmounted path of all directories
+    directories = [
+        get_unmounted_path(cluster_config, dir_path) if is_mounted_filepath(cluster_config, dir_path) else dir_path
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: See above
created_at: 2025-05-13T02:55:18Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Do you mean nemo_skills/utils.py? Cause this is skills/pipeline/utils.py
Agreed to split it. With this PR, a vast majority of codegens utils should be upstreamed, so it should be ok to start splitting the file.
created_at: 2025-05-13T02:56:44Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Agreed, will look into that 
created_at: 2025-05-13T02:57:07Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+    if 'job_dir' in cluster_config['ssh_tunnel']:
+        # Resolve `job_dir` from env if not provided
+        cluster_config['ssh_tunnel']['job_dir'] = os.path.expandvars(cluster_config['ssh_tunnel']['job_dir'])
+        LOG.info(f"Resolved `job_dir` to `{cluster_config['ssh_tunnel']['job_dir']}`")
+
+    if 'identity' in cluster_config['ssh_tunnel']:
+        # Resolve `identity` from env if not provided
+        cluster_config['ssh_tunnel']['identity'] = os.path.expandvars(cluster_config['ssh_tunnel']['identity'])
+        LOG.info(f"Resolved `identity` to `{cluster_config['ssh_tunnel']['identity']}`")
+
+        if "$" in cluster_config['ssh_tunnel']['identity']:
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: identity points to the ssh public key. If after attempting resolution, the resolution fails, then passing it to ssh tunnel gives a vague and non descript error that is hard to debug. So i special cased this part. 
created_at: 2025-05-13T02:58:27Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+    if 'job_dir' in cluster_config['ssh_tunnel']:
+        # Resolve `job_dir` from env if not provided
+        cluster_config['ssh_tunnel']['job_dir'] = os.path.expandvars(cluster_config['ssh_tunnel']['job_dir'])
+        LOG.info(f"Resolved `job_dir` to `{cluster_config['ssh_tunnel']['job_dir']}`")
+
+    if 'identity' in cluster_config['ssh_tunnel']:
+        # Resolve `identity` from env if not provided
+        cluster_config['ssh_tunnel']['identity'] = os.path.expandvars(cluster_config['ssh_tunnel']['identity'])
+        LOG.info(f"Resolved `identity` to `{cluster_config['ssh_tunnel']['identity']}`")
+
+        if "$" in cluster_config['ssh_tunnel']['identity']:
+            raise ValueError(
+                f"Identity file path could not be loaded from env variable `{cluster_config['ssh_tunnel']['identity']}`"
+            )
+
+    # Add mount container list if not present
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Ah its a vestige from some older code, will remove
created_at: 2025-05-13T02:58:42Z
author_association: COLLABORATOR
diff_hunk: @@ -500,3 +500,44 @@ def resolve_python_module_from_file(py_filepath: str, root_module: str = 'nemo_s
     striped_module = str(stripped_module_path).replace("/", ".")
     striped_module = os.path.splitext(striped_module)[0]
     return striped_module
+
+
+def maybe_get_env(value: Union[int, List[int]], env_name, default=None, cast: Callable[[Any], Any] = None):
path: nemo_skills/utils.py
user: login: titu1994
body: Updated
created_at: 2025-05-13T02:59:25Z
author_association: COLLABORATOR
diff_hunk: @@ -121,10 +130,28 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
path: nemo_skills/pipeline/check_contamination.py
user: login: titu1994
body: Just a note, in places where `get_mounted_path` or `get_unmounted_path` exists, it implicitly calls check_if_mounted() - or it will raise an error saying the path is not mounted. So i removed the check_if_mounted() cause it was redundant in those sections
created_at: 2025-05-13T03:04:03Z
author_association: COLLABORATOR
diff_hunk: @@ -758,6 +848,52 @@ def get_cluster_config(cluster=None, config_dir=None):
     return cluster_config
 
 
+def update_ssh_tunnel_config(cluster_config: dict):
+    """
+    Update the ssh tunnel configuration in the cluster config to resolve job dir and username.
+    uses the `user` information to populate `job_dir` if in config.
+
+    Args:
+        cluster_config: dict: The cluster configuration dictionary
+
+    Returns:
+        dict: The updated cluster configuration dictionary
+    """
+    if 'ssh_tunnel' not in cluster_config:
+        return cluster_config
+
+    if 'user' in cluster_config['ssh_tunnel']:
+        # Resolve `user` from env if not provided
+        if cluster_config['ssh_tunnel']['user'] is None:
+            cluster_config['ssh_tunnel']['user'] = os.environ['USER']
+            LOG.info(f"Resolved `user` to `{cluster_config['ssh_tunnel']['user']}`")
+
+        elif isinstance(cluster_config['ssh_tunnel']['user'], str) and '$' in cluster_config['ssh_tunnel']['user']:
+            cluster_config['ssh_tunnel']['user'] = os.path.expandvars(cluster_config['ssh_tunnel']['user'])
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Done
created_at: 2025-05-13T03:18:26Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +45,7 @@ def run_cmd(
         help="One of the configs inside config_dir or NEMO_SKILLS_CONFIG_DIR or ./cluster_configs. "
         "Can also use NEMO_SKILLS_CONFIG instead of specifying as argument.",
     ),
+    command: str = typer.Option(None, help="Command to run in the container"),
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: ```suggestion
    command: str = typer.Option(None, help="Command to run in the container. Can also be specified as extra arguments."),
```
created_at: 2025-05-22T00:15:41Z
author_association: COLLABORATOR
diff_hunk: @@ -63,9 +77,8 @@ def run_cmd(
         help="If specified, will reuse the code from this experiment. "
         "Can provide an experiment name or an experiment object if running from code.",
     ),
-    preprocess_cmd: str = typer.Option(None, help="Command to run before job"),
-    postprocess_cmd: str = typer.Option(None, help="Command to run after job"),
     config_dir: str = typer.Option(None, help="Can customize where we search for cluster configs"),
+    with_sandbox: bool = typer.Option(False, help="Whether to use the sandboxing feature to host the model"),
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: ```suggestion
    with_sandbox: bool = typer.Option(False, help="If True, will start a sandbox container alongside this job"),
```
created_at: 2025-05-22T00:16:24Z
author_association: COLLABORATOR
diff_hunk: @@ -13,21 +13,26 @@
 # limitations under the License.
 
 import logging
-from typing import List
+from typing import List, Optional
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: Would be great to update documentation over here https://github.com/NVIDIA/NeMo-Skills/blob/main/docs/pipelines/run-cmd.md to explain the new server arguments and how the command interacts with them
created_at: 2025-05-22T00:18:24Z
author_association: COLLABORATOR
diff_hunk: @@ -464,6 +448,7 @@ def generate(
         "If not specified, will use the default entrypoint for the server type.",
     ),
     dependent_jobs: int = typer.Option(0, help="Specify this to launch that number of dependent jobs"),
+    mount_paths: str = typer.Option(None, help="Comma separated list of paths to mount on the remote machine"),
path: nemo_skills/pipeline/generate.py
user: login: Kipok
body: Please add some description of these new common parameters over here https://github.com/NVIDIA/NeMo-Skills/blob/main/docs/pipelines/index.md?plain=1#L76
created_at: 2025-05-22T00:31:57Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: oh, but there is still configure_client in pipeline/generate. Should that one be removed then? Also can it be directly used in eval.py and check_contamination.py?
created_at: 2025-05-22T00:36:56Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Hmm that will requite a more refactors. I think it should be a drop in replacement, however the utils one has extra arg `get_random_port` and the generate one has extra unused arg (generation type and server port). So Id like to keep this as a separate pr item (or you can update either later on).
created_at: 2025-05-22T02:20:23Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: But yes technically all 3 server configure calls can be merged into the utils one (if it accepts the union of all extra args from each script?)
created_at: 2025-05-22T02:26:19Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command:
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=/hf_models/Qwen3-1.7B \
path: docs/pipelines/run-cmd.md
user: login: Kipok
body: thanks, great example! Consider replacing this with `Qwen/Qwen3-1.7B` which will make it runnable without a prerequisite for having a model downloaded
```suggestion
    --model=Qwen/Qwen3-1.7B \
```
created_at: 2025-05-22T22:24:46Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command:
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=/hf_models/Qwen3-1.7B \
+    --server_type=vllm \
+    --server_gpus=1 \
+    --with_sandbox \
path: docs/pipelines/run-cmd.md
user: login: Kipok
body: ```suggestion
```
created_at: 2025-05-22T22:25:38Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
path: docs/pipelines/run-cmd.md
user: login: Kipok
body: would be nice to use this example (or something similar) in the gpu tests to make sure this new functionality is checked in ci. Can be done in another PR, but please create an issue to track this if not including here
created_at: 2025-05-22T22:28:11Z
author_association: COLLABORATOR
diff_hunk: @@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
+  - Check if the paths specified in the script are mounted correctly. This is useful if you want to make 
path: docs/pipelines/index.md
user: login: Kipok
body: this doesn't display properly (not a nested list). You can run `mkdocs serve` to get the docs built locally and see how it's visible on the website
created_at: 2025-05-22T22:29:44Z
author_association: COLLABORATOR
diff_hunk: @@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
path: docs/pipelines/index.md
user: login: Kipok
body: this description is a little confusing to me. E.g. sometimes it will raise an error if the path doesn't exist, but sometimes it will create a folder for me. How do I know when it's one or the other?
created_at: 2025-05-22T22:30:52Z
author_association: COLLABORATOR
diff_hunk: @@ -126,10 +135,14 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
-    check_if_mounted(cluster_config, output_file)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
path: nemo_skills/pipeline/check_contamination.py
user: login: Kipok
body: maybe make it a part of get_cluster_config (if mount_paths specified we resolve them)?
created_at: 2025-05-22T22:32:56Z
author_association: COLLABORATOR
diff_hunk: @@ -76,32 +99,80 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
path: nemo_skills/pipeline/run_cmd.py
user: login: Kipok
body: Why do we need this parameter? E.g. all other scripts don't have it and I think the default logic of doing

```
get_random_port = server_gpus != 8 and not exclusive
```

covers majority of possible cases? If you really think we should have it, we should then add it to all scripts, not just this one
created_at: 2025-05-22T22:35:40Z
author_association: COLLABORATOR
diff_hunk: @@ -62,20 +75,28 @@ def start_server(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
     get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
 ):
     """Self-host a model server."""
     setup_logging(disable_hydra_logs=False, use_rich=True)
 
     cluster_config = get_cluster_config(cluster, config_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
 
     try:
         server_type = server_type.value
     except AttributeError:
         pass
 
+    log_dir = check_remote_mounts(cluster_config, log_dir, check_mounted_paths=check_mounted_paths)
+
     if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+        if check_mounted_paths:
path: nemo_skills/pipeline/start_server.py
user: login: Kipok
body: isn't this done in check_remote_mounts already?
created_at: 2025-05-22T22:37:05Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: Better to name it just check_mounts as this is also applicable when running locally and specifying config file
created_at: 2025-05-22T22:38:11Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: this will fail on any files without extension? Can also send a proper check command through the ssh tunnel, but if you think this is enough, I guess it's ok. But can lead to silent issues later on
created_at: 2025-05-22T22:40:13Z
author_association: COLLABORATOR
diff_hunk: @@ -657,6 +682,67 @@ def get_sandox_command(cluster_config):
     return "/entrypoint.sh && /start.sh"
 
 
+def configure_client(
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: ok, please add a TODO or an issue for this
created_at: 2025-05-22T22:41:43Z
author_association: COLLABORATOR
diff_hunk: @@ -782,6 +944,10 @@ def get_cluster_config(cluster=None, config_dir=None):
 
     cluster_config = read_config(config_file)
 
+    # resolve ssh tunnel config
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: should we move this to get_cluster_config function?
created_at: 2025-05-22T22:42:20Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command:
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=/hf_models/Qwen3-1.7B \
+    --server_type=vllm \
+    --server_gpus=1 \
+    --with_sandbox \
path: docs/pipelines/run-cmd.md
user: login: titu1994
body: This is intented to show that you can add sandbox with one flag.
created_at: 2025-05-22T23:06:41Z
author_association: COLLABORATOR
diff_hunk: @@ -126,10 +135,14 @@ def check_contamination(
     get_random_port = server_gpus != 8 and not exclusive
 
     cluster_config = get_cluster_config(cluster, config_dir)
-    check_if_mounted(cluster_config, input_file)
-    check_if_mounted(cluster_config, output_file)
-    if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
path: nemo_skills/pipeline/check_contamination.py
user: login: titu1994
body: This needs to be separate. Otherwise we have to go back and update all get_cluster_config with optional arg resolve mount paths. And on top of that, in python mode, you should call get cluster config for "base config" and then run custom dicts / str of resolve mount paths per task you're executing
created_at: 2025-05-22T23:10:57Z
author_association: COLLABORATOR
diff_hunk: @@ -76,32 +99,80 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
path: nemo_skills/pipeline/run_cmd.py
user: login: titu1994
body: In run cmd, i want it to be fully flexible for dev to control their env. The above logic is fine for our slurm env, may not be the case for others.
created_at: 2025-05-22T23:11:37Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Done
created_at: 2025-05-22T23:12:51Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command:
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=/hf_models/Qwen3-1.7B \
+    --server_type=vllm \
+    --server_gpus=1 \
+    --with_sandbox \
path: docs/pipelines/run-cmd.md
user: login: Kipok
body: Then it's better to add some explanation. Otherwise it appears as if this is required for the example you have (while it's actually not used at all). I'd rather not have it in command and say in another sentence that you can do this
created_at: 2025-05-22T23:12:57Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Its meant to be a dir check - if is a file a filepath with extension or a dir path with no extension. I tested it in python 3.10 - works fine.
created_at: 2025-05-22T23:13:37Z
author_association: COLLABORATOR
diff_hunk: @@ -782,6 +944,10 @@ def get_cluster_config(cluster=None, config_dir=None):
 
     cluster_config = read_config(config_file)
 
+    # resolve ssh tunnel config
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: This is part of get_cluster_config() ?
created_at: 2025-05-22T23:16:26Z
author_association: COLLABORATOR
diff_hunk: @@ -62,20 +75,28 @@ def start_server(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    check_mounted_paths: bool = typer.Option(False, help="Check if mounted paths are available on the remote machine"),
     get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
 ):
     """Self-host a model server."""
     setup_logging(disable_hydra_logs=False, use_rich=True)
 
     cluster_config = get_cluster_config(cluster, config_dir)
+    cluster_config = resolve_mount_paths(cluster_config, mount_paths)
 
     try:
         server_type = server_type.value
     except AttributeError:
         pass
 
+    log_dir = check_remote_mounts(cluster_config, log_dir, check_mounted_paths=check_mounted_paths)
+
     if log_dir:
-        check_if_mounted(cluster_config, log_dir)
+        if check_mounted_paths:
path: nemo_skills/pipeline/start_server.py
user: login: titu1994
body: Good catch, thanks !
created_at: 2025-05-22T23:17:01Z
author_association: COLLABORATOR
diff_hunk: @@ -76,32 +99,80 @@ def run_cmd(
         "--not_exclusive",
         help="If --not_exclusive is used, will NOT use --exclusive flag for slurm",
     ),
+    get_random_port: bool = typer.Option(False, help="If True, will get a random port for the server"),
path: nemo_skills/pipeline/run_cmd.py
user: login: titu1994
body: As to other scripts - it depends. Do we want users to bother with setting this for the most common use cases (like our clusters?) - probably not. So I'm in favor of keeping those automatic. This one is too generic of a script - broader responsibility than just our clusters. So it makes sense to be more flexible in this one.
created_at: 2025-05-22T23:18:17Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: right, but this will treat any files without extension as a directory and try to create it. Then it's unclear how the downstream code will interact with that and might lead to silent errors. We don't normally use file names without extensions, but it's certainly possible to do that, right?
created_at: 2025-05-22T23:18:37Z
author_association: COLLABORATOR
diff_hunk: @@ -782,6 +944,10 @@ def get_cluster_config(cluster=None, config_dir=None):
 
     cluster_config = read_config(config_file)
 
+    # resolve ssh tunnel config
path: nemo_skills/pipeline/utils.py
user: login: Kipok
body: oh, sorry, you're right
created_at: 2025-05-22T23:20:39Z
author_association: COLLABORATOR
diff_hunk: @@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
+  - Check if the paths specified in the script are mounted correctly. This is useful if you want to make 
path: docs/pipelines/index.md
user: login: titu1994
body: Will do in another pr
created_at: 2025-05-22T23:21:01Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,49 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command:
+
+```bash
+ns run_cmd \
+    --cluster=local \
+    --model=/hf_models/Qwen3-1.7B \
+    --server_type=vllm \
+    --server_gpus=1 \
+    --with_sandbox \
path: docs/pipelines/run-cmd.md
user: login: titu1994
body: Done
created_at: 2025-05-22T23:29:01Z
author_association: COLLABORATOR
diff_hunk: @@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
path: docs/pipelines/index.md
user: login: titu1994
body: Will fix in a second pr for nits. That wont need GPU exec hopefully. 
created_at: 2025-05-22T23:35:27Z
author_association: COLLABORATOR
diff_hunk: @@ -88,6 +88,18 @@ All scripts inside pipeline folder have the following parameters.
 - **--run_after**: Can be used in conjunction with `--expname` to chain jobs to
   run one after another (only applicable on slurm). E.g. run training job with
   `--expname my-training-run` and then launch an eval with `--run_after my-training-run`.
+- **--mount_paths**: Can be used to mount additional paths to the cluster config dynamically. 
+  This is useful if you want to access some data that is not mounted in cluster config. E.g. use
+  `--mount_paths /my/remote/workspace:/workspace` to mount `/workspace` folder from the host
+  machine to the slurm job.
+- **--check_mounted_paths**: This flag offers a few different capabilities for convenience:
path: docs/pipelines/index.md
user: login: titu1994
body: The nain thing is `--check_mount_paths` is the flag that enables remote file and dir checks. When not set, we do just local checks of whether the path "looks correct" and mounted. When this flag is set, we confirm with remote calls to the cluster, that indeed the paths are in correct location or not. 

This is set to false by default because it is fine for one off scripts but in a complex chain of scripts - you don't want to assert on existence of data when a previous pipeline stage is going to generate that data file.
created_at: 2025-05-22T23:37:42Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Yes possible to do that.. Hmm this part is actually only necessary for 2 scripts that try to create output directory. If i manually do it in those two scripts, then i can skip this logic here. Let me change that in another pr.
created_at: 2025-05-22T23:43:09Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,54 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command as below. We will also show how to start the code 
+execution server that can be used to run arbitrary code in a sandboxed environment. This can be useful to safely 
+execute untrusted code or to run code that requires a specific environment in a container without affecting the host 
path: docs/pipelines/run-cmd.md
user: login: Kipok
body: I think we need to tone this down, our sandbox is certainly not safe to run untrusted code, and we shouldn't say it in the docs :) It's a little safer than running it without sandbox, but not safe for sure (which we mention here https://nvidia.github.io/NeMo-Skills/basics/sandbox/)
created_at: 2025-05-23T00:43:33Z
author_association: COLLABORATOR
diff_hunk: @@ -14,4 +14,54 @@ code and tries to install it (will finish with error if not running from NeMo-Sk
 ns run_cmd --cluster=local cd /nemo_run/code/ && pip install -e .
 ```
 
-There are many more examples of how to use `run_cmd` throughout our documentation.
\ No newline at end of file
+There are many more examples of how to use `run_cmd` throughout our documentation.
+
+## LLM Server and Sandbox Server
+
+While we can run arbitrary commands with the default `run_cmd` script, we also provide the ability to
+run a LLM server with the `--model` argument and a few extra arguments for the server config. These arguments 
+are similar to the ones used for `start_server` script.
+
+This can be useful to run a server on a local machine or on a cluster with GPUs in a slurm job, while also being able to
+run arbirary code that uses LLM calls. 
+
+### Example
+
+Say you have the following inference file that uses OpenAI API with a vLLM backed server (say to run a 
+project that is compatible with OpenAI API). Imagine a file called `inference.py` with the following code:
+
+```python
+from openai import OpenAI
+client = OpenAI(api_key='EMPTY', base_url=f"http://0.0.0.0:5000/v1", timeout=None)
+api_model = client.models.list().data[0].id
+
+response = client.chat.completions.create(
+    model=api_model,
+    messages=[
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.7,
+    max_tokens=128,
+    top_p=0.95,
+    n=1,
+    stream=False,
+)
+print(response.choices[0].message.content)
+```
+
+Then we can run the server and the inference code in a single command as below. We will also show how to start the code 
+execution server that can be used to run arbitrary code in a sandboxed environment. This can be useful to safely 
+execute untrusted code or to run code that requires a specific environment in a container without affecting the host 
path: docs/pipelines/run-cmd.md
user: login: titu1994
body: Ah ok
created_at: 2025-05-23T02:33:11Z
author_association: COLLABORATOR
diff_hunk: @@ -131,6 +143,82 @@ def check_if_mounted(cluster_config, path_to_check):
         raise ValueError(f"The path '{path_to_check}' is not mounted. Check cluster config.")
 
 
+def check_remote_mounts(
+    cluster_config, log_dir: str, mount_map: Dict[str, Optional[str]] = None, check_mounted_paths: bool = False
+):
+    """
+    Utility method to check if the paths are mounted, whether the remote directories exist, create them if they dont exist
+    and finally resolve their mounted paths.
+
+    Args:
+        cluster_config: cluster config dictionary
+        log_dir: optional str representing the log directory
+        mount_map: a dictionary mapping the paths to their default mount locations.
+            Keys must be the mount source, and values must be the mount destination.
+            If the mount destination is None, the path will not be mounted but still checked and remote directory will still be created.
+            Mount destinations must be absolute paths and begin with '/'.
+        check_mounted_paths: if True, will perform remote calls to dynamically mount the provided paths, and create
+            remote directories if required. Finally, will assert that all remote paths are valid files or directories.
+
+    Returns:
+        tuple: a tuple of the mounted paths for the provided paths in mount_map, and the log_dir
+    """
+    # Check if mount map is provided
+    mount_map = mount_map or {}
+
+    # Compute directory of all files
+    remote_dir_list = []
+
+    # Check paths and add to mount list if not mounted
+    if check_mounted_paths:
+        for idx, (path, default_mount) in enumerate(mount_map.items()):
+            if not is_mounted_filepath(cluster_config, path):
+                # check if the path is a file or a directory
+                # so that the directory can be created
+                is_file = os.path.splitext(path) != ""
path: nemo_skills/pipeline/utils.py
user: login: titu1994
body: Made it just a check and raise error in another PR.
created_at: 2025-05-23T21:45:49Z
author_association: COLLABORATOR

PR comments:
user: login: titu1994
created_at: 2025-05-10T04:51:52Z
body: The other thing is I would normally like to do the check for model paths + auto Mount if not mounted, but HF support makes the code gnarly, so for now I'm not porting that from codegen. 
author_association: COLLABORATOR
user: login: titu1994
created_at: 2025-05-13T03:36:01Z
body: Addressed most of the comments
author_association: COLLABORATOR
user: login: titu1994
created_at: 2025-05-23T03:33:22Z
body: Thanks for the reviews and testing !
author_association: COLLABORATOR