diff --git a/DESCRIPTION b/DESCRIPTION index b1f6bb5..d85c958 100644 --- a/DESCRIPTION +++ b/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")), diff --git a/NEWS.md b/NEWS.md index ecc67ba..a7e85df 100644 --- a/NEWS.md +++ b/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. diff --git a/R/resident.R b/R/resident.R index 6011cf4..29e5eb2 100644 --- a/R/resident.R +++ b/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, ...) { diff --git a/inst/tinytest/test_resident_sdxl.R b/inst/tinytest/test_resident_sdxl.R index f872c83..8a59eaf 100644 --- a/inst/tinytest/test_resident_sdxl.R +++ b/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"))) diff --git a/man/dot-cuda_index.Rd b/man/dot-cuda_index.Rd new file mode 100644 index 0000000..bf60aba --- /dev/null +++ b/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} diff --git a/man/dot-resident_prewarm.Rd b/man/dot-resident_prewarm.Rd index df25031..a8912ad 100644 --- a/man/dot-resident_prewarm.Rd +++ b/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} diff --git a/man/resident_generate.Rd b/man/resident_generate.Rd index cc7c557..b41e4d7 100644 --- a/man/resident_generate.Rd +++ b/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