Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions docs/running-pipelines/resources-gpus-and-services.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,19 @@ Macrodata Cloud supports these GPU types:
- `a10`
- `l4`
- `t4`
- `any`

Pass an ordered list to fall back when the preferred GPU is unavailable:

```python
pipeline.launch_cloud(
name="gpu-inference",
gpu=mdr.GPU(type=["a10", "l4", "t4"], count=1),
)
```

Refiner preserves the order. Choices must be unique, and `any` may only be the
last choice, as in `["a10", "any"]`.

## Runtime services

Expand Down
3 changes: 3 additions & 0 deletions src/refiner/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
"FileAssetConfig": "refiner.pipeline.sinks.assets",
"GPU": "refiner.pipeline",
"GPUType": "refiner.pipeline",
"GPUTypeRequest": "refiner.pipeline",
"SUPPORTED_CUDA_VERSIONS": "refiner.pipeline",
"SUPPORTED_GPU_TYPES": "refiner.pipeline",
"from_items": "refiner.pipeline",
Expand Down Expand Up @@ -71,6 +72,7 @@
"FileAssetConfig",
"GPU",
"GPUType",
"GPUTypeRequest",
"SUPPORTED_CUDA_VERSIONS",
"SUPPORTED_GPU_TYPES",
"Secrets",
Expand Down Expand Up @@ -155,6 +157,7 @@ def __dir__() -> list[str]:
CUDAVersion,
GPU,
GPUType,
GPUTypeRequest,
SUPPORTED_CUDA_VERSIONS,
SUPPORTED_GPU_TYPES,
from_items,
Expand Down
9 changes: 8 additions & 1 deletion src/refiner/cli/jobs/get.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,14 @@ def _stage_gpu_text(runtime_config: Any) -> str:
if not isinstance(runtime_config, dict):
return "-"
count = _safe_text(runtime_config.get("gpuCount"))
gpu_type = _safe_text(runtime_config.get("gpuType"))
gpu_types = runtime_config.get("gpuTypes")
gpu_type = (
" → ".join(gpu_types)
if isinstance(gpu_types, list)
and gpu_types
and all(isinstance(value, str) and value for value in gpu_types)
else _safe_text(runtime_config.get("gpuType"))
)
if count in {"", "-"}:
return "-"
if gpu_type in {"", "-", "None"}:
Expand Down
3 changes: 3 additions & 0 deletions src/refiner/pipeline/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
"CUDAVersion": "refiner.pipeline.resources",
"GPU": "refiner.pipeline.resources",
"GPUType": "refiner.pipeline.resources",
"GPUTypeRequest": "refiner.pipeline.resources",
"SUPPORTED_CUDA_VERSIONS": "refiner.pipeline.resources",
"SUPPORTED_GPU_TYPES": "refiner.pipeline.resources",
}
Expand All @@ -45,6 +46,7 @@
"CUDAVersion",
"GPU",
"GPUType",
"GPUTypeRequest",
"RefinerPipeline",
"Row",
"Shard",
Expand Down Expand Up @@ -118,6 +120,7 @@ def __dir__() -> list[str]:
CUDAVersion,
GPU,
GPUType,
GPUTypeRequest,
SUPPORTED_CUDA_VERSIONS,
SUPPORTED_GPU_TYPES,
)
23 changes: 19 additions & 4 deletions src/refiner/pipeline/resources.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
from dataclasses import dataclass
from typing import Any, Literal, get_args

GPUType = Literal["h100", "l40s", "a100", "a10", "l4", "t4"]
GPUType = Literal["h100", "l40s", "a100", "a10", "l4", "t4", "any"]
GPUTypeRequest = GPUType | tuple[GPUType, ...] | list[GPUType]
CUDAVersion = Literal["12.4", "12.6", "12.8"]

SUPPORTED_GPU_TYPES: tuple[str, ...] = get_args(GPUType)
Expand All @@ -13,15 +14,28 @@
@dataclass(frozen=True, slots=True)
class GPU:
count: int
type: GPUType
type: GPUTypeRequest
cuda_version: CUDAVersion | None = None

