Remote Execution Skill

This skill handles running code on remote GPU or TPU clusters via SkyPilot.

1. Determine Target Device

Identify the target device from the user’s request:

Execution Instructions: Before running the launch script, you must find its absolute path. It is located in the scripts/ directory alongside this skill definition. Use your file search tools (e.g., glob or find) to locate launch_gpu.sh or launch_tpu.sh before executing it.

If the user does not specify a device, ask them which one to use.

2. Prerequisites

• The cluster must already be provisioned. Check that the corresponding cluster name file (.cluster_name_gpu or .cluster_name_tpu) exists and is non-empty in the project root.
• If the file does not exist or is empty, ask the user to provision a cluster first using the appropriate launch script.

3. Cluster Management

Provisioning

# Note: First locate the scripts as instructed above, then run them.
# GPU — common accelerator types: H100:1, A100:1, L4:1
bash <absolute_path_to_launch_gpu.sh> <accelerator_type> <experiment_name>

# TPU — common accelerator types: tpu-v4-8, tpu-v4-16, tpu-v6e-1, tpu-v6e-4
bash <absolute_path_to_launch_tpu.sh> <accelerator_type> <experiment_name>

The launch script automatically updates the corresponding .cluster_name_* file.

Teardown

# GPU
sky down $(cat .cluster_name_gpu) -y

# TPU
sky down $(cat .cluster_name_tpu) -y

4. Execution Command

GPU

sky exec $(cat .cluster_name_gpu) --workdir . "export CUDA_VISIBLE_DEVICES=0; uv run --extra gpu python <PATH_TO_SCRIPT> [ARGS]"

• export CUDA_VISIBLE_DEVICES=0; ensures deterministic single-GPU execution. Adjust for multi-GPU jobs.
• --extra gpu activates GPU optional dependencies (e.g. jax[cuda]).

TPU

sky exec $(cat .cluster_name_tpu) --workdir . "uv run --extra tpu python <PATH_TO_SCRIPT> [ARGS]"

• --extra tpu activates TPU optional dependencies (e.g. jax[tpu]).

Common flags

• --workdir . syncs the current local directory to the remote instance before running.
• For pytest, use python -m pytest <test_path> instead of calling pytest directly.

5. Usage Examples

Run a benchmark on GPU:

sky exec $(cat .cluster_name_gpu) --workdir . "export CUDA_VISIBLE_DEVICES=0; uv run --extra gpu python src/lynx/perf/benchmark_train.py"

Run tests on TPU:

sky exec $(cat .cluster_name_tpu) --workdir . "uv run --extra tpu python -m pytest src/lynx/test/"

6. Operational Notes

• Logs: SkyPilot streams stdout and stderr directly to the terminal.
• Interruption: Ctrl+C may not kill the remote process; check SkyPilot docs for cleanup if needed.