diff --git a/dockerfiles/Dockerfile.sandbox b/dockerfiles/Dockerfile.sandbox index 3cccd8ee6c..ea1d35acc9 100644 --- a/dockerfiles/Dockerfile.sandbox +++ b/dockerfiles/Dockerfile.sandbox @@ -75,7 +75,7 @@ ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ # Install uv (adds to ~/.local/bin), then install deps -RUN if [ "$GITHUB_CI" != "1" ]; then \ +RUN if [ "$GITHUB_CI" != "1" ] && [ "$TARGETARCH" != "arm64" ]; then \ curl -LsSf https://astral.sh/uv/install.sh | sh && \ uv pip install --upgrade pip && \ uv pip install -r /app/stem_requirements.txt --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu; \ diff --git a/nemo_skills/pipeline/generate.py b/nemo_skills/pipeline/generate.py index 95e37abb59..f33796d05c 100644 --- a/nemo_skills/pipeline/generate.py +++ b/nemo_skills/pipeline/generate.py @@ -57,6 +57,7 @@ def _create_commandgroup_from_config( task_name: str, log_dir: str, sbatch_kwargs: Optional[Dict] = None, + sandbox_env_overrides: Optional[List[str]] = None, ) -> CommandGroup: """Create a CommandGroup from server_config. @@ -122,6 +123,14 @@ def _create_commandgroup_from_config( cmd, metadata = sandbox_command(cluster_config=cluster_config, port=sandbox_port) metadata["log_prefix"] = "sandbox" + # Apply user-specified environment overrides for the sandbox + if sandbox_env_overrides: + sandbox_env = metadata.get("environment", {}) + for override in sandbox_env_overrides: + key, value = override.split("=", 1) + sandbox_env[key] = value + metadata["environment"] = sandbox_env + sandbox_cmd = Command( command=cmd, container=cluster_config["containers"]["sandbox"], @@ -242,6 +251,11 @@ 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"), + sandbox_env_overrides: List[str] = typer.Option( + None, + help="Extra environment variables for the sandbox container in KEY=VALUE format. " + "E.g., --sandbox-env-overrides NEMO_SKILLS_SANDBOX_BLOCK_NETWORK=1 to enable network blocking.", + ), keep_mounts_for_sandbox: bool = typer.Option( False, help="If True, will keep the mounts for the sandbox container. Note that, it is risky given that sandbox executes LLM commands and could potentially lead to data loss. So, we advise not to use this unless absolutely necessary.", @@ -455,6 +469,7 @@ def generate( task_name=task_name, log_dir=log_dir, sbatch_kwargs=sbatch_kwargs, + sandbox_env_overrides=sandbox_env_overrides, ) # Use unique internal job name for dependency tracking, but same task_name