def __post_init__(self) -> None:
if self.count <= 0:
raise ValueError("gpu.count must be > 0")
if self.type not in SUPPORTED_GPU_TYPES:
if isinstance(self.type, str):
choices = (self.type,)
elif isinstance(self.type, (list, tuple)):
choices = tuple(self.type)
object.__setattr__(self, "type", choices)
else:
choices = ()
if not 1 <= len(choices) <= 8:
raise ValueError("gpu.type must contain between 1 and 8 choices")
if any(choice not in SUPPORTED_GPU_TYPES for choice in choices):
supported = ", ".join(SUPPORTED_GPU_TYPES)
raise ValueError(f"gpu.type must be one of: {supported}")
if len(set(choices)) != len(choices):
raise ValueError("gpu.type choices must be unique")
if "any" in choices[:-1]:
raise ValueError("gpu.type 'any' must be the final choice")
if (
self.cuda_version is not None
and self.cuda_version not in SUPPORTED_CUDA_VERSIONS
Expand All @@ -32,7 +46,7 @@ def __post_init__(self) -> None:
def to_dict(self) -> dict[str, Any]:
payload: dict[str, Any] = {
"count": self.count,
"type": self.type,
"type": list(self.type) if isinstance(self.type, tuple) else self.type,
}
if self.cuda_version is not None:
payload["cuda_version"] = self.cuda_version
Expand All @@ -43,6 +57,7 @@ def to_dict(self) -> dict[str, Any]:
"CUDAVersion",
"GPU",
"GPUType",
"GPUTypeRequest",
"SUPPORTED_CUDA_VERSIONS",
"SUPPORTED_GPU_TYPES",
]
3 changes: 2 additions & 1 deletion tests/cli/test_job_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ def cli_get_job(self, *, job_id: str) -> dict[str, object]:
"memoryMb": 16384,
"gpuCount": 1,
"gpuType": "a10g",
"gpuTypes": ["a10", "l4", "t4"],
},
"steps": [
{
Expand Down Expand Up @@ -306,7 +307,7 @@ def test_jobs_get_plain_output(monkeypatch, capsys) -> None:
assert "GPU" in out.out
assert " 8 " in out.out
assert "16384" in out.out
assert "1 a10g" in out.out
assert "1 a10 → l4 → t4" in out.out
assert "Progress:" not in out.out
assert "columns=18" in out.out
assert "__meta" not in out.out
Expand Down
31 changes: 31 additions & 0 deletions tests/launchers/test_cloud_launcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -634,11 +634,42 @@ def test_pipeline_launch_cloud_accepts_structured_gpu(
assert runtime.gpu == GPU(count=1, type=gpu_type, cuda_version="12.4")


def test_pipeline_launch_cloud_preserves_ordered_gpu_fallbacks(monkeypatch) -> None:
captured = _stub_cloud_submit(monkeypatch)
monkeypatch.setattr(
"refiner.launchers.cloud.refiner_ref_exists_on_remote",
lambda ref: True,
)

read_jsonl("input.jsonl").launch_cloud(
name="demo cloud",
gpu=GPU(count=1, type=["a10", "l4", "t4"], cuda_version="12.4"),
)

request = cast(CloudRunCreateRequest, captured["submit_request"])
assert request.plan["stages"][0]["gpu"] == {
"count": 1,
"type": ["a10", "l4", "t4"],
"cuda_version": "12.4",
}
assert request.stage_payloads[0].runtime.gpu == GPU(
count=1,
type=("a10", "l4", "t4"),
cuda_version="12.4",
)


def test_gpu_rejects_unsupported_values() -> None:
with pytest.raises(ValueError, match="gpu.count must be > 0"):
GPU(count=0, type="h100")
with pytest.raises(ValueError, match="gpu.type must be one of"):
GPU(count=1, type="v100") # type: ignore[arg-type]
with pytest.raises(ValueError, match="between 1 and 8"):
GPU(count=1, type=[])
with pytest.raises(ValueError, match="unique"):
GPU(count=1, type=["l4", "l4"])
with pytest.raises(ValueError, match="'any' must be the final"):
GPU(count=1, type=["any", "l4"])
with pytest.raises(ValueError, match="gpu.cuda_version must be one of"):
GPU(count=1, type="h100", cuda_version="13.0") # type: ignore[arg-type]

Expand Down
12 changes: 12 additions & 0 deletions tests/platform/test_cloud_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,18 @@ def test_cloud_runtime_config_preserves_positional_resource_order() -> None:
assert runtime.region == ("us", "eu", "ca")


def test_cloud_runtime_config_serializes_ordered_gpu_fallbacks() -> None:
runtime = CloudRuntimeConfig(
num_workers=1,
gpu=GPU(count=1, type=["a10", "l4", "any"]),
)

assert runtime.to_dict()["gpu"] == {
"count": 1,
"type": ["a10", "l4", "any"],
}


def test_cloud_client_cloud_submit_job_posts_to_cloud_runs(monkeypatch) -> None:
captured: dict[str, object] = {}

Expand Down
Loading