Fix pre-warm accumulation; correct resident_generate()'s @return

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.2.2.5
+Version: 0.2.2.6
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NEWS.md

@@ -1,3 +1,35 @@
+# diffuseR 0.2.2.6
+
+* Fixed an allocator pre-warm accumulation introduced in 0.2.2.4.
+  `.resident_prewarm()` requested the full onload need on every
+  activation, which doubled the CUDA caching allocator's pool after each
+  render: a generation fragments the cache, so the next single large
+  request cannot be served from it and takes a fresh allocation beside the
+  old one. Under `resident_deactivate(release = FALSE)` nothing empties
+  the cache, so SDXL went 5.299 GiB after one cycle to 10.322 after two
+  and refused the third. It now measures the free cache on the handle's
+  own device and grows only the shortfall, skipping entirely when the pool
+  already covers the transfer; a cold pool is unchanged, and the
+  cold-start win is intact (2.48 s against 2.51 s before). Measured flat
+  at 5.396 / 5.398 / 5.398 / 5.398 / 5.398 GiB across five cycles, with a
+  phase-offloading family (flux2) untouched because it never takes the
+  bulk branch.
+
+  This also corrected the budget independently of the refusal: both
+  release modes doubled between the first and second cycle, so any peak
+  measured on a single activation understated steady state roughly 2x.
+
+* `resident_generate()`'s documented return value was wrong. It claimed
+  `flux1`, `flux2` and `zimage` return bare image arrays and `sdxl` was
+  the exception. Every family returns a list: the five image families
+  return `list(image, metadata)`, so `$image` unwraps uniformly across all
+  five. `ltx` returns `latents`, `audio_latents`, `latent_shape` and
+  `sample_rate`, plus `video` and `audio` when `decode_video` /
+  `decode_audio` are TRUE — a caller that turns either off gets a list
+  without that field rather than a NULL one. Only visibility differs:
+  `txt2img_sdxl()` and `txt2img_sd21()` use `return()`, the rest
+  `invisible()`.
+
 # diffuseR 0.2.2.5
 
 * `resident_load()` accepts `"sd21"`, the sixth resident family.

R/resident.R

@@ -331,29 +331,109 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx",
 #' the per-tensor path, which is the current behaviour and merely slow, so
 #' the failure is swallowed rather than raised.
 #'
-#' @param bytes Numeric. Host bytes about to be transferred; the pool is
-#'   warmed to this plus a small margin for allocator slack.
-#' @param device Target CUDA device.
+#' Growing the pool is best-effort in the partial case. A request smaller
+#' than a free block already in the cache is served from that block and
+#' grows nothing, so when the pool is short by less than it already holds
+#' the pre-warm may be absorbed rather than add capacity. That is bounded
+#' and harmless -- the onload then falls back to the per-tensor path for
+#' the remainder, which is the old behaviour -- and the alternative, asking
+#' for the whole figure to force a new segment, is the accumulation bug
+#' this function exists to avoid. The cold pool, which is the case worth
+#' optimising and the one a broker's first request hits, is unaffected.
 #'
-#' @return Invisibly NULL.
+#' @param bytes Numeric. Host bytes about to be transferred; the pool is
+#'   warmed toward this plus a small margin for allocator slack.
+#' @param device Target CUDA device, e.g. "cuda" or "cuda:1". Also selects
+#'   which device's allocator is measured.
+#' @param held Free cached bytes the allocator already holds on that
+#'   device, i.e. reserved minus allocated -- bytes that are reserved but
+#'   live belong to something else and cannot serve this transfer. NULL
+#'   measures it. Pass a value to make the decision deterministic: without
+#'   CUDA the measurement is 0, which would always warm, so a test that
+#'   wants the skip has to state what the pool holds rather than depend on
+#'   the machine having a card. Same reason
+#'   \code{\link{.resident_check_fits}} takes \code{free_gb}.
+#'
+#' @return Invisibly, the bytes requested from the allocator: 0 when the
+#'   pool already covers the transfer and nothing was asked for.
 #'
 #' @keywords internal
