feat(runpod): allow requesting a minimum host vCPU / RAM for pods

openweights/cli/common.py

@@ -34,7 +34,13 @@ class StartResult:
 
 class Provider:
     def start(
-        self, image: str, gpu: str, count: int, env: Dict[str, str]
+        self,
+        image: str,
+        gpu: str,
+        count: int,
+        env: Dict[str, str],
+        min_vcpu_count: Optional[int] = None,
+        min_memory_in_gb: Optional[int] = None,
     ) -> StartResult:
         raise NotImplementedError
 
@@ -50,7 +56,13 @@ def __init__(self, key_path: str):
         self.key_path = os.path.expanduser(key_path)
 
     def start(
-        self, image: str, gpu: str, count: int, env: Dict[str, str]
+        self,
+        image: str,
+        gpu: str,
+        count: int,
+        env: Dict[str, str],
+        min_vcpu_count: Optional[int] = None,
+        min_memory_in_gb: Optional[int] = None,
     ) -> StartResult:
         if "RUNPOD_API_KEY" in env:
             os.environ["RUNPOD_API_KEY"] = env["RUNPOD_API_KEY"]
@@ -64,6 +76,8 @@ def start(
             env=env,
             runpod_client=runpod,
             dev_mode=True,  # keep your current choice
+            min_vcpu_count=min_vcpu_count,
+            min_memory_in_gb=min_memory_in_gb,
         )
         assert pod is not None, "Runpod start_worker returned None"
 

openweights/cli/ssh.py

@@ -91,6 +91,22 @@ def add_ssh_parser(parser):
     )
     parser.add_argument("--gpu", default="L40", help="GPU type for provider.")
     parser.add_argument("--count", type=int, default=1, help="GPU count.")
