From 7e46aba3bdb930c098379b579c647af05e32b3d2 Mon Sep 17 00:00:00 2001 From: Gregory Koganovsky Date: Wed, 27 May 2026 15:35:48 +0300 Subject: [PATCH 1/2] feat: support fractional CPU limits in operation resources - Updated the `cpu_limit` parameter to accept fractional values in configuration files and code, enhancing flexibility for resource allocation. - Revised related documentation to clarify the usage of fractional CPU limits. - Added unit tests to validate acceptance and rejection of various CPU limit inputs, ensuring robust error handling and compliance with new specifications. --- docs/operations/map.md | 4 ++-- docs/operations/vanilla.md | 4 ++-- tests/test_client_base.py | 22 ++++++++++++++++++- tests/test_common.py | 6 +++++ yt_framework/operations/common.py | 6 ++++- .../yt/support/operation_resources.py | 20 ++++++++++++----- 6 files changed, 51 insertions(+), 11 deletions(-) diff --git a/docs/operations/map.md b/docs/operations/map.md index c4123bd..5806657 100644 --- a/docs/operations/map.md +++ b/docs/operations/map.md @@ -238,7 +238,7 @@ resources: pool: default # YT pool name pool_tree: null # Optional: pool tree memory_limit_gb: 4 # Memory per job (GB) - cpu_limit: 2 # CPU cores per job + cpu_limit: 2 # CPU cores per job (fractional, e.g. 0.5) job_count: 2 # Number of parallel jobs gpu_limit: 0 # GPU count (default: 0) user_slots: null # Optional: user slots limit @@ -247,7 +247,7 @@ resources: **Resources (rule of thumb):** - Raise `memory_limit_gb` when a single row plus model weights no longer fits. -- `cpu_limit` helps per-task throughput; `job_count` spreads rows across tasks. +- `cpu_limit` helps per-task throughput; fractional values (e.g. `0.5`, `0.1`) are passed through to YTsaurus; `job_count` spreads rows across tasks. - `gpu_limit` > 0 only works with an image that actually exposes GPUs to Python. ### Advanced Configuration diff --git a/docs/operations/vanilla.md b/docs/operations/vanilla.md index f439a03..b2bd34c 100644 --- a/docs/operations/vanilla.md +++ b/docs/operations/vanilla.md @@ -221,7 +221,7 @@ resources: pool: default # YT pool name pool_tree: null # Optional: pool tree memory_limit_gb: 2 # Memory (GB) - cpu_limit: 1 # CPU cores + cpu_limit: 1 # CPU cores (fractional, e.g. 0.5) gpu_limit: 0 # GPU count (default: 0) user_slots: null # Optional: user slots limit ``` @@ -229,7 +229,7 @@ resources: **Resource guidelines:** - **Memory**: Allocate based on task needs -- **CPU**: More CPUs = faster execution +- **CPU**: More cores = faster execution; fractional limits (e.g. `0.5`) are supported - **GPU**: Set to 1+ for GPU workloads (requires custom Docker) ### Advanced Configuration diff --git a/tests/test_client_base.py b/tests/test_client_base.py index 4fadb77..91eb3db 100644 --- a/tests/test_client_base.py +++ b/tests/test_client_base.py @@ -23,10 +23,30 @@ def test_operation_resources_rejects_non_positive_memory_gb() -> None: def test_operation_resources_rejects_non_positive_cpu_limit() -> None: - with pytest.raises(ValueError, match="cpu_limit must be set to a positive integer"): + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): OperationResources(cpu_limit=0) +def test_operation_resources_accepts_fractional_cpu_limit() -> None: + resources = OperationResources(cpu_limit=0.5) + assert resources.cpu_limit == 0.5 + + +def test_operation_resources_rejects_zero_fractional_cpu_limit() -> None: + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=0.0) + + +def test_operation_resources_rejects_negative_fractional_cpu_limit() -> None: + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=-0.5) + + +def test_operation_resources_rejects_bool_cpu_limit() -> None: + with pytest.raises(TypeError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=True) # type: ignore[arg-type] + + def test_operation_resources_rejects_negative_gpu_limit() -> None: with pytest.raises( ValueError, match="gpu_limit must be set to a non-negative integer" diff --git a/tests/test_common.py b/tests/test_common.py index 810aeb8..810590f 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -44,6 +44,12 @@ def test_extract_operation_resources_reads_top_level_keys() -> None: assert res.cpu_limit == 8 +def test_extract_operation_resources_parses_fractional_cpu_limit() -> None: + cfg = OmegaConf.create({"cpu_limit": 0.5}) + res = extract_operation_resources(cfg, _LOG) + assert res.cpu_limit == 0.5 + + def test_extract_operation_resources_reads_nested_resources_block() -> None: cfg = OmegaConf.create({"resources": {"pool": "nest", "job_count": 3}}) res = extract_operation_resources(cfg, _LOG) diff --git a/yt_framework/operations/common.py b/yt_framework/operations/common.py index 5bb3ad6..6df840a 100644 --- a/yt_framework/operations/common.py +++ b/yt_framework/operations/common.py @@ -38,6 +38,10 @@ def _int_from_config_value(value: object) -> int: return int(cast("Any", value)) +def _float_from_config_value(value: object) -> float: + return float(cast("Any", value)) + + def _get_config_value_with_default( config: DictConfig, key: str, @@ -161,7 +165,7 @@ def extract_operation_resources( logger, ), ) - cpu_limit = _int_from_config_value( + cpu_limit = _float_from_config_value( _get_config_value_with_default(resources_config, "cpu_limit", 2, logger), ) gpu_limit = _int_from_config_value( diff --git a/yt_framework/yt/support/operation_resources.py b/yt_framework/yt/support/operation_resources.py index 1b324ed..95daaad 100644 --- a/yt_framework/yt/support/operation_resources.py +++ b/yt_framework/yt/support/operation_resources.py @@ -1,5 +1,6 @@ """Operation resource limits for YT jobs.""" +import math from dataclasses import dataclass @@ -9,6 +10,15 @@ def _require_positive_resource(name: str, value: int) -> None: raise ValueError(msg) +def _require_positive_cpu_limit(name: str, value: object) -> None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + msg = f"{name} must be a positive number, got {value!r}" + raise TypeError(msg) + if value <= 0 or not math.isfinite(float(value)): + msg = f"{name} must be a positive number, got {value}" + raise ValueError(msg) + + def _require_non_negative_gpu(gpu_limit: int) -> None: if gpu_limit < 0: msg = f"gpu_limit must be set to a non-negative integer, got {gpu_limit}" @@ -28,14 +38,14 @@ class OperationResources: pool_tree: Optional pool tree name (default: None). docker_image: Optional Docker image name for containerized execution (default: None). memory_gb: Memory allocation in GB (default: 4). In config files, use `memory_limit_gb`. - cpu_limit: CPU cores allocated (default: 2). + cpu_limit: CPU cores allocated per job; fractional values allowed (default: 2). gpu_limit: Number of GPUs allocated (default: 0). job_count: Number of parallel jobs (default: 1). user_slots: Optional user slots limit (default: None). Raises: - ValueError: If memory_gb, cpu_limit, or job_count are not positive integers, - or if gpu_limit is negative. + ValueError: If memory_gb or job_count are not positive integers, if cpu_limit + is not a positive number, or if gpu_limit is negative. """ @@ -43,7 +53,7 @@ class OperationResources: pool_tree: str | None = None docker_image: str | None = None memory_gb: int = 4 - cpu_limit: int = 2 + cpu_limit: float = 2 gpu_limit: int = 0 job_count: int = 1 user_slots: int | None = None @@ -51,6 +61,6 @@ class OperationResources: def __post_init__(self) -> None: """Validate resource fields after initialization.""" _require_positive_resource("memory_gb", self.memory_gb) - _require_positive_resource("cpu_limit", self.cpu_limit) + _require_positive_cpu_limit("cpu_limit", self.cpu_limit) _require_non_negative_gpu(self.gpu_limit) _require_positive_resource("job_count", self.job_count) From 2d51b8c7d046f68b3208e0be5ba4214fb3e01a57 Mon Sep 17 00:00:00 2001 From: Gregory Koganovsky Date: Wed, 27 May 2026 15:42:34 +0300 Subject: [PATCH 2/2] feat: enhance CPU limit validation in operation resources - Introduced a new `validate_cpu_limit` function to ensure `cpu_limit` values are positive finite numbers and not of type `bool`. - Added unit tests to validate rejection of non-finite and boolean CPU limit values, improving error handling and robustness. - Updated existing tests to reflect changes in CPU limit validation logic, ensuring comprehensive coverage. --- tests/test_client_base.py | 16 ++++++++++++++++ tests/test_common.py | 14 ++++++++++++++ yt_framework/operations/common.py | 2 ++ yt_framework/yt/clients/operation_resources.py | 9 ++++++--- yt_framework/yt/support/operation_resources.py | 14 +++++++++++++- 5 files changed, 51 insertions(+), 4 deletions(-) diff --git a/tests/test_client_base.py b/tests/test_client_base.py index 91eb3db..0c96f95 100644 --- a/tests/test_client_base.py +++ b/tests/test_client_base.py @@ -1,6 +1,7 @@ """Tests for yt_framework.yt.clients.client_base (OperationResources + BaseYTClient helpers).""" import logging +import math from pathlib import Path from typing import Any from unittest.mock import MagicMock @@ -47,6 +48,21 @@ def test_operation_resources_rejects_bool_cpu_limit() -> None: OperationResources(cpu_limit=True) # type: ignore[arg-type] +def test_operation_resources_rejects_non_finite_cpu_limit() -> None: + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=float("inf")) + + +def test_operation_resources_rejects_negative_infinite_cpu_limit() -> None: + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=float("-inf")) + + +def test_operation_resources_rejects_nan_cpu_limit() -> None: + with pytest.raises(ValueError, match="cpu_limit must be a positive number"): + OperationResources(cpu_limit=math.nan) + + def test_operation_resources_rejects_negative_gpu_limit() -> None: with pytest.raises( ValueError, match="gpu_limit must be set to a non-negative integer" diff --git a/tests/test_common.py b/tests/test_common.py index 810590f..3f3f589 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -4,6 +4,7 @@ from pathlib import Path from unittest.mock import MagicMock, patch +import pytest from omegaconf import OmegaConf from yt_framework.core.dependencies import PipelineStageDependencies @@ -50,6 +51,19 @@ def test_extract_operation_resources_parses_fractional_cpu_limit() -> None: assert res.cpu_limit == 0.5 +def test_extract_operation_resources_parses_fractional_cpu_limit_in_nested_resources() -> ( + None +): + cfg = OmegaConf.create({"resources": {"cpu_limit": 0.5}}) + res = extract_operation_resources(cfg, _LOG) + assert res.cpu_limit == 0.5 + + +def test_extract_operation_resources_rejects_bool_cpu_limit() -> None: + with pytest.raises(TypeError, match="cpu_limit must be a positive number"): + extract_operation_resources(OmegaConf.create({"cpu_limit": True}), _LOG) + + def test_extract_operation_resources_reads_nested_resources_block() -> None: cfg = OmegaConf.create({"resources": {"pool": "nest", "job_count": 3}}) res = extract_operation_resources(cfg, _LOG) diff --git a/yt_framework/operations/common.py b/yt_framework/operations/common.py index 6df840a..8944709 100644 --- a/yt_framework/operations/common.py +++ b/yt_framework/operations/common.py @@ -9,6 +9,7 @@ from yt_framework.utils.env import load_secrets from yt_framework.yt.clients.client_base import OperationResources +from yt_framework.yt.clients.operation_resources import validate_cpu_limit from ._internal.tokenizer_artifact import ( init_tokenizer_artifact_directory, @@ -39,6 +40,7 @@ def _int_from_config_value(value: object) -> int: def _float_from_config_value(value: object) -> float: + validate_cpu_limit("cpu_limit", value) return float(cast("Any", value)) diff --git a/yt_framework/yt/clients/operation_resources.py b/yt_framework/yt/clients/operation_resources.py index ff402a9..9fd09fd 100644 --- a/yt_framework/yt/clients/operation_resources.py +++ b/yt_framework/yt/clients/operation_resources.py @@ -1,5 +1,8 @@ -"""Re-export :class:`OperationResources` from ``yt_framework.yt.support``.""" +"""Re-export operation resource types from ``yt_framework.yt.support``.""" -from yt_framework.yt.support.operation_resources import OperationResources +from yt_framework.yt.support.operation_resources import ( + OperationResources, + validate_cpu_limit, +) -__all__ = ["OperationResources"] +__all__ = ["OperationResources", "validate_cpu_limit"] diff --git a/yt_framework/yt/support/operation_resources.py b/yt_framework/yt/support/operation_resources.py index 95daaad..af99f7e 100644 --- a/yt_framework/yt/support/operation_resources.py +++ b/yt_framework/yt/support/operation_resources.py @@ -19,6 +19,17 @@ def _require_positive_cpu_limit(name: str, value: object) -> None: raise ValueError(msg) +def validate_cpu_limit(name: str, value: object) -> None: + """Validate a ``cpu_limit`` value (config or constructor). + + Raises: + TypeError: If ``value`` is not an ``int`` or ``float`` (e.g. ``bool``). + ValueError: If ``value`` is not a positive finite number. + + """ + _require_positive_cpu_limit(name, value) + + def _require_non_negative_gpu(gpu_limit: int) -> None: if gpu_limit < 0: msg = f"gpu_limit must be set to a non-negative integer, got {gpu_limit}" @@ -45,7 +56,8 @@ class OperationResources: Raises: ValueError: If memory_gb or job_count are not positive integers, if cpu_limit - is not a positive number, or if gpu_limit is negative. + is not a positive finite number, or if gpu_limit is negative. + TypeError: If cpu_limit is not an ``int`` or ``float`` (e.g. ``bool``). """