-.resident_prewarm <- function(bytes, device) {
+.resident_prewarm <- function(bytes, device, held = NULL) {
     # isTRUE() rather than a bare is.finite(): a NULL pinned_bytes gives
     # logical(0), and `||` on a zero-length value is an error in R >= 4.3,
     # so the guard meant to skip the pre-warm would instead fail the
     # activation it exists to speed up.
     if (!isTRUE(is.finite(bytes)) || bytes <= 0) {
-        return(invisible(NULL))
+        return(invisible(0))
+    }
+    # Only grow what is missing, and only when something IS missing.
+    #
+    # Asking for the full figure unconditionally doubles the pool on every
+    # activation after the first. A render fragments the cache into its
+    # activation blocks, so the next single large request cannot be served
+    # from it and takes a fresh cudaMalloc alongside the old block. With
+    # release = FALSE -- which a residency broker passes deliberately, to
+    # keep an exclusive grant's blocks off other tenants -- nothing ever
+    # empties the cache, so it grows by one block per cycle until
+    # activation is refused. Measured on SDXL: 5.299 GiB after cycle 1,
+    # 10.322 after cycle 2, refused on cycle 3, with the step (5.023 GiB)
+    # matching the pre-warm block (5.021 GiB) to two thousandths.
+    #
+    # Bare activate/deactivate cycles never showed it: without a render the
+    # block is reused cleanly and the pool holds flat. It takes a
+    # generation in between, which is why one cycle is not a test.
+    if (is.null(held)) {
+        held <- tryCatch({
+            s <- torch::cuda_memory_stats(device = .cuda_index(device))
+            as.numeric(s$reserved_bytes$all$current) -
+                as.numeric(s$allocated_bytes$all$current)
+        }, error = function(e) 0)
+    }
+    if (!isTRUE(is.finite(held)) || held < 0) {
+        held <- 0
     }
+    # One target, used for both the skip and the size, so the two cannot
+    # disagree. Skipping at `held >= bytes` while growing toward
+    # `bytes * 1.05` put a step in the middle: 3.999 GiB held asked for
+    # 0.201 GiB and 4.000 GiB held asked for nothing.
+    target <- as.numeric(bytes) * 1.05
+    if (held >= target) {
+        return(invisible(0))
+    }
+    want <- target - held
     tryCatch({
-        warm <- torch::torch_empty(as.numeric(bytes) * 1.05,
-                                   dtype = torch::torch_uint8(),
+        warm <- torch::torch_empty(want, dtype = torch::torch_uint8(),
                                    device = device)
         rm(warm)
         gc(verbose = FALSE)
     }, error = function(e) invisible(NULL))
-    invisible(NULL)
+    invisible(want)
+}
+
+#' Device ordinal for a torch device string
+#'
+#' \code{torch::cuda_memory_stats()} defaults to
+#' \code{cuda_current_device()}, so reading it without an argument reports
+#' whichever device happens to be current rather than the one a handle is
+#' bound to. \code{resident_load()} binds an explicit \code{"cuda:N"}
+#' precisely so transitions cannot drift, and a handle on \code{cuda:1}
+#' deciding from \code{cuda:0}'s pool would either skip a pre-warm it needs
+#' or repeat one it does not.
+#'
+#' @param device Character, e.g. "cuda", "cuda:0", "cuda:1".
+#'
+#' @return Integer ordinal. An unqualified device gives the current one.
+#'
+#' @keywords internal
+.cuda_index <- function(device) {
+    d <- as.character(device)[[1]]
+    if (grepl(":", d, fixed = TRUE)) {
+        n <- suppressWarnings(as.integer(sub("^.*:", "", d)))
+        if (!is.na(n)) {
+            return(n)
+        }
+    }
+    tryCatch(torch::cuda_current_device(), error = function(e) 0L)
 }
 
 #' Which components a bulk activation puts on the card
@@ -601,13 +681,25 @@ resident_deactivate <- function(res, release = TRUE) {
 #'   \code{sdxl} the handle supplies \code{devices} matching its own
 #'   placement unless the caller names it.
 #'
-#' @return Whatever the family generator returns, and the families do not
-#'   agree: an image array for \code{flux1}, \code{flux2} and
-#'   \code{zimage}, a video array for \code{ltx}, and for \code{sdxl} a
-#'   list of \code{image} and \code{metadata}, because
-#'   \code{\link{txt2img_sdxl}} has always returned that pair and changing
-#'   it would break every existing caller. A broker that wants one shape
-#'   should normalise in its own wrapper.
+#' @return Whatever the family generator returns, which is always a list.
+#'
+#'   The five image families (\code{flux1}, \code{flux2}, \code{zimage},
+#'   \code{sdxl}, \code{sd21}) return \code{list(image, metadata)}, where
+#'   \code{image} is an [H, W, 3] array in [0, 1], so a caller unwraps
+#'   \code{$image} uniformly across all five.
+#'
+#'   \code{ltx} returns \code{latents}, \code{audio_latents},
+#'   \code{latent_shape} and \code{sample_rate}, plus \code{video} and
+#'   \code{audio} -- but those two are produced only when
+#'   \code{decode_video} and \code{decode_audio} are TRUE, which they are
+#'   by default. A caller that turns either off gets a list without that
+#'   field rather than a NULL one, so index it with \code{[[ ]]} and check,
+#'   the way \code{\link{txt2vid_ltx2}} does internally.
+#'
+#'   Only the visibility differs: \code{\link{txt2img_sdxl}} and
+#'   \code{\link{txt2img_sd21}} use \code{return()} while the other three
+#'   image families and \code{ltx} use \code{invisible()}, which affects
+#'   auto-printing at the console and nothing else.
 #'
 #' @export
 resident_generate <- function(res, prompt, ...) {

inst/tinytest/test_resident_sdxl.R

@@ -213,4 +213,61 @@ expect_silent(diffuseR:::.resident_prewarm(NA_real_, "cuda"))
 expect_silent(diffuseR:::.resident_prewarm(NULL, "cuda"))
 # An impossible size on a real card must be swallowed, not raised.
 expect_silent(diffuseR:::.resident_prewarm(1e18, "cuda"))
-expect_null(diffuseR:::.resident_prewarm(1024, "cuda"))
+expect_equal(diffuseR:::.resident_prewarm(0, "cuda"), 0)
+
+# It must GROW the pool, not re-request it. Asking for the full figure on
+# every activation doubled the cache once a render had fragmented it: the
+# single large request could not be served from cache and took a fresh
+# cudaMalloc beside the old block. Under release = FALSE -- which a
+# residency broker passes deliberately -- nothing empties the cache, so
+# SDXL went 5.299 -> 10.322 GiB and the third activation was refused.
+gb <- 1024^3
+
+# One target governs both the skip and the size. Skipping at held >= bytes
+# while growing toward bytes * 1.05 put a step in the middle, so the
+# threshold is the target itself.
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 5 * gb), 0)
+# Exactly at the target still counts as covered.
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 4 * gb * 1.05),
+             0)
+
+# Pool short: grow toward bytes * 1.05, so the 5% margin lands on the
+# FINAL pool. (bytes - held) * 1.05 would ask for 5% of the gap instead
+# and undershoot the target whenever held > 0.
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 3 * gb),
+             4 * gb * 1.05 - 3 * gb)
+
+# No discontinuity around `bytes`: holding a hair under and a hair over the
+# raw need must differ by a hair, not by the whole margin. This is the case
+# the old threshold got wrong.
+lo <- diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 4 * gb - 1)
+hi <- diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 4 * gb + 1)
+expect_true(abs(lo - hi) < 10)
+expect_true(lo > 0 && hi > 0)
+
+# Cold pool: byte-identical to the original behaviour, so the 74x
+# cold-start win is untouched.
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = 0),
+             4 * gb * 1.05)
+
+# A nonsense reading must not be trusted into a negative request.
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = NA_real_),
+             4 * gb * 1.05)
+expect_equal(diffuseR:::.resident_prewarm(4 * gb, "cuda", held = -1),
+             4 * gb * 1.05)
+
+# --- the allocator is read on the handle's own device -------------------------------
+
+# cuda_memory_stats() defaults to cuda_current_device(), so reading it
+# without an argument reports whichever device is current rather than the
+# one the handle bound. resident_load() binds an explicit "cuda:N" so
+# transitions cannot drift; a cuda:1 handle deciding from cuda:0's pool
+# would skip a pre-warm it needs or repeat one it does not.
+expect_equal(diffuseR:::.cuda_index("cuda:0"), 0L)
+expect_equal(diffuseR:::.cuda_index("cuda:1"), 1L)
+expect_equal(diffuseR:::.cuda_index("cuda:7"), 7L)
+# Unqualified falls back to the current device, whatever that is.
+expect_true(is.numeric(diffuseR:::.cuda_index("cuda")))
+# A malformed ordinal must not become NA and poison the stats lookup.
+expect_true(is.numeric(diffuseR:::.cuda_index("cuda:x")))
+expect_false(is.na(diffuseR:::.cuda_index("cuda:x")))

