From 01d2c9f4abc142eea589a78222522485dbbba8a5 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Tue, 1 Sep 2026 20:25:29 +0000 Subject: [PATCH 1/3] [bug fix] Always return one split_dict shard per DP rank Signed-off-by: Stas Bekman --- .gitignore | 4 +++ PR_DESCRIPTION.md | 11 -------- arctic_platform/common/utils/batch.py | 2 +- tests/common/test_split_dict.py | 40 +++++++++++++++++++++++++++ 4 files changed, 45 insertions(+), 12 deletions(-) delete mode 100644 PR_DESCRIPTION.md create mode 100644 tests/common/test_split_dict.py diff --git a/.gitignore b/.gitignore index bfbc22b..79bf9bb 100644 --- a/.gitignore +++ b/.gitignore @@ -170,3 +170,7 @@ recipes/**/outputs/ # Client recipe connection configs — these hold a Snowflake PAT. The checked-in # config.json.template is not matched (it does not end in .json). recipes/*.json + +# Clone-local PR / probe write-ups — never commit +PR_DESCRIPTION.md +PROBE_*.md diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md deleted file mode 100644 index ea8c5be..0000000 --- a/PR_DESCRIPTION.md +++ /dev/null @@ -1,11 +0,0 @@ -## Summary - -- SFT returns a single merged `metrics` dict after `fwd_bwd` + `step`, matching RL `update_actor` (step metrics as the base, fwd_bwd keys win). -- `ArcticSFTClient.train_step` and `merge_sft_step_metrics` live on the unified client (`arctic_platform/client/sft.py`). The worker already emitted `loss` / `grad_norm` (and any extra loss-fn keys); this wires them through as one dict instead of making every caller cherry-pick. -- Demos and `docs/sft.md` use `train_step`. -- Merges `origin/main` client unification (`ArcticSFTClient` subclasses `ArcticClient`; old `arctic_platform/sft/client.py` is gone). - -## Testing - -- CPU on `stas-dev-2-0` (`CUDA_VISIBLE_DEVICES=`, conda `dev`): `tests/sft/test_sft_client_ops.py` + `tests/sft/test_sft_config.py` + `tests/client/test_client_ops.py` — 84 passed after the merge (job `20260824T163728Z-1073948-000-2668422822`). -- No GPU / live-server run for this change (client surface only). diff --git a/arctic_platform/common/utils/batch.py b/arctic_platform/common/utils/batch.py index a4110fd..933c390 100644 --- a/arctic_platform/common/utils/batch.py +++ b/arctic_platform/common/utils/batch.py @@ -79,7 +79,7 @@ def _split_value(val, num_chunks: int): f"{num_chunks}. The client must send at least one sample per " "DP worker." ) - return list(torch.chunk(val, num_chunks, dim=0)) + return list(torch.tensor_split(val, num_chunks, dim=0)) if isinstance(val, list): if len(val) < num_chunks: raise ValueError( diff --git a/tests/common/test_split_dict.py b/tests/common/test_split_dict.py new file mode 100644 index 0000000..0bd012f --- /dev/null +++ b/tests/common/test_split_dict.py @@ -0,0 +1,40 @@ +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``split_dict`` must return one shard per DP rank for any batch size ``B >= n``.""" + +from __future__ import annotations + +import torch + +from arctic_platform.common.utils.batch import split_dict +from arctic_platform.testing_utils import TestCasePlus +from arctic_platform.testing_utils import torch_assert_equal + + +class TestSplitDictRemainder(TestCasePlus): + def test_tensor_split_returns_one_shard_per_rank(self): + # torch.chunk(B, n) can return fewer than n tensors (B=6 n=4 → 3 chunks). + for batch_size, num_chunks in ((6, 4), (5, 4), (9, 4), (13, 8), (4, 4), (7, 4)): + ids = torch.arange(batch_size * 3).view(batch_size, 3) + shards = split_dict({"input_ids": ids}, num_chunks) + self.assertEqual(len(shards), num_chunks, msg=f"B={batch_size} n={num_chunks}") + rows = [int(s["input_ids"].shape[0]) for s in shards] + self.assertEqual(sum(rows), batch_size, msg=f"B={batch_size} n={num_chunks} rows={rows}") + torch_assert_equal(torch.cat([s["input_ids"] for s in shards], dim=0), ids) + + def test_batch_smaller_than_ranks_is_rejected(self): + ids = torch.arange(6).view(3, 2) + with self.assertRaises(ValueError): + split_dict({"input_ids": ids}, 4) From c0ef73b5cdfbc8994e55a424884374a93a8ef8b6 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Tue, 1 Sep 2026 20:27:35 +0000 Subject: [PATCH 2/3] cleanup Signed-off-by: Stas Bekman --- .gitignore | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.gitignore b/.gitignore index 79bf9bb..bfbc22b 100644 --- a/.gitignore +++ b/.gitignore @@ -170,7 +170,3 @@ recipes/**/outputs/ # Client recipe connection configs — these hold a Snowflake PAT. The checked-in # config.json.template is not matched (it does not end in .json). recipes/*.json - -# Clone-local PR / probe write-ups — never commit -PR_DESCRIPTION.md -PROBE_*.md From c895e5418c7c8be472f25424bc0835b19986224e Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Tue, 1 Sep 2026 21:44:43 +0000 Subject: [PATCH 3/3] [CI] adapt to AI's changed Signed-off-by: Stas Bekman --- .github/workflows/unit_tests.yaml | 8 +------- pyproject.toml | 14 +++++--------- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/.github/workflows/unit_tests.yaml b/.github/workflows/unit_tests.yaml index bb8ce9d..dff5ea4 100644 --- a/.github/workflows/unit_tests.yaml +++ b/.github/workflows/unit_tests.yaml @@ -53,13 +53,7 @@ jobs: echo "$VIRTUAL_ENV/bin" >> $GITHUB_PATH echo "VIRTUAL_ENV=$VIRTUAL_ENV" >> $GITHUB_ENV uv pip install torch --index-url https://download.pytorch.org/whl/cpu - # [rl] declares arctic-inference and vllm, so this resolves the whole stack. - # ArcticInference builds a native nanobind extension via CMake. actions/setup-python exports - # Python(3)_ROOT_DIR/pythonLocation pointing at the tool-cache interpreter (no nanobind); CMake's - # find_package(Python) honors those hints and fails to find nanobind. Unset them so CMake uses uv's - # build-env interpreter, which has nanobind from ArcticInference's build-system.requires. - env -u Python_ROOT_DIR -u Python2_ROOT_DIR -u Python3_ROOT_DIR -u pythonLocation \ - uv pip install ".[testing,rl]" + uv pip install ".[sft,testing]" - name: Environment print run: | uv pip list diff --git a/pyproject.toml b/pyproject.toml index 270077b..e247be7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,9 +47,6 @@ classifiers = [ requires-python = ">=3.12" - -# Config models only. Every usable install picks a backend extra: [cortex] to -# drive Cortex, [sft]/[rl]/[onprem] to run against an on-prem server. dependencies = [ "pydantic>=2.10", "typing-extensions", # Self, per the pydantic model_validator convention @@ -70,10 +67,7 @@ dev = [ "arctic_platform[testing]", ] -# Pulls [onprem] because the suite exercises both servers, and because `dev` is -# defined as formatting + testing. testing = [ - "arctic_platform[onprem]", "parameterized", "pytest-instafail", "pytest-flakefinder", @@ -126,15 +120,17 @@ sft = [ "nvidia-ml-py", "psutil", "ray", + "tensordict", "transformers", "uvicorn", ] -# [sft] plus sampling. arctic-inference[vllm] owns the validated vLLM pin. +# [sft] plus sampling. arctic-inference[vllm] pulls vLLM and applies that extra's +# pin; [rl] does not declare vLLM. 0.3.0 is the first release of that extra +# contract. rl = [ "arctic_platform[sft]", - "arctic-inference[server,vllm]>=0.2.0", - "tensordict", + "arctic-inference[server,vllm]>=0.3.0", ] onprem = [