Skypilot Support (#151)

* Create skypilot_train.yaml

* entered template for skypilot job launch

* addressed PR comments

* removed accelerate to avoid double launching

* user doesn't supply accelerate template, but uses it internally still for arbor services

* fixed accidental deletion

* fix for deepspeed testing on scripts

* ruff checks

---------

Co-authored-by: Noah Ziems <nziems2@nd.edu>

Author: zhassan223

README.md

@@ -144,14 +144,6 @@ print(optimized_translate(english="hello"))
 ```
 ---
 
-### Accelerate Configuration
-
-For advanced distributed training setups, you can specify a custom [Hugging Face Accelerate](https://huggingface.co/docs/accelerate/) config in your `~/.arbor/config.yaml`:
-
-```yaml
-accelerate_config: "/path/to/your/accelerate_config.yaml"
-```
-
 ---
 
 ### Troubleshooting

arbor/core/config.py

@@ -11,9 +11,6 @@ class Config(BaseModel):
     # Basic settings
     storage_path: str = str(Path.home() / ".arbor" / "storage")
 
-    # Training settings
-    accelerate_config: Optional[str] = None
-
     # Server settings
     inactivity_timeout: int = 30
 

arbor/server/services/jobs/grpo_job.py

@@ -210,7 +210,6 @@ def _launch_inference_job(
             num_processes=num_processes,
             main_process_port=main_process_port,
             script_args=script_args,
-            accelerate_config=self.config.accelerate_config,
             env=my_env,
             log_callback=self.create_log_callback("GRPO"),
         )

arbor/utils/process_runner.py

@@ -2,6 +2,7 @@
 Clean abstraction for running long-running processes without the Popen ugliness.
 """
 
+import os
 import signal
 import subprocess
 import sys
@@ -210,43 +211,55 @@ def start_training(
         num_processes: int,
         main_process_port: int,
         script_args: list[str],
-        accelerate_config: Optional[str] = None,
         env: Optional[dict[str, str]] = None,
         log_callback: Optional[Callable[[str], None]] = None,
     ) -> subprocess.Popen:
         """
-        Start an accelerate-based training process.
+        Start an accelerate-based training process using the bundled config.
 
         Args:
-            script_path: Path to the training script
+            module: The python module to run
             num_processes: Number of processes for accelerate
             main_process_port: Port for main process
-            script_args: Arguments to pass to the script (everything after script path)
-            accelerate_config: Optional accelerate config file path
+            script_args: Arguments to pass to the script
             env: Environment variables
             log_callback: Function to call with each log line
         """
+        # Path to the bundled deepspeed config within the project
+        project_root = Path(__file__).resolve().parent.parent.parent
+        accelerate_config = str(project_root / "configs" / "deepspeed_config.yaml")
+
         command = [
             sys.executable,
             "-m",
             "accelerate.commands.launch",
+            "--config_file",
+            accelerate_config,
             "--num_processes",
             str(num_processes),
             "--main_process_port",
             str(main_process_port),
         ]
 
-        if accelerate_config:
-            command.extend(["--config_file", accelerate_config])
-
         command.extend(["--module", module])
 
         if script_args:
             # Separator required so accelerate forwards these to the module
             command.append("--")
             command.extend(script_args)
 
-        return self.start(command, env=env, log_callback=log_callback)
+        # Ensure a stable, unique rendezvous port is used for torch.distributed
+        runtime_env = dict(os.environ)
+        if env:
+            runtime_env.update(env)
+        # Pin MASTER_ADDR/MASTER_PORT and Accelerate's port to the provided main_process_port
+        runtime_env.setdefault("MASTER_ADDR", "127.0.0.1")
+        runtime_env["MASTER_PORT"] = str(main_process_port)
+        runtime_env["ACCELERATE_TORCH_DISTRIBUTED_PORT"] = str(main_process_port)
+        # Optionally keep torch distributed debug noise off unless explicitly enabled by caller
+        runtime_env.setdefault("TORCH_DISTRIBUTED_DEBUG", "OFF")
+
+        return self.start(command, env=runtime_env, log_callback=log_callback)
 
     def start_training_from_full_command(
         self,

configs/deepspeed_config.yaml

@@ -0,0 +1,26 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+  bf16:
+    enabled: auto
+  fp16:
+    enabled: auto
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: no
+num_machines: 1
+# num_processes: 4 //infer from command line
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

deploy/README.md

@@ -0,0 +1,189 @@
+# Arbor Deployment Guide
+
+This directory contains SkyPilot configurations for running Arbor GRPO training jobs on cloud infrastructure.
+
+## Quick Start
+
+### Prerequisites
+
+- [SkyPilot](https://skypilot.readthedocs.io/) installed: `pip install skypilot[aws]` (or your cloud provider)
+- AWS/GCP/other cloud credentials configured
+- An S3 bucket (or equivalent) for storing checkpoints
+
+### Running Any Training Script
+
+The `skypilot_example.yaml` configuration is designed to work with any Python training script. Scripts run from the project root directory, so use paths relative to root (e.g., `examples/your_script.py`).
+
+#### Step 1: Configure Your Training Script
+
+```bash
+sky jobs launch -n training-job deploy/skypilot_example.yaml \
+  --env TRAINING_SCRIPT=examples/your_script.py \
+  --infra aws
+```
+
+**Example:**
+```bash
+sky jobs launch -n banking-training deploy/skypilot_example.yaml \
+  --env TRAINING_SCRIPT=examples/banking_grpo.py \
+  --infra aws
+```
+
+#### Step 2: Configure Your Checkpoint S3 Bucket
+
+Edit `skypilot_example.yaml` and update the checkpoint bucket path (around line 25):
+
+```yaml
+file_mounts:
+  /mnt/checkpoints:
+    source: s3://your-bucket-name/  # Change this to your S3 bucket
+    mode: MOUNT_CACHED
+```
+
+Replace `s3://your-bucket-name/` with your actual S3 bucket path. The bucket should already exist and be accessible with your cloud credentials.
+
+#### Step 3: Customize Setup (if needed)
+
+If your script requires additional dependencies or data downloads, edit the `setup` section in `skypilot_example.yaml`:
+
+```yaml
+setup: |
+  # ... common setup ...
+
+  # Add your script-specific dependencies:
+  uv pip install package-name
+
+  # Add your script-specific data downloads:
+  curl -Ls https://example.com/data.tar.gz | tar -xz -C examples
+```
+
+**Example for `multihop_dev.py`:**
+```yaml
+setup: |
+  # ... common setup ...
+  uv pip install ujson bm25s PyStemmer "jax[cpu]"
+  curl -Ls https://huggingface.co/dspy/cache/resolve/main/wiki.abstracts.2017.tar.gz | tar -xz -C examples
+```
+
+#### Step 4: Launch the Training Job
+
+```bash
+sky jobs launch -n training-job deploy/skypilot_example.yaml --infra <aws|gcp|nebius|lambda>
+```
+
+Replace `<aws|gcp|nebius|lambda>` with your cloud provider.
+
+#### Step 5: Monitor Progress
+
+```bash
+# View logs
+sky logs training-job
+
+# Check job status
+sky status
+```
+
+#### Step 6: Stop the Job (if needed)
+
+```bash
+sky jobs cancel training-job
+```
+
+## Configuration Details
+
+### Resources
+
+- **GPUs**: 4x L40S GPUs by default (adjust based on your script's needs)
+- **Spot Instances**: Disabled by default (`use_spot: false`). Set to `true` for cost savings
+- **Disk**: 256 GB
+
+### DeepSpeed Configuration
+
+The training uses DeepSpeed ZeRO-3 (configured in `../configs/deepspeed_config.yaml`):
+
+### What Gets Installed
+
+The setup script automatically installs:
+- Project dependencies via `uv sync`
+- Common packages: `datasets`, `deepspeed`, `accelerate`
+
+You can add script-specific dependencies in the setup section.
+
+## Examples
+
+### Running multihop_dev.py
+
+1. Edit `skypilot_example.yaml` setup section to uncomment multihop-specific setup:
+```yaml
+uv pip install ujson bm25s PyStemmer "jax[cpu]"
+curl -Ls https://huggingface.co/dspy/cache/resolve/main/wiki.abstracts.2017.tar.gz | tar -xz -C examples
+```
+
+2. The script path is already set to `examples/multihop_dev.py` by default
+
+3. Launch:
+```bash
+sky jobs launch -n multihop-dev deploy/skypilot_example.yaml --infra aws
+```
+
+### Running banking_grpo.py
+
+1. Edit the script path in the run section:
+```yaml
+TRAINING_SCRIPT=${TRAINING_SCRIPT:-examples/banking_grpo.py}
+```
+
+2. Launch (without accelerate, since it doesn't need distributed training):
+```bash
+sky jobs launch -n banking-training deploy/skypilot_example.yaml --infra aws
+```
+
+Or if you want to use accelerate anyway:
+```bash
+sky jobs launch -n banking-training deploy/skypilot_example.yaml \
+  --env TRAINING_SCRIPT=examples/banking_grpo.py \
+  --env USE_ACCELERATE=true \
+  --infra aws
+```
+
+No additional setup needed - it only uses standard dependencies.
+
+## Customization
+
+### Changing GPU Count
+
+To use a different number of GPUs, modify `skypilot_example.yaml`:
+
+```yaml
+resources:
+  accelerators: L40S:4  # Change number as needed
+```
+
+**Note**: If you change the GPU count and your script uses accelerate, make sure the number matches your script's expectations.
+
+### Using Spot Instances
+
+For cost savings, enable spot instances:
+
+```yaml
+resources:
+  use_spot: true
+```
+
+### Changing Cloud Providers
+
+SkyPilot supports multiple providers. Specify during launch:
+
+```bash
+sky jobs launch -n training-job deploy/skypilot_example.yaml --infra aws
+sky jobs launch -n training-job deploy/skypilot_example.yaml --infra gcp
+sky jobs launch -n training-job deploy/skypilot_example.yaml --infra azure
+```
+
+
+
+## Additional Resources
+
+- [SkyPilot Documentation](https://skypilot.readthedocs.io/)
+- [DeepSpeed Documentation](https://www.deepspeed.ai/)
+- [Arbor Documentation](../README.md)

deploy/skypilot_example.yaml

@@ -0,0 +1,49 @@
+# Run the full multihop training
+#
+# Launch:
+#   sky jobs launch -n multihop-dev deploy/skypilot_train.yaml --infra <aws|gcp|nebius|lambda|etc>
+#
+# Monitor progress:
+#   sky logs multihop-dev
+#
+# This will train the model using 4x L40S GPUs and save results to S3.
+
+name: arbor-multihop-train
+
+resources:
+  accelerators: L40S:4
+  use_spot: true
+  disk_size: 256
+
+file_mounts:
+  /mnt/checkpoints:
+    source: s3://arbor-multihop-checkpoints/
+    mode: MOUNT_CACHED
+
+#envs:
+  # Fill in your wandb key: copy from https://wandb.ai/authorize
+  # Alternatively, you can use `--env WANDB_API_KEY=$WANDB_API_KEY`
+  # to pass the key in the command line, during `sky jobs launch`.
+#  WANDB_API_KEY:
+
+workdir: .
+
+setup: |
+  echo "Begin setup..."
+  curl -LsSf https://astral.sh/uv/install.sh | sh
+  export PATH="$HOME/.local/bin:$PATH"
+  uv sync
+  echo "Installing dependencies unique to multihop-dev..."
+  uv pip install ujson
+  uv pip install bm25s
+  uv pip install PyStemmer
+  uv pip install "jax[cpu]"
+  uv pip install datasets==3.6.0
+  echo "Downloading wiki.abstracts.2017.tar.gz from HuggingFace..."
+  curl -Ls https://huggingface.co/dspy/cache/resolve/main/wiki.abstracts.2017.tar.gz | tar -xz -C examples
+
+run: |
+  echo "Begin run..."
+  cd examples/
+  uv run python multihop_dev.py
+  echo "End run..."