freesolo-flash-dev 1.4.16__tar.gz → 1.4.17__tar.gz

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.4.16
+Version: 1.4.17
 Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/flash/custom/backend.py

@@ -50,6 +50,7 @@ class TorchBackend:
         return validate_operation(self, operation, arguments)
 
     def execute(self, operation, arguments):
+        _ensure_cuda_headroom()
         return getattr(self, operation)(**arguments)
 
     def _tensor(self, tokens):
@@ -271,6 +272,35 @@ class TorchBackend:
         return {}
 
 
+# cuBLAS handle/workspace creation and Triton kernels allocate outside PyTorch's caching
+# allocator, so they need memory the driver can still hand out.
+CUDA_HEADROOM_BYTES = 6 * 1024**3
+
+
+def _ensure_cuda_headroom(min_free_bytes=CUDA_HEADROOM_BYTES):
+    """Return PyTorch's unused cached blocks to the driver when free GPU memory runs low.
+
+    A session's operations see many sequence lengths, so cached blocks rarely fit the next
+    request and the caching allocator keeps reserving: a resumed Qwen3.5-9B DPO session held
+    ~78 GiB reserved with only ~19 GiB allocated. The first backward then needed memory outside
+    that cache (cuBLAS handle creation, Triton kernels) and failed with CUBLAS_STATUS_ALLOC_FAILED
+    / "Triton Error [CUDA]: out of memory" at <1 GiB free. Releasing the cache before an operation
+    restores that headroom; it changes no tensors. Best effort, a no-op without CUDA.
+    """
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+        free, _ = torch.cuda.mem_get_info()
+        if free >= min_free_bytes:
+            return False
+        torch.cuda.empty_cache()
+        return True
+    except Exception:
+        return False
+
+
 def _enable_checkpointing(model):
     """Checkpoint catalog models without active dropout or stateful normalization.
 

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/flash/custom/controller.py

@@ -10,7 +10,7 @@ import uuid
 from pathlib import Path
 
 from flash.client import client_from_config
-from flash.client.http import ApiError, ClientError
+from flash.client.http import ApiError, ClientError, RequestTimeoutError, ServiceUnreachableError
 from flash.client.specs import spec_payload
 from flash.custom.async_client import AsyncTrainingClient
 from flash.custom.client import TrainingClient
@@ -64,15 +64,16 @@ _ATTACH_RETRY_MAX_SECONDS = 30.0
 def _transient_attach_error(exc):
     """Whether a failed attach lookup can be retried rather than ending the run.
 
-    425 is the documented "GPU still provisioning" answer. A 5xx or a transport failure (no
-    status) means the control plane could not answer this time -- for example a slow Modal
+    425 is the documented "GPU still provisioning" answer. A 5xx or a typed transport failure
+    means the control plane could not answer this time -- for example a slow Modal
     tunnel lookup or an overloaded plane -- not that the run is gone. Treating those as fatal made