+    parser.add_argument(
+        "--min-vcpu",
+        type=int,
+        default=None,
+        help=(
+            "Only place the pod on a host offering at least this many vCPUs "
+            "(total for the pod, not per GPU). Useful for CPU-bound workloads. "
+            "Higher values shrink the pool of eligible hosts."
+        ),
+    )
+    parser.add_argument(
+        "--min-memory-gb",
+        type=int,
+        default=None,
+        help="Only place the pod on a host offering at least this much RAM in GB.",
+    )
     parser.add_argument(
         "--remote-cwd",
         default="/workspace",
@@ -153,7 +169,12 @@ def __init__(self, ssh_spec):
 
         print("[ow] Starting/allocating machine...")
         start_res = provider.start(
-            image=args.image, gpu=args.gpu, count=args.count, env=provider_env
+            image=args.image,
+            gpu=args.gpu,
+            count=args.count,
+            env=provider_env,
+            min_vcpu_count=args.min_vcpu,
+            min_memory_in_gb=args.min_memory_gb,
         )
         ssh = start_res.ssh
 

openweights/cluster/start_runpod.py

@@ -224,6 +224,20 @@ def _env_cooldown_ladder_seconds() -> tuple[int, ...]:
 RUNPOD_MIN_UPLOAD = os.getenv("OW_RUNPOD_MIN_UPLOAD")
 RUNPOD_DATA_CENTER_ID = os.getenv("OW_RUNPOD_DATA_CENTER_ID")
 RUNPOD_COUNTRY_CODE = os.getenv("OW_RUNPOD_COUNTRY_CODE")
+# Minimum host CPU / RAM required when placing a pod.
+#
+# RunPod treats these as machine *filters*, not allocations: they exclude hosts
+# that would give the pod less than this, but never grant more than the host's
+# per-GPU share. Both are absolute per-pod values, not per-GPU, matching the
+# semantics of RunPod's `minVcpuCount` / `minMemoryInGb` deploy fields.
+#
+# Useful for workloads that are CPU-bound rather than VRAM-bound (e.g. RL
+# environments that execute generated code for every rollout), where the
+# default placement can land on a host with far fewer cores than the job needs.
+# Setting these too high shrinks the pool of eligible hosts and makes
+# provisioning failures more likely.
+RUNPOD_MIN_VCPU_COUNT = os.getenv("OW_RUNPOD_MIN_VCPU_COUNT")
+RUNPOD_MIN_MEMORY_GB = os.getenv("OW_RUNPOD_MIN_MEMORY_GB")
 
 
 # Check that GPU name mapping is unique in both directions
@@ -640,6 +654,8 @@ def _start_worker(
     pending_workers=None,
     env=None,
     runpod_client=None,
+    min_vcpu_count=None,
+    min_memory_in_gb=None,
 ):
     client = runpod_client or runpod
     gpu = GPUs[gpu]
@@ -662,6 +678,20 @@ def _start_worker(
     )
     if worker_id is None:
         worker_id = uuid.uuid4().hex[:8]
+
+    # Explicit arguments win over the OW_RUNPOD_* environment defaults.
+    # Leaving both unset omits the fields entirely, preserving current behaviour.
+    if min_vcpu_count is None and RUNPOD_MIN_VCPU_COUNT:
+        min_vcpu_count = int(RUNPOD_MIN_VCPU_COUNT)
+    if min_memory_in_gb is None and RUNPOD_MIN_MEMORY_GB:
+        min_memory_in_gb = int(RUNPOD_MIN_MEMORY_GB)
+    if min_vcpu_count or min_memory_in_gb:
+        logger.info(
+            "Requesting host with min_vcpu_count=%s min_memory_in_gb=%s",
+            min_vcpu_count,
+            min_memory_in_gb,
+        )
+
     pod = client.create_pod(
         name,
         image,
@@ -677,6 +707,8 @@ def _start_worker(
         country_code=RUNPOD_COUNTRY_CODE,
         min_download=int(RUNPOD_MIN_DOWNLOAD) if RUNPOD_MIN_DOWNLOAD else None,
         min_upload=int(RUNPOD_MIN_UPLOAD) if RUNPOD_MIN_UPLOAD else None,
+        min_vcpu_count=min_vcpu_count,
+        min_memory_in_gb=min_memory_in_gb,
         ports="8000/http,10101/http,22/tcp",
         start_ssh=True,
         env=env,
@@ -702,6 +734,8 @@ def start_worker(
     ttl_hours=24,
     env=None,
     runpod_client=None,
+    min_vcpu_count=None,
+    min_memory_in_gb=None,
 ):
     pending_workers = []
     if dev_mode:
@@ -732,6 +766,8 @@ def start_worker(
             pending_workers,
             env,
             runpod_client,
+            min_vcpu_count,
+            min_memory_in_gb,
         )
         if pod is None:
             raise RuntimeError("RunPod create_pod returned no pod")

tests/test_min_vcpu_passthrough.py

@@ -0,0 +1,73 @@
+import pytest
+
+from openweights.cluster import start_runpod
+
+
+class FakeRunpodClient:
+    """Records the kwargs create_pod was called with."""
+
+    def __init__(self):
+        self.create_calls = []
+
+    def create_pod(self, *args, **kwargs):
+        self.create_calls.append(kwargs)
+        return {"id": "fake-pod-id"}
+
+    def terminate_pod(self, pod_id):
+        pass
+
+
+@pytest.fixture
+def client(monkeypatch):
+    monkeypatch.setenv("USER", "tester")
+    monkeypatch.setattr(start_runpod, "RUNPOD_MIN_VCPU_COUNT", None)
+    monkeypatch.setattr(start_runpod, "RUNPOD_MIN_MEMORY_GB", None)
+    return FakeRunpodClient()
+
+
+def start(client, **overrides):
+    params = dict(
+        gpu="H200",
+        image="nielsrolf/ow-unsloth:v0.11",
+        count=4,
+        dev_mode=False,
+        env={},
+        runpod_client=client,
+    )
+    params.update(overrides)
+    start_runpod.start_worker(**params)
+    assert client.create_calls, "create_pod was never called"
+    return client.create_calls[0]
+
+
+def test_min_vcpu_and_memory_are_forwarded_to_create_pod(client):
+    kwargs = start(client, min_vcpu_count=96, min_memory_in_gb=500)
+
+    assert kwargs["min_vcpu_count"] == 96
+    assert kwargs["min_memory_in_gb"] == 500
+    assert kwargs["gpu_count"] == 4
+
+
+def test_fields_are_omitted_when_unset(client):
+    kwargs = start(client)
+
+    assert kwargs["min_vcpu_count"] is None
+    assert kwargs["min_memory_in_gb"] is None
+
+
+def test_environment_variables_provide_defaults(client, monkeypatch):
+    monkeypatch.setattr(start_runpod, "RUNPOD_MIN_VCPU_COUNT", "64")
+    monkeypatch.setattr(start_runpod, "RUNPOD_MIN_MEMORY_GB", "256")
+
+    kwargs = start(client)
+
+    assert kwargs["min_vcpu_count"] == 64
+    assert kwargs["min_memory_in_gb"] == 256
+
+
+def test_explicit_arguments_override_environment_defaults(client, monkeypatch):
+    monkeypatch.setattr(start_runpod, "RUNPOD_MIN_VCPU_COUNT", "64")
+
+    kwargs = start(client, min_vcpu_count=8)
+
+    assert kwargs["min_vcpu_count"] == 8