man/dot-cuda_index.Rd

@@ -0,0 +1,23 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{.cuda_index}
+\alias{.cuda_index}
+\title{Device ordinal for a torch device string}
+\usage{
+.cuda_index(device)
+}
+\arguments{
+\item{device}{Character, e.g. "cuda", "cuda:0", "cuda:1".}
+}
+\value{
+Integer ordinal. An unqualified device gives the current one.
+}
+\description{
+\code{torch::cuda_memory_stats()} defaults to
+\code{cuda_current_device()}, so reading it without an argument reports
+whichever device happens to be current rather than the one a handle is
+bound to. \code{resident_load()} binds an explicit \code{"cuda:N"}
+precisely so transitions cannot drift, and a handle on \code{cuda:1}
+deciding from \code{cuda:0}'s pool would either skip a pre-warm it needs
+or repeat one it does not.
+}
+\keyword{internal}

man/dot-resident_prewarm.Rd

@@ -3,16 +3,27 @@
 \alias{.resident_prewarm}
 \title{Grow the caching allocator's pool in one allocation before a bulk onload}
 \usage{
-.resident_prewarm(bytes, device)
+.resident_prewarm(bytes, device, held = NULL)
 }
 \arguments{
 \item{bytes}{Numeric. Host bytes about to be transferred; the pool is
-warmed to this plus a small margin for allocator slack.}
+warmed toward this plus a small margin for allocator slack.}
 
-\item{device}{Target CUDA device.}
+\item{device}{Target CUDA device, e.g. "cuda" or "cuda:1". Also selects
+which device's allocator is measured.}
+
+\item{held}{Free cached bytes the allocator already holds on that
+device, i.e. reserved minus allocated -- bytes that are reserved but
+live belong to something else and cannot serve this transfer. NULL
+measures it. Pass a value to make the decision deterministic: without
+CUDA the measurement is 0, which would always warm, so a test that
+wants the skip has to state what the pool holds rather than depend on
+the machine having a card. Same reason
+\code{\link{.resident_check_fits}} takes \code{free_gb}.}
 }
 \value{
-Invisibly NULL.
+Invisibly, the bytes requested from the allocator: 0 when the
+  pool already covers the transfer and nothing was asked for.
 }
 \description{
 A cold bulk onload grows the pool one \code{cudaMalloc} per tensor, and
@@ -31,5 +42,15 @@ Best-effort. A card that cannot seat the block in one piece falls back to
 the per-tensor path, which is the current behaviour and merely slow, so
 the failure is swallowed rather than raised.
 
+Growing the pool is best-effort in the partial case. A request smaller
+than a free block already in the cache is served from that block and
+grows nothing, so when the pool is short by less than it already holds
+the pre-warm may be absorbed rather than add capacity. That is bounded
+and harmless -- the onload then falls back to the per-tensor path for
+the remainder, which is the old behaviour -- and the alternative, asking
+for the whole figure to force a new segment, is the accumulation bug
+this function exists to avoid. The cold pool, which is the case worth
+optimising and the one a broker's first request hits, is unaffected.
+
 }
 \keyword{internal}

man/resident_generate.Rd

@@ -17,13 +17,25 @@ resident_generate(res, prompt, ...)
 placement unless the caller names it.}
 }
 \value{
-Whatever the family generator returns, and the families do not
-  agree: an image array for \code{flux1}, \code{flux2} and
-  \code{zimage}, a video array for \code{ltx}, and for \code{sdxl} a
-  list of \code{image} and \code{metadata}, because
-  \code{\link{txt2img_sdxl}} has always returned that pair and changing
-  it would break every existing caller. A broker that wants one shape
-  should normalise in its own wrapper.
+Whatever the family generator returns, which is always a list.
+
+  The five image families (\code{flux1}, \code{flux2}, \code{zimage},
+  \code{sdxl}, \code{sd21}) return \code{list(image, metadata)}, where
+  \code{image} is an [H, W, 3] array in [0, 1], so a caller unwraps
+  \code{$image} uniformly across all five.
+
+  \code{ltx} returns \code{latents}, \code{audio_latents},
+  \code{latent_shape} and \code{sample_rate}, plus \code{video} and
+  \code{audio} -- but those two are produced only when
+  \code{decode_video} and \code{decode_audio} are TRUE, which they are
+  by default. A caller that turns either off gets a list without that
+  field rather than a NULL one, so index it with \code{[[ ]]} and check,
+  the way \code{\link{txt2vid_ltx2}} does internally.
+
+  Only the visibility differs: \code{\link{txt2img_sdxl}} and
+  \code{\link{txt2img_sd21}} use \code{return()} while the other three
+  image families and \code{ltx} use \code{invisible()}, which affects
+  auto-printing at the console and nothing else.
 }
 \description{
 Dispatches to the family's generator with the resident pipeline