diff --git a/openweights/cli/common.py b/openweights/cli/common.py index 783e596..87e6bb2 100644 --- a/openweights/cli/common.py +++ b/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" diff --git a/openweights/cli/ssh.py b/openweights/cli/ssh.py index 1fe34bb..4f66427 100644 --- a/openweights/cli/ssh.py +++ b/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 diff --git a/openweights/cluster/start_runpod.py b/openweights/cluster/start_runpod.py index 3005017..54e4564 100644 --- a/openweights/cluster/start_runpod.py +++ b/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") diff --git a/tests/test_min_vcpu_passthrough.py b/tests/test_min_vcpu_passthrough.py new file mode 100644 index 0000000..c4ae84f --- /dev/null +++ b/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