-    run_custom cancel healthy runs. Every other 4xx (ended run, auth, ownership) stays fatal.
+    run_custom cancel healthy runs. Every other 4xx (ended run, auth, ownership) and untyped
+    ClientError (such as a malformed successful response) stays fatal.
     """
     if isinstance(exc, ApiError):
         status = exc.status or 0
         return status == 425 or status >= 500
-    return isinstance(exc, ClientError)
+    return isinstance(exc, (RequestTimeoutError, ServiceUnreachableError))
 
 
 def connect(api, run_id, *, timeout=1200, sleep=time.sleep, clock=time.monotonic):

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/flash/providers/_lifecycle/net/worker.py

@@ -135,10 +135,12 @@ def build_worker_env(
     # leaves rollout.enable_sleep_mode defaulted True, so the engine always builds a
     # CuMemAllocator -- which asserts outright on "expandable_segments:True"
     # (vllm/device_allocator/cumem.py:132, pytorch#147851). Both therefore take the
-    # non-expandable conf.
+    # non-expandable conf. SFT and custom runs never build a vLLM engine (custom samples with
+    # HF generate), so they take expandable segments: a custom session's varied sequence lengths
+    # otherwise fragment the caching allocator until cuBLAS/Triton allocations outside it fail.
     _alloc_conf = (
         "expandable_segments:True"
-        if str(getattr(spec, "algorithm", "")).lower() == "sft"
+        if str(getattr(spec, "algorithm", "")).lower() in ("sft", "custom")
         else "garbage_collection_threshold:0.8,max_split_size_mb:256"
     )
     env: dict[str, str] = {

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.4.16"
+version = "1.4.17"
 description ="Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI"
 # The training and serving stack is validated on Python 3.11 and 3.12.
 requires-python = ">=3.11,<3.13"
@@ -210,7 +210,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "1.4.16"
+version = "1.4.17"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/tests/test_custom_controller_connect.py

@@ -2,7 +2,13 @@
 
 import pytest
 
-from flash.client.http import ApiError, ClientError
+from flash.client.http import (
+    ApiClient,
+    ApiError,
+    ClientError,
+    RequestTimeoutError,
+    ServiceUnreachableError,
+)
 from flash.custom import controller
 
 
@@ -73,7 +79,9 @@ def test_provisioning_then_ready(clock):
         ApiError(500, "Internal Server Error"),
         ApiError(502, "Bad Gateway"),
         ApiError(503, "Service Unavailable"),
-        ClientError("request to the Flash service timed out"),
+        ApiError(504, "Gateway Timeout"),
+        RequestTimeoutError("request to the Flash service timed out"),
+        ServiceUnreachableError("connection reset"),
     ],
 )
 def test_transient_control_plane_failures_are_retried(clock, transient):
@@ -92,6 +100,21 @@ def test_backoff_is_capped(clock):
     assert max(clock.sleeps) == controller._ATTACH_RETRY_MAX_SECONDS
 
 
+@pytest.mark.parametrize("body", [b"<html>wrong service</html>", b"[]", b"null"])
+def test_malformed_control_plane_response_fails_without_retry(clock, body):
+    decoder = ApiClient("https://plane.example", "test-key")
+    with pytest.raises(ClientError) as decoded:
+        decoder._decode_response("/v1/runs/run-1/custom", body)
+    api = FakeApi([decoded.value, CONNECTION])
+
+    with pytest.raises(ClientError) as caught:
+        attach(api, clock)
+
+    assert caught.value is decoded.value
+    assert api.calls == 1
+    assert clock.sleeps == []
+
+
 @pytest.mark.parametrize("status", [400, 401, 403, 404, 409])
 def test_definitive_client_errors_stay_fatal(clock, status):
     api = FakeApi([ApiError(status, "no")])

freesolo_flash_dev-1.4.17/tests/test_custom_cuda_headroom.py

@@ -0,0 +1,61 @@
+"""The custom backend returns cached CUDA blocks to the driver before an operation when free
+memory is low, so allocations outside PyTorch's cache (cuBLAS, Triton) can still succeed.
+
+A fake torch module stands in for CUDA, so these run on CPU-only CI without torch installed.
+"""
+
+import sys
+import types
+
+from flash.custom import backend
+
+GIB = 1024**3
+
+
+def fake_cuda(monkeypatch, *, available=True, free=1 * GIB, total=80 * GIB, fail=False):
+    calls = []
+
+    def mem_get_info():
+        if fail:
+            raise RuntimeError("driver unavailable")
+        return free, total
+
+    cuda = types.SimpleNamespace(
+        is_available=lambda: available,
+        mem_get_info=mem_get_info,
+        empty_cache=lambda: calls.append("empty_cache"),
+    )
+    monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(cuda=cuda))
+    return calls
+
+
+def test_low_free_memory_releases_the_cache(monkeypatch):
+    calls = fake_cuda(monkeypatch, free=1 * GIB)
+    assert backend._ensure_cuda_headroom() is True
+    assert calls == ["empty_cache"]
+
+
+def test_enough_free_memory_keeps_the_cache(monkeypatch):
+    calls = fake_cuda(monkeypatch, free=backend.CUDA_HEADROOM_BYTES)
+    assert backend._ensure_cuda_headroom() is False
+    assert calls == []
+
+
+def test_no_cuda_is_a_no_op(monkeypatch):
+    calls = fake_cuda(monkeypatch, available=False)
+    assert backend._ensure_cuda_headroom() is False
+    assert calls == []
+
+
+def test_driver_errors_never_fail_an_operation(monkeypatch):
+    calls = fake_cuda(monkeypatch, fail=True)
+    assert backend._ensure_cuda_headroom() is False
+    assert calls == []
+
+
+def test_every_operation_checks_headroom_first(monkeypatch):
+    order = []
+    monkeypatch.setattr(backend, "_ensure_cuda_headroom", lambda: order.append("headroom"))
+    fake = types.SimpleNamespace(forward=lambda **kw: order.append(("forward", kw)) or {"ok": True})
+    assert backend.TorchBackend.execute(fake, "forward", {"data": []}) == {"ok": True}
+    assert order == ["headroom", ("forward", {"data": []})]

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/tests/test_flash_worker.py

@@ -354,6 +354,22 @@ def test_alloc_conf_default_expandable_for_sft(monkeypatch):
     assert env["PYTORCH_ALLOC_CONF"] == "expandable_segments:True"
 
 
+def test_alloc_conf_expandable_for_custom(monkeypatch):
+    # custom runs build no vLLM engine; without expandable segments a session's varied sequence
+    # lengths fragment the cache until cuBLAS/Triton allocations outside it fail.
+    from flash.core.spec import JobSpec, TrainSpec
+    from flash.providers._lifecycle.net.worker import build_worker_env
+
+    monkeypatch.delenv("PYTORCH_ALLOC_CONF", raising=False)
+    monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False)
+    spec = JobSpec(
+        model="Qwen/Qwen3.5-9B", algorithm="custom", train=TrainSpec(max_steps=4), seed=0
+    )
+    env = build_worker_env(spec)
+    assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"
+    assert env["PYTORCH_ALLOC_CONF"] == env["PYTORCH_CUDA_ALLOC_CONF"]
+
+
 def test_error_artifact_name_is_per_phase_and_attempt():
     """Error files are scoped per-phase and per-attempt so a stale prior-attempt
     traceback can't be mistaken for the current attempt's crash on a retry."""

{freesolo_flash_dev-1.4.16 → freesolo_flash_dev-1.4.17}/uv.lock

@@ -1297,7 +1297,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.4.16"
+version = "1.4.17"
 source = { editable = "." }
 
 [package.optional-dependencies]